Coverage for snekql/_runtime_selection.py: 91%

1"""Backend Runtime Adapter selection for Database initialization.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from dataclasses import dataclass 

7from pathlib import Path 

8from typing import Any, Literal 

9 

10from snekql.errors import DatabaseRuntimeError 

11from snekql.mariadb.config import Config as MariaDBConfig 

12from snekql.model import BackendFamily, Table, require_model_backend 

13from snekql.sqlite.config import Config as SQLiteConfig 

14from snekql.storage import SchemaPolicy 

15from snekql.structured_logging import ResolvedStructuredLogger 

16from snekql.validation import NonNegativeFloat, PositiveInt, validate_boundary 

17 

18type RuntimeConfig = SQLiteConfig | MariaDBConfig 

19 

20 

21@validate_boundary(error_type=DatabaseRuntimeError) 

22def _build_legacy_sqlite_config( 

23 *, 

24 acquire_timeout: NonNegativeFloat, 

25 database: Path | Literal[":memory:"], 

26 pool_size: PositiveInt, 

27) -> SQLiteConfig: 

28 """Build an explicit SQLite config for the legacy initializer shape.""" 

29 

30 return SQLiteConfig( 

31 acquire_timeout=acquire_timeout, 

32 database=database, 

33 pool_size=pool_size, 

34 ) 

35 

36 

37@dataclass(frozen=True) 

38class RuntimeSelection: 

39 """Selected backend config and runtime initializer for one Database startup.""" 

40 

41 backend_family: BackendFamily 

42 config: RuntimeConfig 

43 

44 def validate_model_backends(self, models: Sequence[type[Table[Any]]]) -> None: 

45 """Reject models from another backend before runtime startup mutates state.""" 

46 

47 for model in models: 

48 received_backend = require_model_backend(model) 

49 if received_backend != self.backend_family: 

50 msg = ( 

51 f"backend mismatch: expected {self.backend_family} model, " 

52 f"received {received_backend} model {model.__name__}" 

53 ) 

54 raise DatabaseRuntimeError(msg) 

55 

56 async def initialize_runtime( 

57 self, 

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

59 schema_policy: SchemaPolicy, 

60 *, 

61 logger: ResolvedStructuredLogger, 

62 ) -> object: 

63 """Import and initialize the selected Backend Runtime Adapter lazily.""" 

64 

65 if isinstance(self.config, MariaDBConfig): 

66 from snekql.mariadb.runtime import ( # noqa: PLC0415 

67 initialize_runtime as initialize_mariadb_runtime, 

68 ) 

69 

70 return await initialize_mariadb_runtime( 

71 self.config, 

72 models, 

73 schema_policy, 

74 logger=logger, 

75 ) 

76 try: 

77 from snekql.sqlite.runtime import ( # noqa: PLC0415 

78 initialize_runtime as initialize_sqlite_runtime, 

79 ) 

80 except ModuleNotFoundError as error: 

81 if error.name == "aiosqlite": 

82 msg = "SQLite runtime requires the aiosqlite extra; install with snekql[aiosqlite]" 

83 raise DatabaseRuntimeError(msg) from error 

84 raise 

85 

86 return await initialize_sqlite_runtime( 

87 self.config, 

88 models, 

89 schema_policy, 

90 logger=logger, 

91 ) 

92 

93 

94def _selection_from_config(config: RuntimeConfig) -> RuntimeSelection: 

95 backend_family: BackendFamily = ( 

96 "mariadb" if isinstance(config, MariaDBConfig) else "sqlite" 

97 ) 

98 return RuntimeSelection(backend_family=backend_family, config=config) 

99 

100 

101def resolve_runtime_selection( 

102 *, 

103 backend: object | None, 

104 database: Path | Literal[":memory:"] | None, 

105 pool_size: PositiveInt, 

106 acquire_timeout: NonNegativeFloat, 

107) -> RuntimeSelection: 

108 """Resolve public Database.initialize arguments to a Runtime Adapter choice.""" 

109 

110 if backend is not None: 

111 if not isinstance(backend, SQLiteConfig | MariaDBConfig): 

112 msg = "unsupported database backend config" 

113 raise DatabaseRuntimeError(msg) 

114 if database is not None: 

115 msg = "backend config cannot be combined with database" 

116 raise DatabaseRuntimeError(msg) 

117 return _selection_from_config(backend) 

118 if database is None: 

119 msg = "Database.initialize requires a backend config or database" 

120 raise DatabaseRuntimeError(msg) 

121 config = _build_legacy_sqlite_config( 

122 acquire_timeout=acquire_timeout, 

123 database=database, 

124 pool_size=pool_size, 

125 ) 

126 return RuntimeSelection(backend_family="sqlite", config=config)