Coverage for src/typedal/config.py: 100%

134 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-01-29 16:15 +0100

1""" 

2TypeDAL can be configured by a combination of pyproject.toml (static), env (dynamic) and code (programmic). 

3""" 

4 

5import os 

6import re 

7import typing 

8import warnings 

9from collections import defaultdict 

10from pathlib import Path 

11from typing import Any, Optional 

12 

13import black.files 

14import tomli 

15from configuraptor import TypedConfig, alias 

16from dotenv import dotenv_values, find_dotenv 

17 

18from .types import AnyDict 

19 

20if typing.TYPE_CHECKING: # pragma: no cover 

21 from edwh_migrate import Config as MigrateConfig 

22 from pydal2sql.typer_support import Config as P2SConfig 

23 

24 

25class TypeDALConfig(TypedConfig): 

26 """ 

27 Unified config for TypeDAL runtime behavior and migration utilities. 

28 """ 

29 

30 # typedal: 

31 database: str 

32 dialect: str 

33 folder: str = "databases" 

34 caching: bool = True 

35 pool_size: int = 0 

36 pyproject: str 

37 connection: str = "default" 

38 

39 # pydal2sql: 

40 input: str = "" 

41 output: str = "" 

42 noop: bool = False 

43 magic: bool = True 

44 tables: Optional[list[str]] = None 

45 function: str = "define_tables" 

46 

47 # edwh-migrate: 

48 # migrate uri = database 

49 database_to_restore: Optional[str] 

50 migrate_cat_command: Optional[str] 

51 schema_version: Optional[str] 

52 redis_host: Optional[str] 

53 migrate_table: str = "typedal_implemented_features" 

54 flag_location: str 

55 create_flag_location: bool = True 

56 schema: str = "public" 

57 

58 # typedal (depends on properties above) 

59 migrate: bool = True 

60 fake_migrate: bool = False 

61 

62 # aliases: 

63 db_uri: str = alias("database") 

64 db_type: str = alias("dialect") 

65 db_folder: str = alias("folder") 

66 

67 # repr set by @beautify (by inheriting from TypedConfig) 

68 

69 def to_pydal2sql(self) -> "P2SConfig": 

70 """ 

71 Convert the config to the format required by pydal2sql. 

72 """ 

73 from pydal2sql.typer_support import Config, get_pydal2sql_config 

74 

75 if self.pyproject: # pragma: no cover 

76 project = Path(self.pyproject).read_text() 

77 

78 if "[tool.typedal]" not in project and "[tool.pydal2sql]" in project: 

79 # no typedal config, but existing p2s config: 

80 return get_pydal2sql_config(self.pyproject) 

81 

82 return Config.load( 

83 { 

84 "db_type": self.dialect, 

85 "format": "edwh-migrate", 

86 "tables": self.tables, 

87 "magic": self.magic, 

88 "function": self.function, 

89 "input": self.input, 

90 "output": self.output, 

91 "pyproject": self.pyproject, 

92 } 

93 ) 

94 

95 def to_migrate(self) -> "MigrateConfig": 

96 """ 

97 Convert the config to the format required by edwh-migrate. 

98 """ 

99 from edwh_migrate import Config, get_config 

100 

101 if self.pyproject: # pragma: no cover 

102 project = Path(self.pyproject).read_text() 

103 

104 if "[tool.typedal]" not in project and "[tool.migrate]" in project: 

105 # no typedal config, but existing p2s config: 

106 return get_config() 

107 

108 return Config.load( 

109 { 

110 "migrate_uri": self.database, 

111 "schema_version": self.schema_version, 

112 "redis_host": self.redis_host, 

113 "migrate_cat_command": self.migrate_cat_command, 

114 "database_to_restore": self.database_to_restore, 

115 "migrate_table": self.migrate_table, 

116 "flag_location": self.flag_location, 

117 "create_flag_location": self.create_flag_location, 

118 "schema": self.schema, 

119 "db_folder": self.folder, 

120 "migrations_file": self.output, 

121 } 

122 ) 

123 

124 

125def find_pyproject_toml(directory: str | None = None) -> typing.Optional[str]: 

126 """ 

127 Find the project's config toml, looks up until it finds the project root (black's logic). 

128 """ 

129 return black.files.find_pyproject_toml((directory or os.getcwd(),)) 

130 

131 

132def _load_toml(path: str | bool | None = True) -> tuple[str, AnyDict]: 

