Files
trex-core/trex_core/adapters/persistence/postgresql.py
T
genichinandClaude Opus 4.8 13837acc60 Implement BAT-006: portfolio-hub ledger integration (REQ-017~019)
- REQ-017: LedgerPort contract (application/ledger_port.py) + portfolio-hub
  adapter + in-memory stub (adapters/ledger/), TREX_PORTFOLIO_HUB_* config
- REQ-018: fill→ledger recording in execution loop; FillStatus carries
  execution_id (deterministic synth) + fee; paper/kis adapters updated
- REQ-019: ledger outbox + retry worker (application/ledger_outbox.py),
  correlation_id idempotency, outcome vocabulary; 002_create_ledger_outbox.sql
- Tests: in-memory/portfolio-hub ledger, execution ledger hook, outbox,
  domain import sweep

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-04 22:57:31 +09:00

421 lines
18 KiB
Python

"""PostgreSQL persistence adapters (REQ-006).
- PostgreSQLSlotRepository: implements the domain SlotRepository protocol
(single-writer, pg_advisory_xact_lock serialization).
- PostgreSQLIntentStore: implements the application IntentStore port
(intent_log dedup gate + audit).
psycopg (v3) is imported lazily so the domain/test paths that use the
in-memory adapters don't require it. External I/O lives here, never in the
domain layer (NFR-1).
"""
from __future__ import annotations
import json
from contextlib import contextmanager
from datetime import date, datetime, timezone
from typing import Any, Iterator, Optional
from uuid import UUID, uuid4
from trex_core.application.ledger_outbox import (
LedgerWorkItem,
OutboxOutcome,
WorkItemStatus,
)
from trex_core.application.ledger_port import LedgerTrade
from trex_core.application.ports import IntentRecord
from trex_core.domain.slot import Slot, SlotNotFoundError, SlotState, SlotTransitionRecord
_SLOT_COLUMNS = (
"slot_id, trading_date, status, symbol, qty, avg_price, entry_price, "
"allocated_capital, order_id, strategy_id, broker_order_id, portfolio_id, "
"sell_strategy_id, forced_close_at, created_at, updated_at"
)
def _require_psycopg(module: Any = None) -> Any:
if module is not None:
return module
try:
import psycopg # type: ignore[import-not-found]
except ImportError as exc: # pragma: no cover - environment guard
raise RuntimeError(
"psycopg is required for PostgreSQL adapters. Install 'trex-core[postgres]'."
) from exc
return psycopg
class PostgreSQLSlotRepository:
"""psycopg-backed SlotRepository. Connection-per-call (orbit-proven)."""
def __init__(self, dsn: str, *, psycopg_module: Any = None) -> None:
self._psycopg = _require_psycopg(psycopg_module)
self._dsn = dsn
@contextmanager
def transactional_advisory_lock(self, lock_id: int) -> Iterator[None]:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
try:
with conn.cursor() as cur:
cur.execute("SELECT pg_advisory_xact_lock(%s)", (lock_id,))
yield
conn.commit()
except BaseException:
conn.rollback()
raise
def pre_allocate_slots(self, trading_date: date, count: int) -> int:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT COUNT(*) FROM daily_slot WHERE trading_date = %s", (trading_date,)
)
row = cur.fetchone()
if row and row[0] > 0:
conn.commit()
return 0
now = datetime.now(tz=timezone.utc)
rows = [
(uuid4(), trading_date, SlotState.IDLE.value, now, now) for _ in range(count)
]
cur.executemany(
"INSERT INTO daily_slot (slot_id, trading_date, status, created_at, updated_at) "
"VALUES (%s, %s, %s, %s, %s)",
rows,
)
conn.commit()
return count
def list_slots(self, trading_date: date) -> list[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot WHERE trading_date = %s",
(trading_date,),
)
return [self._row_to_slot(r) for r in cur.fetchall()]
def fetch_slot(self, slot_id: UUID) -> Optional[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot WHERE slot_id = %s", (slot_id,)
)
row = cur.fetchone()
return self._row_to_slot(row) if row else None
def find_idle_slot(self, trading_date: date) -> Optional[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot "
"WHERE trading_date = %s AND status = 'idle' LIMIT 1",
(trading_date,),
)
row = cur.fetchone()
return self._row_to_slot(row) if row else None
def update_slot(self, slot: Slot) -> None:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"UPDATE daily_slot SET status=%s, symbol=%s, qty=%s, avg_price=%s, "
"entry_price=%s, allocated_capital=%s, order_id=%s, strategy_id=%s, "
"broker_order_id=%s, portfolio_id=%s, sell_strategy_id=%s, "
"forced_close_at=%s, updated_at=%s WHERE slot_id=%s",
(
slot.status.value, slot.symbol, slot.qty, slot.avg_price,
slot.entry_price, slot.allocated_capital, slot.order_id,
slot.strategy_id, slot.broker_order_id, slot.portfolio_id,
slot.sell_strategy_id, slot.forced_close_at,
slot.updated_at or datetime.now(tz=timezone.utc), slot.slot_id,
),
)
if cur.rowcount == 0:
conn.rollback()
raise SlotNotFoundError(f"slot not found: {slot.slot_id}")
conn.commit()
def append_transition_log(self, record: SlotTransitionRecord) -> None:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO slot_transition_log "
"(slot_id, from_state, to_state, ts, reason, actor, extra) "
"VALUES (%s, %s, %s, %s, %s, %s, %s)",
(
record.slot_id, record.from_state.value, record.to_state.value,
record.ts, record.reason, record.actor, json.dumps(record.extra),
),
)
conn.commit()
def increment_symbol_entry(self, trading_date: date, symbol: str) -> int:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO daily_symbol_entries (trading_date, symbol, entry_count, last_entry_at) "
"VALUES (%s, %s, 1, %s) "
"ON CONFLICT (trading_date, symbol) DO UPDATE "
"SET entry_count = daily_symbol_entries.entry_count + 1, "
" last_entry_at = EXCLUDED.last_entry_at RETURNING entry_count",
(trading_date, symbol, datetime.now(tz=timezone.utc)),
)
row = cur.fetchone()
conn.commit()
return int(row[0]) if row else 0
def get_symbol_entry_count(self, trading_date: date, symbol: str) -> int:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT entry_count FROM daily_symbol_entries "
"WHERE trading_date = %s AND symbol = %s",
(trading_date, symbol),
)
row = cur.fetchone()
return int(row[0]) if row else 0
# --- order execution queue (BAT-004 / REQ-013, option A) ---
def find_pending_for_dispatch(self) -> Optional[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot "
"WHERE status IN ('pending_buy','pending_sell') AND broker_order_id IS NULL "
"ORDER BY updated_at LIMIT 1 FOR UPDATE SKIP LOCKED"
)
row = cur.fetchone()
return self._row_to_slot(row) if row else None
def mark_dispatched(self, slot_id: UUID, broker_order_id: str) -> bool:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"UPDATE daily_slot SET broker_order_id=%s, updated_at=%s "
"WHERE slot_id=%s AND broker_order_id IS NULL",
(broker_order_id, datetime.now(tz=timezone.utc), slot_id),
)
ok = cur.rowcount == 1
conn.commit()
return ok
def list_dispatched_pending(self) -> list[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot "
"WHERE status IN ('pending_buy','pending_sell') AND broker_order_id IS NOT NULL"
)
return [self._row_to_slot(r) for r in cur.fetchall()]
def clear_dispatch(self, slot_id: UUID) -> None:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"UPDATE daily_slot SET broker_order_id=NULL, updated_at=%s WHERE slot_id=%s",
(datetime.now(tz=timezone.utc), slot_id),
)
conn.commit()
def count_slots_by_status(self, trading_date: date) -> dict[SlotState, int]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT status, COUNT(*) FROM daily_slot WHERE trading_date = %s GROUP BY status",
(trading_date,),
)
return {SlotState(r[0]): int(r[1]) for r in cur.fetchall()}
@staticmethod
def _row_to_slot(row: tuple) -> Slot:
f = lambda v: float(v) if v is not None else None # noqa: E731
return Slot(
slot_id=row[0] if isinstance(row[0], UUID) else UUID(str(row[0])),
trading_date=row[1],
status=SlotState(row[2]),
symbol=row[3],
qty=f(row[4]),
avg_price=f(row[5]),
entry_price=f(row[6]),
allocated_capital=f(row[7]),
order_id=row[8],
strategy_id=row[9],
broker_order_id=row[10],
portfolio_id=row[11],
sell_strategy_id=row[12],
forced_close_at=row[13],
created_at=row[14],
updated_at=row[15],
)
class PostgreSQLIntentStore:
"""intent_log-backed IntentStore. intent_id UNIQUE is the dedup gate."""
def __init__(self, dsn: str, *, psycopg_module: Any = None) -> None:
self._psycopg = _require_psycopg(psycopg_module)
self._dsn = dsn
def try_insert(self, intent_id: str, intent_type: str, payload: dict) -> bool:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO intent_log (intent_id, intent_type, payload, status) "
"VALUES (%s, %s, %s, 'processing') ON CONFLICT (intent_id) DO NOTHING",
(intent_id, intent_type, json.dumps(payload)),
)
inserted = cur.rowcount == 1
conn.commit()
return inserted
def get(self, intent_id: str) -> Optional[IntentRecord]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT intent_id, intent_type, status, slot_id, reason "
"FROM intent_log WHERE intent_id = %s",
(intent_id,),
)
row = cur.fetchone()
if row is None:
return None
return _PgIntentRecord(
intent_id=row[0], intent_type=row[1], status=row[2],
slot_id=(row[3] if isinstance(row[3], UUID) or row[3] is None else UUID(str(row[3]))),
reason=row[4],
)
def finalize(
self, intent_id: str, status: str, slot_id: Optional[UUID], reason: Optional[str]
) -> None:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"UPDATE intent_log SET status=%s, slot_id=%s, reason=%s, updated_at=%s "
"WHERE intent_id=%s",
(status, slot_id, reason, datetime.now(tz=timezone.utc), intent_id),
)
conn.commit()
class _PgIntentRecord:
def __init__(self, intent_id, intent_type, status, slot_id, reason):
self.intent_id = intent_id
self.intent_type = intent_type
self.status = status
self.slot_id = slot_id
self.reason = reason
# --- ledger outbox (BAT-006 / REQ-019) ---------------------------------------
_OUTBOX_COLUMNS = (
"correlation_id, asset_id, order_id, execution_id, side, qty, price, "
"fee, tax, executed_at, status, attempts, last_outcome, last_error, "
"created_at, updated_at"
)
class PostgreSQLLedgerOutbox:
"""ledger_outbox-backed LedgerOutbox. correlation_id PK is the dedup gate."""
def __init__(self, dsn: str, *, psycopg_module: Any = None) -> None:
self._psycopg = _require_psycopg(psycopg_module)
self._dsn = dsn
def enqueue(self, trade: LedgerTrade) -> LedgerWorkItem:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO ledger_outbox "
"(correlation_id, asset_id, order_id, execution_id, side, "
" qty, price, fee, tax, executed_at) "
"VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) "
"ON CONFLICT (correlation_id) DO NOTHING",
(
trade.correlation_id, trade.asset_id, trade.order_id,
trade.execution_id, trade.side,
trade.qty, trade.price, trade.fee, trade.tax,
trade.executed_at,
),
)
conn.commit()
existing = self.get(trade.correlation_id)
assert existing is not None
return existing
def list_pending(self, limit: int = 32) -> list[LedgerWorkItem]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_OUTBOX_COLUMNS} FROM ledger_outbox "
"WHERE status = 'pending' ORDER BY created_at LIMIT %s",
(limit,),
)
return [self._row_to_item(r) for r in cur.fetchall()]
def get(self, correlation_id: str) -> Optional[LedgerWorkItem]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_OUTBOX_COLUMNS} FROM ledger_outbox WHERE correlation_id = %s",
(correlation_id,),
)
row = cur.fetchone()
return self._row_to_item(row) if row else None
def mark_outcome(
self,
correlation_id: str,
outcome: OutboxOutcome,
*,
error: Optional[str] = None,
) -> LedgerWorkItem:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
if outcome is OutboxOutcome.COMPLETED:
status_value = WorkItemStatus.COMPLETED.value
attempts_delta = 1
elif outcome is OutboxOutcome.CONFLICT_BLOCKED:
status_value = WorkItemStatus.CONFLICT_BLOCKED.value
attempts_delta = 1
elif outcome is OutboxOutcome.DUPLICATE_COMPLETED:
status_value = WorkItemStatus.COMPLETED.value
attempts_delta = 0
else: # retryable
status_value = WorkItemStatus.PENDING.value
attempts_delta = 1
cur.execute(
"UPDATE ledger_outbox SET status=%s, attempts=attempts+%s, "
"last_outcome=%s, last_error=%s, updated_at=%s "
"WHERE correlation_id=%s",
(
status_value, attempts_delta, outcome.value, error,
datetime.now(tz=timezone.utc), correlation_id,
),
)
conn.commit()
row = self.get(correlation_id)
assert row is not None
return row
@staticmethod
def _row_to_item(row: tuple) -> LedgerWorkItem:
f = lambda v: float(v) if v is not None else 0.0 # noqa: E731
trade = LedgerTrade(
asset_id=row[1], order_id=row[2], execution_id=row[3], side=row[4],
qty=f(row[5]), price=f(row[6]), fee=f(row[7]), tax=f(row[8]),
executed_at=row[9],
)
return LedgerWorkItem(
correlation_id=row[0],
trade=trade,
status=WorkItemStatus(row[10]),
attempts=int(row[11]) if row[11] is not None else 0,
last_outcome=OutboxOutcome(row[12]) if row[12] else None,
last_error=row[13],
created_at=row[14],
updated_at=row[15],
)