Coverage for src/typedal/cli.py: 82%

125 statements  

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

1""" 

2Typer CLI for TypeDAL. 

3""" 

4 

5import sys 

6import typing 

7import warnings 

8from pathlib import Path 

9from typing import Optional 

10 

11import tomli 

12from configuraptor import asdict 

13from configuraptor.alias import is_alias 

14from configuraptor.helpers import is_optional 

15 

16from .types import AnyDict 

17 

18try: 

19 import edwh_migrate 

20 import pydal2sql # noqa: F401 

21 import questionary 

22 import rich 

23 import tomlkit 

24 import typer 

25 from tabulate import tabulate 

26except ImportError as e: # pragma: no cover 

27 # ImportWarning is hidden by default 

28 warnings.warn( 

29 "`migrations` extra not installed. Please run `pip install typedal[migrations]` to fix this.", 

30 source=e, 

31 category=RuntimeWarning, 

32 ) 

33 exit(127) # command not found 

34 

35from pydal2sql.typer_support import IS_DEBUG, with_exit_code 

36from pydal2sql.types import ( 

37 DBType_Option, 

38 OptionalArgument, 

39 OutputFormat_Option, 

40 Tables_Option, 

41) 

42from pydal2sql_core import core_alter, core_create 

43from typing_extensions import Never 

44 

45from . import caching 

46from .__about__ import __version__ 

47from .config import TypeDALConfig, _fill_defaults, load_config, transform 

48from .core import TypeDAL 

49 

50app = typer.Typer( 

51 no_args_is_help=True, 

52) 

53 

54questionary_types: dict[typing.Hashable, Optional[AnyDict]] = { 

55 str: { 

56 "type": "text", 

57 "validate": lambda text: True if len(text) > 0 else "Please enter a value", 

58 }, 

59 Optional[str]: { 

60 "type": "text", 

61 # no validate because it's optional 

62 }, 

63 bool: { 

64 "type": "confirm", 

65 }, 

66 int: {"type": "text", "validate": lambda text: True if text.isdigit() else "Please enter a number"}, 

67 # specific props: 

68 "dialect": { 

69 "type": "select", 

70 "choices": ["sqlite", "postgres", "mysql"], 

71 }, 

72 "folder": { 

73 "type": "path", 

74 "message": "Database directory:", 

75 "only_directories": True, 

76 # "default": "", 

77 }, 

78 "input": { 

79 "type": "path", 

80 "message": "Python file containing table definitions.", 

81 "file_filter": lambda file: "." not in file or file.endswith(".py"), 

82 }, 

83 "output": { 

84 "type": "path", 

85 "message": "Python file where migrations will be written to.", 

86 "file_filter": lambda file: "." not in file or file.endswith(".py"), 

87 }, 

88 # disabled props: 

89 "pyproject": None, # internal 

90 "noop": None, # only for debugging 

91 "connection": None, # internal 

92 "migrate": None, # will probably conflict 

93 "fake_migrate": None, # only enable via config if required 

94} 

95 

96T = typing.TypeVar("T") 

97 

98notfound = object() 

99 

100 

101def _get_question(prop: str, annotation: typing.Type[T]) -> Optional[AnyDict]: # pragma: no cover 

102 question = questionary_types.get(prop, notfound) 

103 if question is notfound: 

104 # None means skip the question, notfound means use the type default! 

105 question = questionary_types.get(annotation) # type: ignore 

106 

107 if not question: 

108 return None 

109 # make a copy so the original is not overwritten: 

110 return question.copy() # type: ignore 

111 

112 

113def get_question(prop: str, annotation: typing.Type[T], default: T | None) -> Optional[T]: # pragma: no cover 

114 """ 

115 Generate a question based on a config property and prompt the user for it. 

116 """ 

117 if not (question := _get_question(prop, annotation)): 

118 return default 

119 

120 question["name"] = prop 

121 question["message"] = question.get("message", f"{prop}? ") 

122 default = typing.cast(T, default or question.get("default") or "") 

123 

124 if annotation == int: 

125 default = typing.cast(T, str(default)) 

126 

127 response = questionary.unsafe_prompt([question], default=default)[prop] 

128 return typing.cast(T, response) 

129 

130 

131@app.command() 

132@with_exit_code(hide_tb=IS_DEBUG) 

133def setup( 

134 config_file: typing.Annotated[Optional[str], typer.Option("--config", "-c")] = None, 

135 minimal: bool = False, 

136) -> None: # pragma: no cover 

