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

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

2import os 

3from asyncio import run 

4 

5import typer 

6from rich import print 

7from typer import Typer 

8 

9from kwai.core.db.database import Database 

10from kwai.core.dependencies import container 

11from kwai.core.settings import ENV_SETTINGS_FILE, Settings 

12 

13 

14def check(): 

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

16 if ENV_SETTINGS_FILE not in os.environ: 

17 print( 

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

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

20 ) 

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

22 env_file = os.environ[ENV_SETTINGS_FILE] 

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

24 

25 

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

27 

28 

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

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

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

32 

33 Args: 

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

35 """ 

36 try: 

37 settings = container[Settings] 

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

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

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

41 if password: 

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

43 except Exception as ex: 

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

45 print(ex) 

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

47 

48 

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

50def test(): 

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

52 

53 async def _main(): 

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

55 try: 

56 database = container[Database] 

57 await database.check_connection() 

58 await database.close() 

59 except Exception as ex: 

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

61 print(ex) 

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

63 

64 print( 

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

66 ) 

67 

68 run(_main())