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

135 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-12-18 13:45 +0100

1""" 

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

3""" 

4import os 

5import re 

6import typing 

7import warnings 

8from collections import defaultdict 

9from pathlib import Path 

10from typing import Any, Optional 

11 

12import black.files 

13import tomli 

14from configuraptor import TypedConfig, alias 

15from dotenv import dotenv_values, find_dotenv 

16 

17if typing.TYPE_CHECKING: # pragma: no cover 

18 from edwh_migrate import Config as MigrateConfig 

19 from pydal2sql.typer_support import Config as P2SConfig 

20 

21 

22class TypeDALConfig(TypedConfig): 

23 """ 

24 Unified config for TypeDAL runtime behavior and migration utilities. 

25 """ 

26 

27 # typedal: 

28 database: str 

29 dialect: str 

30 folder: str = "databases" 

31 caching: bool = True 

32 pool_size: int = 0 

33 pyproject: str 

34 connection: str = "default" 

35 

36 # pydal2sql: 

37 input: str = "" # noqa: A003 

38 output: str = "" 

39 noop: bool = False 

40 magic: bool = True 

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

42 function: str = "define_tables" 

43 

44 # edwh-migrate: 

45 # migrate uri = database 

46 database_to_restore: Optional[str] 

47 migrate_cat_command: Optional[str] 

48 schema_version: Optional[str] 

49 redis_host: Optional[str] 

50 migrate_table: str = "typedal_implemented_features" 

51 flag_location: str 

52 create_flag_location: bool = True 

53 schema: str = "public" 

54 

55 # typedal (depends on properties above) 

56 migrate: bool = True 

57 fake_migrate: bool = False 

58 

59 # aliases: 

60 db_uri: str = alias("database") 

61 db_type: str = alias("dialect") 

62 db_folder: str = alias("folder") 

63 

64 def __repr__(self) -> str: 

65 """ 

66 Dump the config to a (fancy) string. 

67 """ 

68 return f"<TypeDAL {self.__dict__}>" 

69 

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

71 """ 

72 Convert the config to the format required by pydal2sql. 

73 """ 

74 from pydal2sql.typer_support import Config, get_pydal2sql_config 

75 

76 if self.pyproject: # pragma: no cover 

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

78 

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

80 # no typedal config, but existing p2s config: 

81 return get_pydal2sql_config(self.pyproject) 

82 

83 return Config.load( 

84 { 

85 "db_type": self.dialect, 

86 "format": "edwh-migrate", 

87 "tables": self.tables, 

88 "magic": self.magic, 

89 "function": self.function, 

90 "input": self.input, 

91 "output": self.output, 

92 "pyproject": self.pyproject, 

93 } 

94 ) 

95 

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

97 """ 

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

99 """ 

100 from edwh_migrate import Config, get_config 

101 

102 if self.pyproject: # pragma: no cover 

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

104 

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

106 # no typedal config, but existing p2s config: 

107 return get_config() 

108 

109 return Config.load( 

110 { 

111 "migrate_uri": self.database, 

112 "schema_version": self.schema_version, 

113 "redis_host": self.redis_host, 

114 "migrate_cat_command": self.migrate_cat_command, 

115 "database_to_restore": self.database_to_restore, 

116 "migrate_table": self.migrate_table, 

117 "flag_location": self.flag_location, 

118 "create_flag_location": self.create_flag_location, 

119 "schema": self.schema, 

120 "db_folder": self.folder, 

121 "migrations_file": self.output, 

122 } 

123 ) 

124 

125 

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

127 """ 

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

129 """ 

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

131 

132 

133def _load_toml(path: str | bool | None = True) -> tuple[str, dict[str, Any]]: 

134 """ 

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

136 

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

138 If it is False, no data is loaded. 

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

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

141 """ 

142 if path is False: 

143 toml_path = None 

144 elif path in (True, None): 

145 toml_path = find_pyproject_toml() 

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

147 toml_path = str(path) 

148 else: 

149 toml_path = find_pyproject_toml(str(path)) 

150 

151 if not toml_path: 

152 # nothing to load 

153 return "", {} 

154 

155 try: 

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

157 data = tomli.load(f) 

158 

159 return toml_path or "", typing.cast(dict[str, Any], data["tool"]["typedal"]) 

160 except Exception as e: 

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

162 return toml_path or "", {} 

163 

164 

165def _load_dotenv(path: str | bool | None = True) -> tuple[str, dict[str, Any]]: 

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