133 """ 

134 Path can be a file, a directory, a bool or None. 

135 

136 If it is True or None, the default logic is used. 

137 If it is False, no data is loaded. 

138 if it is a directory, the pyproject.toml will be searched there. 

139 If it is a path, that file will be used. 

140 """ 

141 if path is False: 

142 toml_path = None 

143 elif path in (True, None): 

144 toml_path = find_pyproject_toml() 

145 elif Path(str(path)).is_file(): 

146 toml_path = str(path) 

147 else: 

148 toml_path = find_pyproject_toml(str(path)) 

149 

150 if not toml_path: 

151 # nothing to load 

152 return "", {} 

153 

154 try: 

155 with open(toml_path, "rb") as f: 

156 data = tomli.load(f) 

157 

158 return toml_path or "", typing.cast(AnyDict, data["tool"]["typedal"]) 

159 except Exception as e: 

160 warnings.warn(f"Could not load typedal config toml: {e}", source=e) 

161 return toml_path or "", {} 

162 

163 

164def _load_dotenv(path: str | bool | None = True) -> tuple[str, AnyDict]: 

165 fallback_data = {k.lower().removeprefix("typedal_"): v for k, v in os.environ.items()} 

166 if path is False: 

167 dotenv_path = None 

168 fallback_data = {} 

169 elif path in (True, None): 

170 dotenv_path = find_dotenv(usecwd=True) 

171 elif Path(str(path)).is_file(): 

172 dotenv_path = str(path) 

173 else: 

174 dotenv_path = str(Path(str(path)) / ".env") 

175 

176 if not dotenv_path: 

177 return "", fallback_data 

178 

179 # 1. find everything with TYPEDAL_ prefix 

180 # 2. remove that prefix 

181 # 3. format values if possible 

182 data = dotenv_values(dotenv_path) 

183 data |= os.environ # higher prio than .env 

184 

185 typedal_data = {k.lower().removeprefix("typedal_"): v for k, v in data.items()} 

186 

187 return dotenv_path, typedal_data 

188 

189 

190DB_ALIASES = { 

191 "postgresql": "postgres", 

192 "psql": "postgres", 

193 "sqlite3": "sqlite", 

194} 

195 

196 

197def get_db_for_alias(db_name: str) -> str: 

198 """ 

199 Convert a db dialect alias to the standard name. 

200 """ 

201 return DB_ALIASES.get(db_name, db_name) 

202 

203 

204DEFAULTS: dict[str, Any | typing.Callable[[AnyDict], Any]] = { 

205 "database": lambda data: data.get("db_uri") or "sqlite:memory", 

206 "dialect": lambda data: ( 

207 get_db_for_alias(data["database"].split(":")[0]) if ":" in data["database"] else data.get("db_type") 

208 ), 

209 "migrate": lambda data: not (data.get("input") or data.get("output")), 

210 "folder": lambda data: data.get("db_folder"), 

211 "flag_location": lambda data: ( 

212 f"{db_folder}/flags" if (db_folder := (data.get("folder") or data.get("db_folder"))) else "/flags" 

213 ), 

214 "pool_size": lambda data: 1 if data.get("dialect", "sqlite") == "sqlite" else 3, 

215} 

216 

217 

218def _fill_defaults(data: AnyDict, prop: str, fallback: Any = None) -> None: 

219 default = DEFAULTS.get(prop, fallback) 

220 if callable(default): 

221 default = default(data) 

222 data[prop] = default 

223 

224 

225def fill_defaults(data: AnyDict, prop: str) -> None: 

226 """ 

227 Fill missing property defaults with (calculated) sane defaults. 

228 """ 

229 if data.get(prop, None) is None: 

230 _fill_defaults(data, prop) 

231 

232 

233TRANSFORMS: dict[str, typing.Callable[[AnyDict], Any]] = { 

234 "database": lambda data: ( 

235 data["database"] 

236 if (":" in data["database"] or not data.get("dialect")) 

237 else (data["dialect"] + "://" + data["database"]) 

238 ) 

239} 

240 

241 

242def transform(data: AnyDict, prop: str) -> bool: 

243 """ 

244 After the user has chosen a value, possibly transform it. 

245 """ 

246 if fn := TRANSFORMS.get(prop): 

247 data[prop] = fn(data) 

248 return True 

249 return False 

250 

251 

252def expand_posix_vars(posix_expr: str, context: dict[str, str]) -> str: 

