zhixing-system/zhixing-server/src/zhixing_server/modules/selection/infrastructure/postgres_reader.py

"""Read-only PostgreSQL adapter for selection history."""

from __future__ import annotations

from datetime import date, datetime
from decimal import Decimal, InvalidOperation
from typing import cast

import psycopg

from ....bootstrap.config import Settings
from ..domain.models import SelectionBar, SelectionDailyBasic, StockHistory
from ..domain.ports import MarketDataReaderError
from ..domain.runs import SelectionExecutionSource, SelectionStock


class SelectionReaderError(MarketDataReaderError):
    """Database read failure with stock and target-date context."""


class SelectionMarketDataNotReady(MarketDataReaderError):
    """The requested date has no market-data batch eligible for selection."""


_HISTORY_QUERY = """
SELECT
    bar.ts_code,
    stock.name,
    bar.trade_date,
    bar.open,
    bar.high,
    bar.low,
    bar.close,
    bar.vol,
    basic.turnover_rate,
    basic.total_mv
FROM market_daily_bar AS bar
LEFT JOIN market_stock AS stock
  ON stock.ts_code = bar.ts_code
LEFT JOIN market_daily_basic AS basic
  ON basic.ts_code = bar.ts_code
 AND basic.trade_date = bar.trade_date
WHERE bar.ts_code = %s
  AND bar.source_adj = 'qfq'
  AND bar.trade_date <= %s
ORDER BY bar.trade_date ASC
"""

_SOURCE_QUERY = """
SELECT id, target_count, valid_count, coverage
FROM market_sync_batch
WHERE target_trade_date = %s
  AND strategy_eligible = true
  AND status IN ('success', 'partial_success')
ORDER BY finished_at DESC NULLS LAST, created_at DESC, id DESC
LIMIT 1
"""

_ELIGIBLE_STOCKS_QUERY = """
SELECT stock.ts_code, stock.name
FROM market_stock AS stock
WHERE stock.is_active = true
  AND EXISTS (
      SELECT 1
      FROM market_daily_bar AS bar
      WHERE bar.ts_code = stock.ts_code
        AND bar.trade_date = %s
        AND bar.source_adj = 'qfq'
        AND bar.open IS NOT NULL
        AND bar.high IS NOT NULL
        AND bar.low IS NOT NULL
        AND bar.close IS NOT NULL
        AND bar.vol IS NOT NULL
  )
  AND EXISTS (
      SELECT 1
      FROM market_daily_basic AS basic
      WHERE basic.ts_code = stock.ts_code
        AND basic.trade_date = %s
  )
ORDER BY stock.ts_code
"""


def _as_date(value: object) -> date:
    """Convert a PostgreSQL date-like scalar to a date."""

    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, date):
        return value
    return date.fromisoformat(str(value)[:10])


def _as_float(value: object) -> float | None:
    """Convert nullable PostgreSQL numerics to finite floats."""

    if value is None:
        return None
    try:
        number = Decimal(str(value))
    except (InvalidOperation, ValueError) as exc:
        raise ValueError(f"invalid market-data numeric value: {value!r}") from exc
    if number.is_nan():
        return None
    if not number.is_finite():
        raise ValueError(f"market-data numeric value must be finite: {value!r}")
    return float(number)


