Coverage for src/finbot/execution/state.py: 91%

1import sqlite3 

2from dataclasses import dataclass 

3from datetime import UTC, datetime 

4from pathlib import Path 

5 

6from finbot.execution.models import BrokerPosition, Order, OrderSide, OrderState, OrderStatus 

7 

8DEFAULT_EXECUTION_STATE_PATH = Path(".finbot/state/execution.sqlite3") 

9 

10 

11class ExecutionStateError(RuntimeError): 

12 """Base error for persistent execution state""" 

13 

14 

15class DuplicateClientOrderError(ExecutionStateError): 

16 """Raised when a client order ID has already been reserved""" 

17 

18 

19class StoredOrderNotFoundError(ExecutionStateError): 

20 """Raised when an expected persisted order doesn't exist.""" 

21 

22 

23@dataclass(frozen=True) 

24class PersistedOrder: 

25 order: Order 

26 status: OrderStatus 

27 broker_order_id: str | None 

28 created_at: datetime 

29 updated_at: datetime 

30 

31 

32class ExecutionStateStore: 

33 def __init__(self, path: str | Path = DEFAULT_EXECUTION_STATE_PATH): 

34 self.path = Path(path) 

35 self.path.parent.mkdir(parents=True, exist_ok=True) 

36 self._initialize() 

37 

38 def _connect(self) -> sqlite3.Connection: 

39 connection = sqlite3.connect(self.path) 

40 connection.row_factory = sqlite3.Row 

41 return connection 

42 

43 def _initialize(self) -> None: 

44 with self._connect() as conn: 

45 conn.execute( 

46 """ 

47 create table if not exists orders 

48 ( 

49 client_order_id text primary key, 

50 broker_order_id text unique, 

51 symbol text not null, 

52 side text not null, 

53 quantity real not null, 

54 status text not null, 

55 created_at text not null, 

56 updated_at text not null 

57 ) 

58 """, 

59 ) 

60 conn.execute( 

61 """ 

62 create table if not exists positions 

63 ( 

64 symbol text primary key, 

65 quantity real not null, 

66 average_entry_price real not null, 

67 updated_at text not null 

68 ) 

69 """ 

70 ) 

71 

72 def reserve(self, order: Order) -> PersistedOrder: 

73 client_order_id = order.client_order_id 

74 if client_order_id is None: 

75 raise ValueError("Client order ID required for persistent execution state") 

76 

77 now = datetime.now(UTC).isoformat() 

78 

79 try: 

80 with self._connect() as conn: 

81 conn.execute( 

82 """ 

83 insert into orders (client_order_id, broker_order_id, symbol, side, quantity, 

84 status, created_at, updated_at) 

85 values (?, null, ?, ?, ?, ?, ?, ?) 

86 """, 

87 ( 

88 client_order_id, order.symbol, order.side.value, order.quantity, 

89 OrderStatus.PENDING.value, now, now, 

90 ), 

91 ) 

92 except sqlite3.IntegrityError as exc: 

93 raise DuplicateClientOrderError( 

94 f"Client order ID {client_order_id!r} has already been reserved.", 

95 ) from exc 

96 

97 persisted = self.get_by_client_order_id(client_order_id) 

98 assert persisted is not None 

99 return persisted 

100 

101 def update(self, state: OrderState) -> PersistedOrder: 

102 client_order_id = state.order.client_order_id 

103 if client_order_id is None: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true

104 raise ValueError("Client order ID required for persistent execution state") 

105 

106 now = datetime.now(UTC).isoformat() 

107 

108 try: 

109 with self._connect() as conn: 

110 cursor = conn.execute( 

111 """ 

112 update orders 

113 set broker_order_id = ?, 

114 status = ?, 

115 updated_at = ? 

116 where client_order_id = ? 

117 """, 

118 (state.order_id, state.status.value, now, client_order_id), 

119 ) 

120 except sqlite3.IntegrityError as exc: 

121 raise ExecutionStateError( 

122 f"Broker order ID {state.order_id!r} is already associated with another order", 

123 ) from exc 

124 

125 if cursor.rowcount != 1: 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true

126 raise StoredOrderNotFoundError( 

127 f"Couldn't find order ID {client_order_id!r} in execution state", 

128 ) 

129 

130 persisted = self.get_by_client_order_id(client_order_id) 

131 assert persisted is not None 

132 return persisted 

133 

134 def get_by_client_order_id(self, client_order_id: str) -> PersistedOrder | None: 

135 with self._connect() as conn: 

136 row = conn.execute( 

137 """ 

138 select * 

139 from orders 

140 where client_order_id = ? 

141 """, 

142 (client_order_id,), 

143 ).fetchone() 

144 

145 return self._from_row(row) 

146 

147 def get_by_broker_order_id(self, broker_order_id: str) -> PersistedOrder | None: 

148 with self._connect() as conn: 

149 row = conn.execute( 

150 """ 

151 select * 

152 from orders 

153 where broker_order_id = ? 

154 """, 

155 (broker_order_id,), 

156 ).fetchone() 

157 

158 return self._from_row(row) 

159 

160 def list_reconcilable(self) -> list[PersistedOrder]: 

161 statuses = ( 

162 OrderStatus.PENDING.value, 

163 OrderStatus.OPEN.value, 

164 OrderStatus.PARTIALLY_FILLED.value, 

165 ) 

166 

167 with self._connect() as conn: 

168 rows = conn.execute( 

169 """ 

170 select * 

171 from orders 

172 where status in (?, ?, ?) 

173 order by created_at, client_order_id 

174 """, 

175 statuses, 

176 ).fetchall() 

177 

178 return [persisted for row in rows if (persisted := self._from_row(row)) is not None] 

179 

180 @staticmethod 

181 def _from_row(row: sqlite3.Row | None) -> PersistedOrder | None: 

182 if row is None: 182 ↛ 183line 182 didn't jump to line 183 because the condition on line 182 was never true

183 return None 

184 

185 return PersistedOrder( 

186 order=Order( 

187 symbol=row["symbol"], 

188 side=OrderSide(row["side"]), 

189 quantity=float(row["quantity"]), 

190 client_order_id=row["client_order_id"], 

191 ), 

192 status=OrderStatus(row["status"]), 

193 broker_order_id=row["broker_order_id"], 

194 created_at=datetime.fromisoformat(row["created_at"]), 

195 updated_at=datetime.fromisoformat(row["updated_at"]), 

196 ) 

197 

198 def list_positions(self) -> list[BrokerPosition]: 

199 with self._connect() as conn: 

200 rows = conn.execute( 

201 """ 

202 select symbol, quantity, average_entry_price 

203 from positions 

204 order by symbol 

205 """, 

206 ).fetchall() 

207 

208 return [ 

209 BrokerPosition( 

210 symbol=row["symbol"], 

211 quantity=float(row["quantity"]), 

212 average_entry_price=float(row["average_entry_price"]), 

213 ) 

214 for row in rows 

215 ] 

216 

217 def replace_positions(self, positions: list[BrokerPosition]) -> None: 

218 now = datetime.now(UTC).isoformat() 

219 

220 with self._connect() as conn: 

221 conn.execute("delete from positions") 

222 conn.executemany( 

223 """ 

224 insert into positions (symbol, quantity, average_entry_price, updated_at) 

225 values (?, ?, ?, ?) 

226 """, 

227 [ 

228 (position.symbol, position.quantity, position.average_entry_price, now) 

229 for position in positions 

230 ], 

231 )