Stage 04.3 - repositories, balance snapshots and environment mode fix

app/src/storage/repositories/balance_snapshots.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from src.storage.session import get_connection
+
+
+class BalanceSnapshotRepository:
+    def add_snapshot(
+        self,
+        *,
+        source: str,
+        payload: dict[str, Any],
+    ) -> None:
+        payload_json = json.dumps(payload, ensure_ascii=False)
+
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    '''
+                    INSERT INTO balance_snapshots (source, payload_json)
+                    VALUES (%s, %s::jsonb)
+                    ''',
+                    (source, payload_json),
+                )
+
+    def count_snapshots(self) -> int:
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute("SELECT COUNT(*) FROM balance_snapshots")
+                row = cursor.fetchone()
+
+        return int(row[0]) if row else 0
+
+    def list_recent_snapshots(self, limit: int = 5) -> list[dict[str, str]]:
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    '''
+                    SELECT created_at, source, payload_json::text
+                    FROM balance_snapshots
+                    ORDER BY created_at DESC, id DESC
+                    LIMIT %s
+                    ''',
+                    (limit,),
+                )
+                rows = cursor.fetchall()
+
+        items: list[dict[str, str]] = []
+        for row in rows:
+            items.append(
+                {
+                    "created_at": str(row[0]),
+                    "source": str(row[1]),
+                    "payload_json": str(row[2]),
+                }
+            )
+
+        return items

app/src/storage/repositories/order_drafts.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from src.storage.session import get_connection
+
+
+class OrderDraftRepository:
+    def add_draft(
+        self,
+        *,
+        symbol: str,
+        side: str,
+        order_type: str,
+        quantity: str,
+        status: str = "draft",
+        payload: dict[str, Any] | None = None,
+    ) -> None:
+        payload_json = json.dumps(payload, ensure_ascii=False) if payload is not None else None
+
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    '''
+                    INSERT INTO order_drafts (symbol, side, order_type, quantity, status, payload_json)
+                    VALUES (%s, %s, %s, %s, %s, %s::jsonb)
+                    ''',
+                    (symbol, side, order_type, quantity, status, payload_json),
+                )
+
+    def list_recent_drafts(self, limit: int = 10) -> list[dict[str, str]]:
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    '''
+                    SELECT created_at, symbol, side, order_type, quantity::text, status
+                    FROM order_drafts
+                    ORDER BY created_at DESC, id DESC
+                    LIMIT %s
+                    ''',
+                    (limit,),
+                )
+                rows = cursor.fetchall()
+
+        items: list[dict[str, str]] = []
+        for row in rows:
+            items.append(
+                {
+                    "created_at": str(row[0]),
+                    "symbol": str(row[1]),
+                    "side": str(row[2]),
+                    "order_type": str(row[3]),
+                    "quantity": str(row[4]),
+                    "status": str(row[5]),
+                }
+            )
+        return items
+
+    def count_drafts(self) -> int:
+        with get_connection() as connection:
+            with connection.cursor() as cursor:
+                cursor.execute("SELECT COUNT(*) FROM order_drafts")
+                row = cursor.fetchone()
+
+        return int(row[0]) if row else 0