class PostgresMarketDataReader:
    """Load qfq bars and same-day basic facts without writing market data."""

    def __init__(self, settings: Settings | str) -> None:
        """Create a reader from injected settings or a compatible URL string."""

        self.database_url = settings.database_url if isinstance(settings, Settings) else settings

    def load_history(self, ts_code: str, target_trade_date: date) -> StockHistory:
        """Read all retained qfq rows through the explicit target date.

        Args:
            ts_code: Tushare stock identifier.
            target_trade_date: Historical date to which rows are truncated.

        Returns:
            A sorted ``StockHistory``.  An empty history is a normal missing
            target-data result and is interpreted by the application layer.

        Raises:
            SelectionReaderError: If PostgreSQL cannot complete the read.
            ValueError: If a returned date or numeric field is malformed.
        """

        try:
            with psycopg.connect(self.database_url) as connection:
                rows = connection.execute(
                    _HISTORY_QUERY,
                    (ts_code, target_trade_date),
                ).fetchall()
        except psycopg.Error as exc:
            raise SelectionReaderError(
                f"failed to load market history for {ts_code} at {target_trade_date.isoformat()}"
            ) from exc

        bars: dict[date, SelectionBar] = {}
        daily_basic: dict[date, SelectionDailyBasic] = {}
        name = ""
        for raw_row in rows:
            row = cast(tuple[object, ...], raw_row)
            row_code, row_name, bar, basic = self._map_row(row, ts_code)
            if row_code != ts_code:
                raise ValueError(f"reader returned unexpected stock code: {row_code}")
            name = row_name or name
            if bar.trade_date <= target_trade_date:
                bars[bar.trade_date] = bar
                daily_basic[bar.trade_date] = basic
        return StockHistory(
            ts_code=ts_code,
            name=name,
            bars=tuple(bars[trade_date] for trade_date in sorted(bars)),
            daily_basic={trade_date: daily_basic[trade_date] for trade_date in sorted(daily_basic)},
        )

    def load_execution_source(
        self,
        strategy: str,
        target_trade_date: date,
    ) -> SelectionExecutionSource:
        """Load the qualified market-data snapshot for a strategy run.

        Args:
            strategy: Supported strategy identity.  The current reader accepts
                ``zhixing_b1`` and keeps the parameter explicit for future
                strategy-specific eligibility rules.
            target_trade_date: Historical trading date to evaluate.

        Returns:
            The eligible stock snapshot and its source synchronization facts.

        Raises:
            SelectionMarketDataNotReady: If no eligible synchronization batch
                or complete active stock exists for the requested date.
            SelectionReaderError: If PostgreSQL cannot complete the read.
        """

        if strategy != "zhixing_b1":
            raise SelectionMarketDataNotReady(f"unsupported selection strategy: {strategy}")
        try:
            with psycopg.connect(self.database_url) as connection:
                source_row = connection.execute(_SOURCE_QUERY, (target_trade_date,)).fetchone()
                if source_row is None:
                    raise SelectionMarketDataNotReady(
                        f"market data is not strategy-eligible for {target_trade_date.isoformat()}"
                    )
                stock_rows = connection.execute(
                    _ELIGIBLE_STOCKS_QUERY,
                    (target_trade_date, target_trade_date),
                ).fetchall()
        except SelectionMarketDataNotReady:
            raise
        except psycopg.Error as exc:
            raise SelectionReaderError(
                f"failed to load selection source at {target_trade_date.isoformat()}"
            ) from exc

        stocks = tuple(
            SelectionStock(ts_code=str(row[0]), name=str(row[1] or "")) for row in stock_rows
        )
        if not stocks:
            raise SelectionMarketDataNotReady(
                f"no eligible stocks have complete market data for {target_trade_date.isoformat()}"
            )
        return SelectionExecutionSource(
            market_sync_batch_id=str(source_row[0]),
            target_trade_date=target_trade_date,
            target_count=int(source_row[1]),
            valid_count=int(source_row[2]),
            coverage=Decimal(str(source_row[3])),
            stocks=stocks,
        )

    @staticmethod
    def _map_row(
        row: tuple[object, ...],
        expected_code: str,
    ) -> tuple[str, str, SelectionBar, SelectionDailyBasic]:
        """Map the current query row, tolerating a legacy test row without name."""

        if len(row) >= 10:
            code, raw_name, raw_date = row[0], row[1], row[2]
            values = row[3:]
        elif len(row) >= 9:
            code, raw_name, raw_date = row[0], "", row[1]
            values = row[2:]
        else:
            raise ValueError("market history row has too few columns")
        row_code = str(code or expected_code)
        name = str(raw_name or "")
        trade_date = _as_date(raw_date)
        if len(values) < 7:
            raise ValueError("market history row is missing OHLCV/basic columns")
        bar = SelectionBar(
            trade_date=trade_date,
            open=_as_float(values[0]),
            high=_as_float(values[1]),
            low=_as_float(values[2]),
            close=_as_float(values[3]),
            volume=_as_float(values[4]),
        )
        basic = SelectionDailyBasic(
            trade_date=trade_date,
            turnover_rate=_as_float(values[5]),
            total_mv=_as_float(values[6]),
        )
        return row_code, name, bar, basic