137 """ 

138 Setup a [tool.typedal] entry in the local pyproject.toml. 

139 """ 

140 # 1. check if [tool.typedal] in pyproject.toml and ask missing questions (excl .env vars) 

141 # 2. else if [tool.migrate] and/or [tool.pydal2sql] exist in the config, ask the user with copied defaults 

142 # 3. else: ask the user every question or minimal questions based on cli arg 

143 

144 config = load_config(config_file) 

145 

146 toml_path = Path(config.pyproject) 

147 

148 if not (config.pyproject and toml_path.exists()): 

149 # no pyproject.toml found! 

150 toml_path = toml_path if config.pyproject else Path("pyproject.toml") 

151 rich.print(f"[blue]Config toml doesn't exist yet, creating {toml_path}[/blue]", file=sys.stderr) 

152 toml_path.touch() 

153 

154 toml_contents = toml_path.read_text() 

155 # tomli has native Python types, tomlkit doesn't but preserves comments 

156 toml_obj: AnyDict = tomli.loads(toml_contents) 

157 

158 if "[tool.typedal]" in toml_contents: 

159 section = toml_obj["tool"]["typedal"] 

160 config.update(**section, _overwrite=True) 

161 

162 if "[tool.pydal2sql]" in toml_contents: 

163 mapping = {"": ""} # <- placeholder 

164 

165 extra_config = toml_obj["tool"]["pydal2sql"] 

166 extra_config = {mapping.get(k, k): v for k, v in extra_config.items()} 

167 extra_config.pop("format", None) # always edwh-migrate 

168 config.update(**extra_config) 

169 

170 if "[tool.migrate]" in toml_contents: 

171 mapping = {"migrate_uri": "database"} 

172 

173 extra_config = toml_obj["tool"]["migrate"] 

174 extra_config = {mapping.get(k, k): v for k, v in extra_config.items()} 

175 

176 config.update(**extra_config) 

177 

178 data = asdict(config, with_top_level_key=False) 

179 data["migrate"] = None # determined based on existence of input/output file. 

180 

181 for prop, annotation in TypeDALConfig.__annotations__.items(): 

182 if is_alias(config.__class__, prop): 

183 # don't store aliases! 

184 data.pop(prop, None) 

185 continue 

186 

187 if minimal and getattr(config, prop, None) not in (None, "") or is_optional(annotation): 

188 # property already present or not required, SKIP! 

189 data[prop] = getattr(config, prop, None) 

190 continue 

191 

192 _fill_defaults(data, prop, data.get(prop)) 

193 default_value = data.get(prop, None) 

194 answer: typing.Any = get_question(prop, annotation, default_value) 

195 

196 if isinstance(answer, str): 

197 answer = answer.strip() 

198 

199 if annotation == bool: 

200 answer = bool(answer) 

201 elif annotation == int: 

202 answer = int(answer) 

203 

204 config.update(**{prop: answer}) 

205 data[prop] = answer 

206 

207 for prop in TypeDALConfig.__annotations__: 

208 transform(data, prop) 

209 

210 with toml_path.open("r") as f: 

211 old_contents: AnyDict = tomlkit.load(f) 

212 

213 if "tool" not in old_contents: 

214 old_contents["tool"] = {} 

215 

216 data.pop("pyproject", None) 

217 data.pop("connection", None) 

218 

219 # ignore any None: 

220 old_contents["tool"]["typedal"] = {k: v for k, v in data.items() if v is not None} 

221 

222 with toml_path.open("w") as f: 

223 tomlkit.dump(old_contents, f) 

224 

225 rich.print(f"[green]Wrote updated config to {toml_path}![/green]") 

226 

227 

228@app.command(name="migrations.generate") 

229@with_exit_code(hide_tb=IS_DEBUG) 

230def generate_migrations( 

231 connection: typing.Annotated[str, typer.Option("--connection", "-c")] = None, 

232 filename_before: OptionalArgument[str] = None, 

233 filename_after: OptionalArgument[str] = None, 

234 dialect: DBType_Option = None, 

235 tables: Tables_Option = None, 

236 magic: Optional[bool] = None, 

237 noop: Optional[bool] = None, 

238 function: Optional[str] = None, 

239 output_format: OutputFormat_Option = None, 

240 output_file: Optional[str] = None, 

241 dry_run: bool = False, 

242) -> bool: 

243 """ 

244 Run pydal2sql based on the typedal config. 

245 """ 

