Coverage for src/kwai/cli/commands/db.py: 0%

1"""db contains subcommands for the database.""" 

2 

3import os 

4 

5from asyncio import run 

6 

7import inject 

8import typer 

9 

10from rich import print 

11from typer import Typer 

12 

13from kwai.core.db.database import Database 

14from kwai.core.settings import ENV_SETTINGS_FILE, get_settings 

15 

16 

17def check(): 

18 """Check if the environment variable is set. If not, stop the cli.""" 

19 if ENV_SETTINGS_FILE not in os.environ: 

20 print( 

21 f"[bold red]Please set env variable {ENV_SETTINGS_FILE} to " 

22 f"the configuration file.[/bold red]" 

23 ) 

24 raise typer.Exit(code=1) from None 

25 env_file = os.environ[ENV_SETTINGS_FILE] 

26 print(f"Settings will be loaded from [bold green]{env_file}[/bold green].") 

27 

28 

29app = Typer(pretty_exceptions_short=True, callback=check) 

30 

31 

32@app.command(help="Show the database settings.") 

33def show(password: bool = typer.Option(False, help="Show the password")): 

34 """Command for showing the active database settings. 

35 

36 Args: 

37 password: show or hide the password (default is hide). 

38 """ 

39 try: 

40 settings = get_settings() 

41 print(f"Host: [bold]{settings.db.host}[/bold]") 

42 print(f"Name: [bold]{settings.db.name}[/bold]") 

43 print(f"User: [bold]{settings.db.user}[/bold]") 

44 if password: 

45 print(f"Password: [bold]{settings.db.password}[/bold]") 

46 except Exception as ex: 

47 print("[bold red]Failed! [/bold red] Could not load the settings!") 

48 print(ex) 

49 raise typer.Exit(code=1) from None 

50 

51 

52@app.command(help="Test the database connection.") 

53def test(): 

54 """Command for testing the database connection.""" 

55 

56 @inject.autoparams() 

57 async def _main(database: Database): 

58 """Closure for handling the async code.""" 

59 try: 

60 await database.check_connection() 

61 await database.close() 

62 except Exception as ex: 

63 print("[bold red]Failed! [/bold red] Could not connect to the database!") 

64 print(ex) 

65 raise typer.Exit(code=1) from None 

66 

67 print( 

68 "[bold green]Success! [/bold green] Connection to the database established!" 

69 ) 

70 

71 run(_main())