253 """ 

254 Replace case-insensitive POSIX and Docker Compose-like environment variables in a string with their values. 

255 

256 Args: 

257 posix_expr (str): The input string containing case-insensitive POSIX or Docker Compose-like variables. 

258 context (dict): A dictionary containing variable names and their respective values. 

259 

260 Returns: 

261 str: The string with replaced variable values. 

262 

263 See Also: 

264 https://stackoverflow.com/questions/386934/how-to-evaluate-environment-variables-into-a-string-in-python 

265 and ChatGPT 

266 """ 

267 env = defaultdict(lambda: "") 

268 for key, value in context.items(): 

269 env[key.lower()] = value 

270 

271 # Regular expression to match "${VAR:default}" pattern 

272 pattern = r"\$\{([^}]+)\}" 

273 

274 def replace_var(match: re.Match[Any]) -> str: 

275 var_with_default = match.group(1) 

276 var_name, default_value = var_with_default.split(":") if ":" in var_with_default else (var_with_default, "") 

277 return env.get(var_name.lower(), default_value) 

278 

279 return re.sub(pattern, replace_var, posix_expr) 

280 

281 

282def expand_env_vars_into_toml_values(toml: AnyDict, env: AnyDict) -> None: 

283 """ 

284 Recursively expands POSIX/Docker Compose-like environment variables in a TOML dictionary. 

285 

286 This function traverses a TOML dictionary and expands POSIX/Docker Compose-like 

287 environment variables (${VAR:default}) using values provided in the 'env' dictionary. 

288 It performs in-place modification of the 'toml' dictionary. 

289 

290 Args: 

291 toml (dict): A TOML dictionary with string values possibly containing environment variables. 

292 env (dict): A dictionary containing environment variable names and their respective values. 

293 

294 Returns: 

295 None: The function modifies the 'toml' dictionary in place. 

296 

297 Notes: 

298 The function recursively traverses the 'toml' dictionary. If a value is a string or a list of strings, 

299 it attempts to substitute any environment variables found within those strings using the 'env' dictionary. 

300 

301 Example: 

302 toml_data = { 

303 'key1': 'This has ${ENV_VAR:default}', 

304 'key2': ['String with ${ANOTHER_VAR}', 'Another ${YET_ANOTHER_VAR}'] 

305 } 

306 environment = { 

307 'ENV_VAR': 'replaced_value', 

308 'ANOTHER_VAR': 'value_1', 

309 'YET_ANOTHER_VAR': 'value_2' 

310 } 

311 

312 expand_env_vars_into_toml_values(toml_data, environment) 

313 # 'toml_data' will be modified in place: 

314 # { 

315 # 'key1': 'This has replaced_value', 

316 # 'key2': ['String with value_1', 'Another value_2'] 

317 # } 

318 """ 

319 if not toml or not env: 

320 return 

321 

322 for key, var in toml.items(): 

323 if isinstance(var, dict): 

324 expand_env_vars_into_toml_values(var, env) 

325 elif isinstance(var, list): 

326 toml[key] = [expand_posix_vars(_, env) for _ in var if isinstance(_, str)] 

327 elif isinstance(var, str): 

328 toml[key] = expand_posix_vars(var, env) 

329 else: 

330 # nothing to substitute 

331 continue 

332 

333 

334def load_config( 

335 connection_name: Optional[str] = None, 

336 _use_pyproject: bool | str | None = True, 

337 _use_env: bool | str | None = True, 

338 **fallback: Any, 

339) -> TypeDALConfig: 

340 """ 

341 Combines multiple sources of config into one config instance. 

342 """ 

343 # load toml data 

344 # load .env data 

345 # combine and fill with fallback values 

346 # load typedal config or fail 

347 toml_path, toml = _load_toml(_use_pyproject) 

348 dotenv_path, dotenv = _load_dotenv(_use_env) 

349 

350 expand_env_vars_into_toml_values(toml, dotenv) 

351 

352 connection_name = connection_name or dotenv.get("connection", "") or toml.get("default", "") 

353 connection: AnyDict = (toml.get(connection_name) if connection_name else toml) or {} 

354 

355 combined = connection | dotenv | fallback 

356 combined = {k.replace("-", "_"): v for k, v in combined.items()} 

357 

358 combined["pyproject"] = toml_path 

359 combined["connection"] = connection_name 

360 

361 for prop in TypeDALConfig.__annotations__: 

362 fill_defaults(combined, prop) 

363 

364 for prop in TypeDALConfig.__annotations__: 

365 transform(combined, prop) 

366 

367 return TypeDALConfig.load(combined, convert_types=True)