trex/trex-core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
09b1b4f4c3d78c985796d3275e2b04815ee984f9
trex-core/trex_core/adapters/persistence/postgresql.py
T

876 lines
37 KiB
Python
Raw Blame History

"""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 collections.abc import Iterator
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, 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.application.strategy_state import (
_STATE_CAP_BYTES,
SetResult,
StrategyStateEntry,
)
from trex_core.domain.slot import Slot, SlotNotFoundError, SlotState, SlotTransitionRecord
_SLOT_COLUMNS = (
"slot_id, status, symbol, qty, avg_price, limit_price, "
"order_id, strategy_id, broker_order_id, portfolio_id, "
"sell_strategy_id, forced_close_at, stop_loss_price, take_profit_price, "
"created_at, updated_at, instance_key, retired_at, "
"timeout_sec, slippage_pct, occupancy_seq, venue, "
"cancel_state, cancel_requested_at, order_submitted_at, order_date, cancel_attempts"
)
_KST = timezone(timedelta(hours=9))
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 list_slots(self, instance_key: str, status_filter: Optional[set[str]] = None) -> list[Slot]:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
if status_filter:
# issue #25: filter by status at the DB instead of fetching
# every row and filtering in Python. ANY(%s) binds the set
# as a single PG array parameter.
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot "
"WHERE instance_key = %s AND status = ANY(%s)",
(instance_key, list(status_filter)),
)
else:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot WHERE instance_key = %s",
(instance_key,),
)
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
# --- asset-centric (dcy-020 / BAT-028) ---
def find_slot_by_portfolio(self, instance_key: str, portfolio_id: str) -> 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 instance_key = %s AND portfolio_id = %s LIMIT 1",
(instance_key, portfolio_id),
)
row = cur.fetchone()
return self._row_to_slot(row) if row else None
def count_active(self, instance_key: str) -> int:
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT count(*) FROM daily_slot WHERE instance_key = %s "
"AND status IN ('holding','pending_buy','pending_sell')",
(instance_key,),
)
row = cur.fetchone()
return int(row[0]) if row else 0
def list_holding_slots(self, instance_key: str) -> 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 instance_key = %s AND status = 'holding'",
(instance_key,),
)
return [self._row_to_slot(r) for r in cur.fetchall()]
def find_managed_slot_by_symbol(self, instance_key: str, symbol: str) -> 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 instance_key = %s AND symbol = %s "
"ORDER BY updated_at DESC NULLS LAST, slot_id ASC LIMIT 1",
(instance_key, symbol),
)
row = cur.fetchone()
return self._row_to_slot(row) if row else None
def reconcile_universe(
self,
instance_key: str,
universe: dict[str, str],
*,
allow_retire: bool = True,
) -> dict[str, int]:
now = datetime.now(tz=timezone.utc)
created = retired = tombstoned = 0
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT portfolio_id, status, retired_at FROM daily_slot "
"WHERE instance_key = %s AND portfolio_id IS NOT NULL",
(instance_key,),
)
existing = {r[0]: (r[1], r[2]) for r in cur.fetchall()}
# backfill symbol on slots that were created before this field was populated
cur.execute(
"SELECT portfolio_id FROM daily_slot "
"WHERE instance_key = %s AND portfolio_id IS NOT NULL AND symbol IS NULL",
(instance_key,),
)
null_symbol_pids = {r[0] for r in cur.fetchall()}
for pid in null_symbol_pids:
if pid in universe:
cur.execute(
"UPDATE daily_slot SET symbol = %s, updated_at = %s "
"WHERE instance_key = %s AND portfolio_id = %s AND symbol IS NULL",
(universe[pid], now, instance_key, pid),
)
# create new asset slots (idempotent via partial unique)
for pid, symbol in universe.items():
if pid in existing:
continue
cur.execute(
"INSERT INTO daily_slot "
"(slot_id, status, symbol, instance_key, portfolio_id, "
"created_at, updated_at) VALUES (%s, 'inactive', %s, %s, %s, %s, %s) "
"ON CONFLICT (instance_key, portfolio_id) "
"WHERE portfolio_id IS NOT NULL DO NOTHING",
(uuid4(), symbol, instance_key, pid, now, now),
)
created += cur.rowcount
# retire / tombstone assets that left the universe (skipped under
# the mass-disappearance guard, dcy-020/OQ-2 D3)
if allow_retire:
for pid, (status, retired_at) in existing.items():
if pid in universe:
continue
if status in ("holding", "pending_buy", "pending_sell"):
tombstoned += 1
elif retired_at is None:
cur.execute(
"UPDATE daily_slot SET retired_at = %s, updated_at = %s "
"WHERE instance_key = %s AND portfolio_id = %s",
(now, now, instance_key, pid),
)
retired += 1
conn.commit()
return {"created": created, "retired": retired, "tombstoned": tombstoned}
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, "
"limit_price=%s, order_id=%s, strategy_id=%s, "
"broker_order_id=%s, portfolio_id=%s, sell_strategy_id=%s, "
"forced_close_at=%s, stop_loss_price=%s, take_profit_price=%s, "
"instance_key=%s, retired_at=%s, timeout_sec=%s, slippage_pct=%s, "
"occupancy_seq=%s, venue=%s, cancel_state=%s, "
"cancel_requested_at=%s, order_submitted_at=%s, order_date=%s, "
"cancel_attempts=%s, updated_at=%s WHERE slot_id=%s",
(
slot.status.value,
slot.symbol,
slot.qty,
slot.avg_price,
slot.limit_price,
slot.order_id,
slot.strategy_id,
slot.broker_order_id,
slot.portfolio_id,
slot.sell_strategy_id,
slot.forced_close_at,
slot.stop_loss_price,
slot.take_profit_price,
slot.instance_key,
slot.retired_at,
slot.timeout_sec,
slot.slippage_pct,
slot.occupancy_seq,
slot.venue,
slot.cancel_state,
slot.cancel_requested_at,
slot.order_submitted_at,
slot.order_date,
slot.cancel_attempts,
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()
# --- order execution queue (BAT-004 / REQ-013, option A) ---
def list_undispatched_pending(
self,
instance_key: str = "default",
limit: int = 64,
) -> list[Slot]:
# dcy-020/REQ-056 (BAT-030): instance-scoped so two daemons sharing a
# DB never dispatch each other's orders. List-based (vs single-row
# claim) so one skipped sell (broker flat) cannot head-block the queue.
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
f"SELECT {_SLOT_COLUMNS} FROM daily_slot "
"WHERE instance_key = %s "
"AND status IN ('pending_buy','pending_sell') AND broker_order_id IS NULL "
"ORDER BY updated_at LIMIT %s",
(instance_key, limit),
)
return [self._row_to_slot(r) for r in cur.fetchall()]
def mark_dispatched(
self,
slot_id: UUID,
broker_order_id: str,
*,
venue: Optional[str] = None,
submitted_at: Optional[datetime] = None,
) -> bool:
submitted_at = submitted_at or datetime.now(tz=timezone.utc)
order_date = submitted_at.astimezone(_KST).strftime("%Y%m%d")
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, "
"venue=COALESCE(%s, venue), order_submitted_at=%s, order_date=%s, "
"cancel_attempts=0, updated_at=%s "
"WHERE slot_id=%s AND broker_order_id IS NULL",
(
broker_order_id,
venue,
submitted_at,
order_date,
datetime.now(tz=timezone.utc),
slot_id,
),
)
ok = cur.rowcount == 1
conn.commit()
return ok
def replace_dispatch(
self,
slot_id: UUID,
expected_broker_order_id: str,
broker_order_id: str,
) -> bool:
"""Atomically persist a broker order ID resolved after submit timeout."""
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=%s",
(
broker_order_id,
datetime.now(tz=timezone.utc),
slot_id,
expected_broker_order_id,
),
)
ok = cur.rowcount == 1
conn.commit()
return ok
def mark_cancel_outcome(
self,
slot_id: UUID,
broker_order_id: str,
state: str,
*,
only_if_unset: bool = False,
) -> bool:
now = datetime.now(tz=timezone.utc)
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"UPDATE daily_slot SET cancel_state=%s, "
"cancel_requested_at=%s, updated_at=%s "
"WHERE slot_id=%s AND broker_order_id=%s "
"AND status IN ('pending_buy','pending_sell') "
"AND (%s = FALSE OR cancel_state IS NULL)",
(
state,
now,
now,
slot_id,
broker_order_id,
only_if_unset,
),
)
ok = cur.rowcount == 1
conn.commit()
return ok
def reserve_cancel_attempt(
self,
slot_id: UUID,
broker_order_id: str,
*,
retry_before: Optional[datetime] = None,
) -> bool:
"""Atomically reserve first cancel or an authoritative backoff-gated retry.
Two separate SQL paths avoid nullable-placeholder type-inference errors
(PostgreSQL IndeterminateDatatype when binding None to %s IS NOT NULL).
"""
now = datetime.now(tz=timezone.utc)
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
if retry_before is None:
# First reservation: only match when no cancel has been
# attempted yet (cancel_state IS NULL).
cur.execute(
"UPDATE daily_slot SET cancel_state='unknown', cancel_requested_at=%s, "
"cancel_attempts=cancel_attempts+1, updated_at=%s "
"WHERE slot_id=%s AND broker_order_id=%s "
"AND status IN ('pending_buy','pending_sell') "
"AND cancel_state IS NULL",
(now, now, slot_id, broker_order_id),
)
else:
# Retry: backoff-gated, only match retryable states.
cur.execute(
"UPDATE daily_slot SET cancel_state='unknown', cancel_requested_at=%s, "
"cancel_attempts=cancel_attempts+1, updated_at=%s "
"WHERE slot_id=%s AND broker_order_id=%s "
"AND status IN ('pending_buy','pending_sell') "
"AND cancel_state IN ('rejected','unknown') "
"AND cancel_attempts < 4 "
"AND cancel_requested_at <= %s",
(now, now, slot_id, broker_order_id, retry_before),
)
ok = cur.rowcount == 1
conn.commit()
return ok
def list_dispatched_pending(self, instance_key: str = "default") -> 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 instance_key = %s "
"AND status IN ('pending_buy','pending_sell') AND broker_order_id IS NOT NULL",
(instance_key,),
)
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, cancel_state=NULL, "
"cancel_requested_at=NULL, order_submitted_at=NULL, order_date=NULL, "
"cancel_attempts=0, updated_at=%s WHERE slot_id=%s",
(datetime.now(tz=timezone.utc), slot_id),
)
conn.commit()
def count_slots_by_status(self, instance_key: str) -> 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 instance_key = %s "
"GROUP BY status",
(instance_key,),
)
return {SlotState(r[0]): int(r[1]) for r in cur.fetchall()}
@staticmethod
def _row_to_slot(row: tuple) -> Slot:
from decimal import Decimal
f = lambda v: float(v) if v is not None else None # noqa: E731
d = lambda v: Decimal(str(v)) if v is not None else None # noqa: E731
# Column order: slot_id(0) status(1) symbol(2) qty(3) avg_price(4)
# limit_price(5) order_id(6) strategy_id(7)
# broker_order_id(8) portfolio_id(9) sell_strategy_id(10)
# forced_close_at(11) stop_loss_price(12) take_profit_price(13)
# created_at(14) updated_at(15) instance_key(16) retired_at(17)
# timeout_sec(18) slippage_pct(19) [dcy-022/REQ-068]
# occupancy_seq(20) [dcy-028/REQ-081], venue(21),
# cancel_state(22), cancel_requested_at(23), order_submitted_at(24),
# order_date(25), cancel_attempts(26)
return Slot(
slot_id=row[0] if isinstance(row[0], UUID) else UUID(str(row[0])),
status=SlotState(row[1]),
symbol=row[2],
qty=f(row[3]),
avg_price=f(row[4]),
limit_price=f(row[5]),
order_id=row[6],
strategy_id=row[7],
broker_order_id=row[8],
portfolio_id=row[9],
sell_strategy_id=row[10],
forced_close_at=row[11],
stop_loss_price=d(row[12]),
take_profit_price=d(row[13]),
created_at=row[14],
updated_at=row[15],
instance_key=row[16] if row[16] is not None else "default",
retired_at=row[17],
timeout_sec=int(row[18]) if row[18] is not None else None,
slippage_pct=f(row[19]),
occupancy_seq=int(row[20]) if len(row) > 20 and row[20] is not None else 0,
venue=str(row[21]) if len(row) > 21 and row[21] is not None else "KRX",
cancel_state=str(row[22]) if len(row) > 22 and row[22] is not None else None,
cancel_requested_at=row[23] if len(row) > 23 else None,
order_submitted_at=row[24] if len(row) > 24 else None,
order_date=str(row[25]) if len(row) > 25 and row[25] is not None else None,
cancel_attempts=int(row[26]) if len(row) > 26 and row[26] is not None else 0,
)
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, next_retry_at, venue, order_venue, instance_key"
)
class PostgreSQLLedgerOutbox:
"""ledger_outbox-backed LedgerOutbox. correlation_id PK is the dedup gate."""
def __init__(
self,
dsn: str,
*,
instance_key: str = "default",
claim_ttl_sec: float = 60.0,
worker_id: str | None = None,
clock=None,
psycopg_module: Any = None,
) -> None:
self._psycopg = _require_psycopg(psycopg_module)
self._dsn = dsn
self._instance_key = instance_key
self._claim_ttl_sec = claim_ttl_sec
self._worker_id = worker_id or f"ledger-writer-{uuid4().hex}"
self._clock = clock or (lambda: datetime.now(tz=timezone.utc))
def _now(self) -> datetime:
return self._clock()
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, venue, order_venue, instance_key) "
"VALUES (%s, %s, %s, %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,
trade.venue,
trade.order_venue,
self._instance_key,
),
)
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' "
"AND instance_key = %s "
"AND (next_retry_at IS NULL OR next_retry_at <= now()) "
"AND (claim_expires_at IS NULL OR claim_expires_at <= now()) "
"ORDER BY created_at, correlation_id LIMIT %s",
(self._instance_key, limit),
)
return [self._row_to_item(r) for r in cur.fetchall()]
def claim_pending(
self,
limit: int = 32,
*,
correlation_id: str | None = None,
) -> list[LedgerWorkItem]:
now = self._now()
lease_until = now + timedelta(seconds=self._claim_ttl_sec)
correlation_sql = ""
params: list[Any] = [self._instance_key]
if correlation_id is not None:
correlation_sql = "AND correlation_id = %s "
params.append(correlation_id)
params.extend([limit, self._worker_id, lease_until, now])
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute(
"WITH candidates AS ("
"SELECT correlation_id FROM ledger_outbox "
"WHERE status = 'pending' "
"AND instance_key = %s "
"AND (next_retry_at IS NULL OR next_retry_at <= now()) "
"AND (claim_expires_at IS NULL OR claim_expires_at <= now()) "
f"{correlation_sql}"
"ORDER BY created_at, correlation_id "
"LIMIT %s "
"FOR UPDATE SKIP LOCKED"
") "
"UPDATE ledger_outbox o "
"SET claim_owner = %s, claim_expires_at = %s, updated_at = %s "
"FROM candidates "
"WHERE o.correlation_id = candidates.correlation_id "
f"RETURNING {_OUTBOX_COLUMNS}",
params,
)
rows = [self._row_to_item(r) for r in cur.fetchall()]
conn.commit()
return rows
def pending_cash_delta(self) -> dict[str, float]:
# dcy-018 / REQ-029 (OQ-4): estimated cash delta of all pending rows by
# asset_id. SUM in SQL: sell → +(qty*price - fee - tax), buy → negative.
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT asset_id, "
"SUM(CASE WHEN side='sell' "
" THEN (qty*price - fee - tax) "
" ELSE -(qty*price + fee + tax) END) "
"FROM ledger_outbox "
"WHERE status='pending' AND instance_key=%s "
"GROUP BY asset_id",
(self._instance_key,),
)
return {r[0]: float(r[1]) for r in cur.fetchall() if r[1] is not None}
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,
next_retry_at: Optional[datetime] = 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 in (OutboxOutcome.PERMANENT_FAILURE, OutboxOutcome.DEAD_LETTERED):
# dcy-015 / REQ-019: terminal delivery_failed (reason in last_outcome).
status_value = WorkItemStatus.DELIVERY_FAILED.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
# next_retry_at only persists for a retryable backoff; cleared otherwise.
nra = next_retry_at if outcome is OutboxOutcome.RETRYABLE_FAILURE else None
cur.execute(
"UPDATE ledger_outbox SET status=%s, attempts=attempts+%s, "
"last_outcome=%s, last_error=%s, next_retry_at=%s, updated_at=%s, "
"claim_owner=NULL, claim_expires_at=NULL "
"WHERE correlation_id=%s AND instance_key=%s "
"AND (claim_owner=%s OR claim_owner IS NULL)",
(
status_value,
attempts_delta,
outcome.value,
error,
nra,
self._now(),
correlation_id,
self._instance_key,
self._worker_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],
venue=row[17],
order_venue=row[18],
)
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],
next_retry_at=row[16],
instance_key=row[19],
)
class PgStrategyStateStore:
"""PostgreSQL slot-scoped strategy-instance map store."""
def __init__(self, dsn: str, *, psycopg_module: Any = None) -> None:
self._psycopg = _require_psycopg(psycopg_module)
self._dsn = dsn
def set(
self,
slot_id: UUID,
strategy_instance_id: str,
*,
expected_version: int,
data: str,
slot_active: bool,
) -> SetResult:
if len(data.encode("utf-8")) > _STATE_CAP_BYTES:
return SetResult(ok=False, reason="state_too_large")
now = datetime.now(tz=timezone.utc)
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
try:
with conn.cursor() as cur:
# Lock the lifecycle row before changing state. A close UPDATE
# and this write therefore have one serial order: a writer that
# follows a close observes INACTIVE and cannot recreate state.
cur.execute(
"SELECT status FROM daily_slot WHERE slot_id = %s::uuid FOR UPDATE",
(slot_id,),
)
slot = cur.fetchone()
if slot is None or slot[0] not in {"pending_buy", "holding", "pending_sell"}:
return SetResult(ok=False, reason="slot_inactive")
cur.execute(
"SELECT data FROM slot_strategy_state WHERE slot_id = %s::uuid",
(slot_id,),
)
state_row = cur.fetchone()
state_map = state_row[0] if state_row is not None else {}
entry = state_map.get(strategy_instance_id)
version = int(entry["version"]) if entry is not None else 0
if expected_version != version:
return SetResult(
ok=False,
reason="version_conflict",
version=version,
)
entry_json = "jsonb_build_object('version', %s::bigint, " \
"'data', %s::text, 'updated_at', %s::timestamptz)"
if state_row is None:
cur.execute(
"INSERT INTO slot_strategy_state (slot_id, data) VALUES ("
"%s::uuid, jsonb_build_object(%s::text, " + entry_json + "))",
(slot_id, strategy_instance_id, version + 1, data, now),
)
else:
cur.execute(
"UPDATE slot_strategy_state SET data = jsonb_set(data, "
"ARRAY[%s::text]::text[], " + entry_json + ", true) "
"WHERE slot_id = %s::uuid",
(strategy_instance_id, version + 1, data, now, slot_id),
)
conn.commit()
return SetResult(ok=True, version=version + 1)
except BaseException:
conn.rollback()
raise
def get_for_slots(
self, slot_ids: list[UUID]
) -> dict[UUID, dict[str, StrategyStateEntry]]:
if not slot_ids:
return {}
out: dict[UUID, dict[str, StrategyStateEntry]] = {}
with self._psycopg.connect(self._dsn) as conn:
with conn.cursor() as cur:
cur.execute(
"SELECT slot_id, data "
"FROM slot_strategy_state WHERE slot_id = ANY(%s)",
(list(slot_ids),),
)
for row in cur.fetchall():
sid = row[0] if isinstance(row[0], UUID) else UUID(str(row[0]))
for instance_id, entry in row[1].items():
out.setdefault(sid, {})[instance_id] = StrategyStateEntry(
version=int(entry["version"]),
data=entry["data"],
updated_at=datetime.fromisoformat(entry["updated_at"]),
)
return out
Reference in New Issue View Git Blame Copy Permalink