246 # 1. choose CREATE or ALTER based on whether 'output' exists? 

247 # 2. pass right args based on 'config' to function chosen in 1. 

248 generic_config = load_config(connection) 

249 pydal2sql_config = generic_config.to_pydal2sql() 

250 pydal2sql_config.update( 

251 magic=magic, 

252 noop=noop, 

253 tables=tables, 

254 db_type=dialect.value if dialect else None, 

255 function=function, 

256 format=output_format, 

257 input=filename_before, 

258 output=output_file, 

259 ) 

260 

261 if pydal2sql_config.output and Path(pydal2sql_config.output).exists(): 

262 if dry_run: 

263 print("Would run `pyda2sql alter` with config", asdict(pydal2sql_config), file=sys.stderr) 

264 sys.stderr.flush() 

265 

266 return True 

267 else: # pragma: no cover 

268 return core_alter( 

269 pydal2sql_config.input, 

270 filename_after or pydal2sql_config.input, 

271 db_type=pydal2sql_config.db_type, 

272 tables=pydal2sql_config.tables, 

273 noop=pydal2sql_config.noop, 

274 magic=pydal2sql_config.magic, 

275 function=pydal2sql_config.function, 

276 output_format=pydal2sql_config.format, 

277 output_file=pydal2sql_config.output, 

278 ) 

279 else: 

280 if dry_run: 

281 print("Would run `pyda2sql create` with config", asdict(pydal2sql_config), file=sys.stderr) 

282 sys.stderr.flush() 

283 

284 return True 

285 else: # pragma: no cover 

286 return core_create( 

287 filename=pydal2sql_config.input, 

288 db_type=pydal2sql_config.db_type, 

289 tables=pydal2sql_config.tables, 

290 noop=pydal2sql_config.noop, 

291 magic=pydal2sql_config.magic, 

292 function=pydal2sql_config.function, 

293 output_format=pydal2sql_config.format, 

294 output_file=pydal2sql_config.output, 

295 ) 

296 

297 

298@app.command(name="migrations.run") 

299@with_exit_code(hide_tb=IS_DEBUG) 

300def run_migrations( 

301 connection: typing.Annotated[str, typer.Option("--connection", "-c")] = None, 

302 migrations_file: OptionalArgument[str] = None, 

303 db_uri: Optional[str] = None, 

304 db_folder: Optional[str] = None, 

305 schema_version: Optional[str] = None, 

306 redis_host: Optional[str] = None, 

307 migrate_cat_command: Optional[str] = None, 

308 database_to_restore: Optional[str] = None, 

309 migrate_table: Optional[str] = None, 

310 flag_location: Optional[str] = None, 

311 schema: Optional[str] = None, 

312 create_flag_location: Optional[bool] = None, 

313 dry_run: bool = False, 

314) -> bool: 

315 """ 

316 Run edwh-migrate based on the typedal config. 

317 """ 

318 # 1. build migrate Config from TypeDAL config 

319 # 2. import right file 

320 # 3. `activate_migrations` 

321 generic_config = load_config(connection) 

322 migrate_config = generic_config.to_migrate() 

323 

324 migrate_config.update( 

325 migrate_uri=db_uri, 

326 schema_version=schema_version, 

327 redis_host=redis_host, 

328 migrate_cat_command=migrate_cat_command, 

329 database_to_restore=database_to_restore, 

330 migrate_table=migrate_table, 

331 flag_location=flag_location, 

332 schema=schema, 

333 create_flag_location=create_flag_location, 

334 db_folder=db_folder, 

335 migrations_file=migrations_file, 

336 ) 

337 

338 if dry_run: 

339 print("Would run `migrate` with config", asdict(migrate_config), file=sys.stderr) 

340 else: # pragma: no cover 

341 edwh_migrate.console_hook([], config=migrate_config) 

342 return True 

343 

344 

345def tabulate_data(data: dict[str, AnyDict]) -> None: 

346 """ 

347 Print a nested dict of data in a nice, human-readable table. 

348 """ 

349 flattened_data = [] 

350 for key, inner_dict in data.items(): 

351 temp_dict = {"": key} 

352 temp_dict.update(inner_dict) 

353 flattened_data.append(temp_dict) 

354 

355 # Display the tabulated data from the transposed dictionary 

356 print(tabulate(flattened_data, headers="keys")) 

357 

358 

359FormatOptions: typing.TypeAlias = typing.Literal["plaintext", "json", "yaml", "toml"] 

360 

