Files
trex-core/trex_core/adapters/persistence/postgresql.py
T

760 lines
32 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 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.application.strategy_state import (
SetResult,
StrategyStateEntry,
_cas,
)
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"
)
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, 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.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,
) -> 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, "
"venue=COALESCE(%s, venue), updated_at=%s "
"WHERE slot_id=%s AND broker_order_id IS NULL",
(
broker_order_id,
venue,
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,
) -> 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')",
(state, now, now, slot_id, broker_order_id),
)
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, 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)
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,
)
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"
)
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, venue) "
"VALUES (%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,
),
)
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 (next_retry_at IS NULL OR next_retry_at <= now()) "
"ORDER BY created_at LIMIT %s",
(limit,),
)
return [self._row_to_item(r) for r in cur.fetchall()]
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' GROUP BY asset_id"
)
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 "
"WHERE correlation_id=%s",
(
status_value,
attempts_delta,
outcome.value,
error,
nra,
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],
venue=row[17],
)
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],
)
class PgStrategyStateStore:
"""slot_strategy_state-backed StrategyStateStore (dcy-028 / REQ-079).
Stores opaque per-strategy state keyed by (slot_id, strategy_id). The CAS
decision is the same `_cas` used by the in-memory double; occupancy_seq and
slot_active are supplied by the caller (fetched from the slot) so the store
only owns the slot_strategy_state table.
"""
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_id: str,
*,
expected_occupancy_seq: int,
expected_version: int,
data: str,
current_occupancy_seq: int,
slot_active: bool,
) -> SetResult:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
try:
with conn.cursor() as cur:
cur.execute(
"SELECT occupancy_seq, version, data, updated_at "
"FROM slot_strategy_state "
"WHERE slot_id = %s AND strategy_id = %s FOR UPDATE",
(slot_id, strategy_id),
)
row = cur.fetchone()
existing = (
StrategyStateEntry(
strategy_id=strategy_id,
occupancy_seq=int(row[0]),
version=int(row[1]),
data=row[2],
updated_at=row[3],
)
if row is not None
else None
)
result = _cas(
existing,
expected_occupancy_seq=expected_occupancy_seq,
expected_version=expected_version,
data=data,
current_occupancy_seq=current_occupancy_seq,
slot_active=slot_active,
)
if not result.ok:
conn.rollback()
return result
cur.execute(
"INSERT INTO slot_strategy_state "
"(slot_id, strategy_id, occupancy_seq, version, data, updated_at) "
"VALUES (%s, %s, %s, %s, %s, %s) "
"ON CONFLICT (slot_id, strategy_id) DO UPDATE SET "
"occupancy_seq = EXCLUDED.occupancy_seq, "
"version = EXCLUDED.version, data = EXCLUDED.data, "
"updated_at = EXCLUDED.updated_at",
(
slot_id,
strategy_id,
current_occupancy_seq,
result.version,
data,
datetime.now(tz=timezone.utc),
),
)
conn.commit()
return result
except BaseException:
conn.rollback()
raise
def get_for_slots(
self, slot_ids: list[UUID], current_occupancy_seq: dict[UUID, int]
) -> 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, strategy_id, occupancy_seq, version, data, updated_at "
"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]))
# stale-occupancy entries are invisible (dcy-028 FR-3 defense)
if int(row[2]) != current_occupancy_seq.get(sid):
continue
out.setdefault(sid, {})[row[1]] = StrategyStateEntry(
strategy_id=row[1],
occupancy_seq=int(row[2]),
version=int(row[3]),
data=row[4],
updated_at=row[5],
)
return out
def delete_all(self, slot_id: UUID) -> None:
with self._psycopg.connect(self._dsn, autocommit=False) as conn:
with conn.cursor() as cur:
cur.execute("DELETE FROM slot_strategy_state WHERE slot_id = %s", (slot_id,))
conn.commit()