big-qmt/py-client/libs/orderbook.py

"""SQLite positions and deals, aligned with SDK models; one writer per database."""

import math
import sqlite3
from contextlib import closing
from dataclasses import asdict, fields
from datetime import datetime
from pathlib import Path
from itertools import chain

from sdk import DealItem, PositionItem

SCHEMA = """
CREATE TABLE IF NOT EXISTS positions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    stock_code TEXT NOT NULL,
    stock_name TEXT NOT NULL DEFAULT '',
    direction INTEGER,
    volume INTEGER NOT NULL DEFAULT 0 CHECK (volume >= 0),
    open_price REAL NOT NULL DEFAULT 0,
    open_cost REAL NOT NULL DEFAULT 0,
    float_profit REAL NOT NULL DEFAULT 0,
    market_value REAL NOT NULL DEFAULT 0,
    stock_holder TEXT NOT NULL DEFAULT '',
    frozen_volume INTEGER NOT NULL DEFAULT 0,
    can_use_volume INTEGER NOT NULL DEFAULT 0,
    on_road_volume INTEGER NOT NULL DEFAULT 0,
    yesterday_volume INTEGER NOT NULL DEFAULT 0,
    last_price REAL NOT NULL DEFAULT 0,
    profit_rate REAL NOT NULL DEFAULT 0,
    future_trade_type INTEGER,
    expire_date TEXT NOT NULL DEFAULT ''
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_positions_stock_code ON positions (stock_code);

CREATE TABLE IF NOT EXISTS deals (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    stock_code TEXT NOT NULL,
    order_sys_id TEXT NOT NULL CHECK (order_sys_id <> ''),
    order_local_id TEXT NOT NULL CHECK (order_local_id <> ''),
    ref INTEGER NOT NULL DEFAULT 0,
    order_ref TEXT NOT NULL DEFAULT '',
    direction INTEGER NOT NULL DEFAULT 0,
    offset_flag INTEGER NOT NULL CHECK (offset_flag IN (23, 24, 48, 49)),
    price REAL NOT NULL CHECK (price >= 0),
    volume INTEGER NOT NULL CHECK (volume > 0),
    trade_amount REAL NOT NULL CHECK (trade_amount > 0),
    trade_date TEXT NOT NULL,
    trade_time TEXT NOT NULL,
    remark TEXT NOT NULL DEFAULT '',
    close_profit REAL NOT NULL DEFAULT 0
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_deals_order_sys_id ON deals (order_sys_id);
CREATE INDEX IF NOT EXISTS idx_deals_order_ref ON deals (order_local_id);
CREATE INDEX IF NOT EXISTS idx_deals_stock_code_date ON deals (stock_code);
CREATE INDEX IF NOT EXISTS idx_deals_date_time ON deals (trade_date);
"""

# SQL is assembled once; order_local_id is derived from the SDK remark property.
POSITION_FIELDS = tuple(field.name for field in fields(PositionItem))
DEAL_FIELDS = tuple(field.name for field in fields(DealItem)) + ('order_local_id',)
POSITION_DEFAULTS = asdict(PositionItem())
POSITION_UPSERT = (
    f"INSERT INTO positions ({', '.join(POSITION_FIELDS)}) "
    f"VALUES ({', '.join(':' + key for key in POSITION_FIELDS)}) "
    "ON CONFLICT(stock_code) DO UPDATE SET "
    + ', '.join(f'{key} = excluded.{key}' for key in POSITION_FIELDS if key != 'stock_code')
)
DEAL_INSERT = (
    f"INSERT INTO deals ({', '.join(DEAL_FIELDS)}) "
    f"VALUES ({', '.join(':' + key for key in DEAL_FIELDS)})"
)


class OrderBook:
    """Position snapshots and append-only deals. No schema migration."""

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.positions: dict[str, dict] = {}
        self.deals: dict[str, dict] = {}
        self.deals_sys_ids: set[str] = set()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with closing(self._connect()) as db:
            db.executescript(SCHEMA)
        self.load()

    def _connect(self) -> sqlite3.Connection:
        db = sqlite3.connect(self.path, timeout=30)
        db.row_factory = sqlite3.Row
        return db

    def load(self) -> None:
        with closing(self._connect()) as db, db:
            db.execute('BEGIN')
            positions = {row['stock_code']: dict(row) for row in db.execute('SELECT * FROM positions')}
            deals = {row['order_sys_id']: dict(row) for row in db.execute('SELECT * FROM deals ORDER BY id')}
        self.positions = positions
        self.deals = deals
        self.deals_sys_ids = set(deals)

    @staticmethod
    def deal_record(deal: DealItem) -> dict:
        if not deal.order_sys_id or deal.volume <= 0:
            raise ValueError('System order ID and positive volume are required')
        
        row = asdict(deal)
        row['order_local_id'] = deal.local_order_id
        if not row['order_local_id']:
            raise ValueError('Local order ID is required')
        
        if any(isinstance(value, float) and not math.isfinite(value) for value in row.values()):
            raise ValueError('Numeric values must be finite')
        amount = deal.trade_amount if deal.trade_amount > 0 else deal.price * deal.volume
        if not math.isfinite(amount) or amount <= 0:
            raise ValueError('Trade amount must be positive and finite')
        row['trade_amount'] = amount
        date = deal.trade_date or datetime.now().date().isoformat()
        if len(date) == 8 and date.isdigit():
            date = f'{date[:4]}-{date[4:6]}-{date[6:]}'
        row['trade_date'] = date
        return row

    def sync_deals(self, deals: list[DealItem]) -> None:
        new_deals = {}
        for deal in deals:
            if deal.order_sys_id not in self.deals_sys_ids and deal.order_sys_id not in new_deals:
                new_deals[deal.order_sys_id] = self.deal_record(deal)
        if not new_deals:
            return
        with closing(self._connect()) as db, db:
            db.executemany(DEAL_INSERT, new_deals.values())
        self.load()

    def sync_positions(self, positions: list[PositionItem]) -> None:
        """Replace the complete position snapshot, retaining IDs for existing stocks."""
        self.save({item.stock_code: asdict(item) for item in positions})

    def save(self, items: dict, deals: list[dict] | None = None) -> None:
        """保存持仓快照，可同时追加成交；省略 deals 时仅更新持仓。"""
        new_deals = []
        if deals is not None:
            if len(deals) < len(self.deals):
                raise ValueError('Execution history is append-only')
            new_deals = deals[len(self.deals):]
        positions = [{**POSITION_DEFAULTS, **item} for item in items.values()]
        for row in chain(positions, new_deals):
            if any(isinstance(value, float) and not math.isfinite(value) for value in row.values()):
                raise ValueError('Numeric values must be finite')
        with closing(self._connect()) as db, db:
            removed = [
                (row['stock_code'],)
                for row in db.execute('SELECT stock_code FROM positions')
                if row['stock_code'] not in items
            ]
            db.executemany('DELETE FROM positions WHERE stock_code = ?', removed)
            db.executemany(POSITION_UPSERT, positions)
            db.executemany(DEAL_INSERT, new_deals)
        self.load()