361 

362def get_output_format(fmt: FormatOptions) -> typing.Callable[[AnyDict], None]: 

363 """ 

364 This function takes a format option as input and \ 

365 returns a function that can be used to output data in the specified format. 

366 """ 

367 match fmt: 

368 case "plaintext": 

369 output = tabulate_data 

370 case "json": 

371 

372 def output(_data: AnyDict) -> None: 

373 import json 

374 

375 print(json.dumps(_data, indent=2)) 

376 

377 case "yaml": 

378 

379 def output(_data: AnyDict) -> None: 

380 import yaml 

381 

382 print(yaml.dump(_data)) 

383 

384 case "toml": 

385 

386 def output(_data: AnyDict) -> None: 

387 import tomli_w 

388 

389 print(tomli_w.dumps(_data)) 

390 

391 case _: 

392 options = typing.get_args(FormatOptions) 

393 raise ValueError(f"Invalid format '{fmt}'. Please choose one of {options}.") 

394 

395 return output 

396 

397 

398@app.command(name="cache.stats") 

399@with_exit_code(hide_tb=IS_DEBUG) 

400def cache_stats( 

401 identifier: typing.Annotated[str, typer.Argument()] = "", 

402 connection: typing.Annotated[str, typer.Option("--connection", "-c")] = None, 

403 fmt: typing.Annotated[ 

404 str, typer.Option("--format", "--fmt", "-f", help="plaintext (default) or json") 

405 ] = "plaintext", 

406) -> None: 

407 """ 

408 Collect caching stats. 

409 

410 Examples: 

411 typedal cache.stats 

412 typedal cache.stats user 

413 typedal cache.stats user.3 

414 """ 

415 config = load_config(connection) 

416 db = TypeDAL(config=config, migrate=False, fake_migrate=False) 

417 

418 output = get_output_format(typing.cast(FormatOptions, fmt)) 

419 

420 data: AnyDict 

421 parts = identifier.split(".") 

422 match parts: 

423 case [] | [""]: 

424 # generic stats 

425 data = caching.calculate_stats(db) # type: ignore 

426 case [table]: 

427 # table stats 

428 data = caching.table_stats(db, table) # type: ignore 

429 case [table, row_id]: 

430 # row stats 

431 data = caching.row_stats(db, table, row_id) # type: ignore 

432 case _: 

433 raise ValueError("Please use the format `table` or `table.id` for this command.") 

434 

435 output(data) 

436 

437 # todo: 

438 # - sort by most dependencies 

439 # - sort by biggest data 

440 # - include size for table_stats, row_stats 

441 # - group by table 

442 # - output format (e.g. json) 

443 

444 

445@app.command(name="cache.clear") 

446@with_exit_code(hide_tb=IS_DEBUG) 

447def cache_clear( 

448 connection: typing.Annotated[str, typer.Option("--connection", "-c")] = None, 

449 purge: typing.Annotated[bool, typer.Option("--all", "--purge", "-p")] = False, 

450) -> None: 

451 """ 

452 Clear (expired) items from the cache. 

453 

454 Args: 

455 connection (optional): [tool.typedal.<connection>] 

456 purge (default: no): remove all items, not only expired 

457 """ 

458 config = load_config(connection) 

459 db = TypeDAL(config=config, migrate=False, fake_migrate=False) 

460 

461 if purge: 

462 caching.clear_cache() 

463 print("Emptied cache") 

464 else: 

465 n = caching.clear_expired() 

466 print(f"Removed {n} expired from cache") 

467 

468 db.commit() 

469 

470 

471def version_callback() -> Never: 

472 """ 

473 --version requested! 

474 """ 

475 print(f"pydal2sql Version: {__version__}") 

476 

477 raise typer.Exit(0) 

478 

479 

480def config_callback() -> Never: 

481 """ 

482 --show-config requested. 

483 """ 

484 config = load_config() 

485 

486 print(repr(config)) 

487 

488 raise typer.Exit(0) 

489 

490 

491@app.callback(invoke_without_command=True) 

492def main( 

493 _: typer.Context, 

494 # stops the program: 

495 show_config: bool = False, 

496 version: bool = False, 

497) -> None: 

498 """ 

499 This script can be used to generate the create or alter sql from pydal or typedal. 

500 """ 

501 if show_config: 

502 config_callback() 

503 elif version: 

504 version_callback() 

505 # else: just continue 

506 

507 

508if __name__ == "__main__": # pragma: no cover 

509 app()