Coverage for snekql/sqlite/runtime.py: 98%

1"""SQLite adapter for the backend-neutral query runtime.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from typing import Any, Literal, cast 

7 

8from aiosqlite import Connection, Cursor 

9 

10from snekql._pool import ( 

11 SQLiteConnectionPool, 

12 close_sqlite_connection, 

13 normalize_sqlite_database, 

14 open_sqlite_connection, 

15) 

16from snekql.errors import DatabaseRuntimeError 

17from snekql.model import Table 

18from snekql.query import AnySelectQuery, compile_select_sql, compile_write_sql 

19from snekql.query import materialize_select_row as materialize_sqlite_select_row 

20from snekql.schema import ( 

21 initialize_sqlite_schema, 

22 validate_schema_models, 

23 validate_schema_policy, 

24) 

25from snekql.sqlite.config import Config 

26from snekql.storage import SchemaPolicy 

27from snekql.structured_logging import ResolvedStructuredLogger 

28from snekql.validation import NonNegativeFloat 

29 

30 

31class SQLiteCursorAdapter: 

32 """Runtime cursor adapter backed by an aiosqlite cursor.""" 

33 

34 def __init__(self, cursor: Cursor) -> None: 

35 self.cursor: Cursor = cursor 

36 

37 async def fetchone(self) -> Sequence[object] | None: 

38 row = await self.cursor.fetchone() 

39 if row is None: 

40 return None 

41 return cast("Sequence[object]", row) 

42 

43 async def fetchall(self) -> Sequence[Sequence[object]]: 

44 rows = await self.cursor.fetchall() 

45 return [cast("Sequence[object]", row) for row in rows] 

46 

47 async def close(self) -> None: 

48 await self.cursor.close() 

49 

50 

51class SQLiteConnectionAdapter: 

52 """Runtime connection adapter backed by an aiosqlite connection.""" 

53 

54 def __init__(self, connection: Connection) -> None: 

55 self.connection: Connection = connection 

56 

57 async def begin(self) -> None: 

58 await self._execute_control_sql("BEGIN") 

59 

60 async def commit(self) -> None: 

61 await self._execute_control_sql("COMMIT") 

62 

63 async def rollback(self) -> None: 

64 await self._execute_control_sql("ROLLBACK") 

65 

66 async def execute( 

67 self, 

68 sql: str, 

69 params: tuple[object, ...], 

70 ) -> SQLiteCursorAdapter: 

71 cursor = await self.connection.execute(sql, params) 

72 return SQLiteCursorAdapter(cursor) 

73 

74 async def _execute_control_sql(self, sql: str) -> None: 

75 cursor = await self.connection.execute(sql, ()) 

76 try: 

77 return 

78 finally: 

79 await cursor.close() 

80 

81 

82class SQLiteRuntime: 

83 """SQLite adapter satisfying the backend-neutral runtime seam.""" 

84 

85 backend_family: Literal["sqlite"] = "sqlite" 

86 

87 def __init__( 

88 self, 

89 *, 

90 logger: ResolvedStructuredLogger, 

91 acquire_timeout: NonNegativeFloat, 

92 connection_pool: SQLiteConnectionPool, 

93 ) -> None: 

94 self.acquire_timeout: NonNegativeFloat = acquire_timeout 

95 self.connection_pool: SQLiteConnectionPool = connection_pool 

96 self.logger: ResolvedStructuredLogger = logger 

97 

98 async def acquire( 

99 self, 

100 acquisition_timeout: NonNegativeFloat, 

101 ) -> SQLiteConnectionAdapter: 

102 connection = await self.connection_pool.acquire(acquisition_timeout) 

103 return SQLiteConnectionAdapter(connection) 

104 

105 async def release(self, connection: object) -> None: 

106 if not isinstance(connection, SQLiteConnectionAdapter): 

107 msg = "SQLite runtime cannot release a foreign connection" 

108 raise DatabaseRuntimeError(msg) 

109 await self.connection_pool.release(connection.connection) 

110 

111 async def close(self, close_timeout: NonNegativeFloat) -> None: 

112 await self.connection_pool.close(close_timeout) 

113 

114 def check_accepting_work(self) -> None: 

115 self.connection_pool.check_accepting_work() 

116 

117 def compile_select_sql( 

118 self, 

119 query: AnySelectQuery, 

120 ) -> tuple[str, tuple[object, ...]]: 

121 return compile_select_sql(query) 

122 

123 def compile_write_sql(self, query: object) -> tuple[str, tuple[object, ...]]: 

124 return compile_write_sql(query) 

125 

126 def materialize_select_row( 

127 self, 

128 query: AnySelectQuery, 

129 row: Sequence[object], 

130 ) -> object: 

131 return materialize_sqlite_select_row(query, row) 

132 

133 

134async def initialize_runtime( 

135 config: Config, 

136 models: Sequence[type[Table[Any]]], 

137 schema_policy: SchemaPolicy, 

138 *, 

139 logger: ResolvedStructuredLogger, 

140) -> SQLiteRuntime: 

141 """Initialize SQLite connectivity, schema startup, and pool lifecycle.""" 

142 

143 validate_schema_policy(schema_policy) 

144 validate_schema_models(models) 

145 database_path = normalize_sqlite_database(config.database) 

146 logger.debug("sqlite connection opening", database_path=database_path) 

147 connection = await open_sqlite_connection(database_path) 

148 try: 

149 await initialize_sqlite_schema( 

150 connection, 

151 models, 

152 schema_policy, 

153 logger=logger, 

154 ) 

155 except Exception: 

156 await close_sqlite_connection(connection) 

157 raise 

158 return SQLiteRuntime( 

159 logger=logger, 

160 acquire_timeout=config.acquire_timeout, 

161 connection_pool=SQLiteConnectionPool( 

162 logger=logger, 

163 database_path=database_path, 

164 initial_connection=connection, 

165 pool_size=config.pool_size, 

166 ), 

167 ) 

168 

169 

170__all__ = [ 

171 "SQLiteConnectionAdapter", 

172 "SQLiteCursorAdapter", 

173 "SQLiteRuntime", 

174 "initialize_runtime", 

175]