167 if path is False: 

168 dotenv_path = None 

169 fallback_data = {} 

170 elif path in (True, None): 

171 dotenv_path = find_dotenv(usecwd=True) 

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

173 dotenv_path = str(path) 

174 else: 

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

176 

177 if not dotenv_path: 

178 return "", fallback_data 

179 

180 # 1. find everything with TYPEDAL_ prefix 

181 # 2. remove that prefix 

182 # 3. format values if possible 

183 data = dotenv_values(dotenv_path) 

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

185 

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

187 

188 return dotenv_path, typedal_data 

189 

190 

191DB_ALIASES = { 

192 "postgresql": "postgres", 

193 "psql": "postgres", 

194 "sqlite3": "sqlite", 

195} 

196 

197 

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

199 """ 

200 Convert a db dialect alias to the standard name. 

201 """ 

202 return DB_ALIASES.get(db_name, db_name) 

203 

204 

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

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

207 "dialect": lambda data: get_db_for_alias(data["database"].split(":")[0]) 

208 if ":" in data["database"] 

209 else data.get("db_type"), 

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

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

212 "flag_location": lambda data: f"{db_folder}/flags" 

213 if (db_folder := (data.get("folder") or data.get("db_folder"))) 

214 else "/flags", 

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

216} 

217 

218 

219def _fill_defaults(data: dict[str, Any], prop: str, fallback: Any = None) -> None: 

220 default = DEFAULTS.get(prop, fallback) 

221 if callable(default): 

222 default = default(data) 

223 data[prop] = default 

224 

225 

226def fill_defaults(data: dict[str, Any], prop: str) -> None: 

227 """ 

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

229 """ 

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

231 _fill_defaults(data, prop) 

232 

233 

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

235 "database": lambda data: data["database"] 

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

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

238} 

239 

240 

241def transform(data: dict[str, Any], prop: str) -> bool: 

242 """ 

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

244 """ 

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

246 data[prop] = fn(data) 

247 return True 

248 return False 

249 

250 

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

252 """ 

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

254 

255 Args: 

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

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

258 

259 Returns: 

260 str: The string with replaced variable values. 

261 

262 See Also: 

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

264 and ChatGPT 

265 """ 

266 env = defaultdict(lambda: "") 

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

268 env[key.lower()] = value 

269 

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

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

272 

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

274 var_with_default = match.group(1) 

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

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

277 

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

279 

280 

281def expand_env_vars_into_toml_values(toml: dict[str, Any], env: dict[str, Any]) -> None: 

282 """ 

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

284 

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

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

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

288 

289 Args: 

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

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

292 

293 Returns: 

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

295 

296 Notes: 

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

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

299 

300 Example: 

301 toml_data = { 

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

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

304 } 

305 environment = { 

306 'ENV_VAR': 'replaced_value', 

307 'ANOTHER_VAR': 'value_1', 

308 'YET_ANOTHER_VAR': 'value_2' 

309 } 

310 

311 expand_env_vars_into_toml_values(toml_data, environment) 

312 # 'toml_data' will be modified in place: 

313 # { 

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

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

316 # } 

317 """ 

318 if not toml or not env: 

319 return 

320 

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

322 if isinstance(var, dict): 

323 expand_env_vars_into_toml_values(var, env) 

324 elif isinstance(var, list): 

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

326 elif isinstance(var, str): 

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

328 else: 

329 # nothing to substitute 

330 continue 

331 

332 

333def load_config( 

334 _use_pyproject: bool | str | None = True, _use_env: bool | str | None = True, **fallback: Any 

335) -> TypeDALConfig: 

336 """ 

337 Combines multiple sources of config into one config instance. 

338 """ 

339 # load toml data 

340 # load .env data 

341 # combine and fill with fallback values 

342 # load typedal config or fail 

343 toml_path, toml = _load_toml(_use_pyproject) 

344 dotenv_path, dotenv = _load_dotenv(_use_env) 

345 

346 expand_env_vars_into_toml_values(toml, dotenv) 

347 

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

349 connection: dict[str, Any] = (toml.get(connection_name) if connection_name else toml) or {} 

350 

351 combined = connection | dotenv | fallback 

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

353 

354 combined["pyproject"] = toml_path 

355 combined["connection"] = connection_name 

356 

357 for prop in TypeDALConfig.__annotations__: 

358 fill_defaults(combined, prop) 

359 

360 for prop in TypeDALConfig.__annotations__: 

361 transform(combined, prop) 

362 

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