diff --git a/storage.py b/storage.py new file mode 100644 index 0000000..d820bec --- /dev/null +++ b/storage.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import sqlite3 +import time +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class DecisionRow: + ts: float + symbol: str + market_ticker: str + close_time: str + side: str + prob: float + yes_bid: float + yes_ask: float + spread: float + stake_dollars: float + count: int + limit_cents: int + reason: str + + +class Storage: + def __init__(self, path: str): + self.path = path + self._init() + + def _init(self) -> None: + con = sqlite3.connect(self.path) + cur = con.cursor() + cur.execute( + """ + CREATE TABLE IF NOT EXISTS decisions ( + ts REAL, + symbol TEXT, + market_ticker TEXT, + close_time TEXT, + side TEXT, + prob REAL, + yes_bid REAL, + yes_ask REAL, + spread REAL, + stake_dollars REAL, + count INTEGER, + limit_cents INTEGER, + reason TEXT + ) + """ + ) + cur.execute( + """ + CREATE TABLE IF NOT EXISTS orders ( + ts REAL, + market_ticker TEXT, + mode TEXT, + client_order_id TEXT, + order_id TEXT, + status TEXT, + details TEXT + ) + """ + ) + con.commit() + con.close() + + def log_decision(self, row: DecisionRow) -> None: + con = sqlite3.connect(self.path) + cur = con.cursor() + cur.execute( + """ + INSERT INTO decisions VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?) + """, + ( + row.ts, + row.symbol, + row.market_ticker, + row.close_time, + row.side, + row.prob, + row.yes_bid, + row.yes_ask, + row.spread, + row.stake_dollars, + row.count, + row.limit_cents, + row.reason, + ), + ) + con.commit() + con.close() + + def log_order( + self, + market_ticker: str, + mode: str, + client_order_id: str, + order_id: Optional[str], + status: str, + details: str, + ) -> None: + con = sqlite3.connect(self.path) + cur = con.cursor() + cur.execute( + "INSERT INTO orders VALUES (?,?,?,?,?,?,?)", + (time.time(), market_ticker, mode, client_order_id, order_id, status, details), + ) + con.commit() + con.close()