1from __future__ import annotations
 2
 3import subprocess
 4import sys
 5import time
 6from typing import TYPE_CHECKING, cast
 7
 8import click
 9
10from plain.cli import register_cli
11
12from ..backups.cli import cli as backups_cli
13from ..db import OperationalError
14from ..db import db_connection as _db_connection
15
16if TYPE_CHECKING:
17    from ..backends.base.base import BaseDatabaseWrapper
18
19    db_connection = cast("BaseDatabaseWrapper", _db_connection)
20else:
21    db_connection = _db_connection
22
23
24@register_cli("db")
25@click.group()
26def cli() -> None:
27    """Database operations"""
28
29
30cli.add_command(backups_cli)
31
32
33@cli.command()
34@click.argument("parameters", nargs=-1)
35def shell(parameters: tuple[str, ...]) -> None:
36    """Open an interactive database shell"""
37    try:
38        db_connection.client.runshell(list(parameters))
39    except FileNotFoundError:
40        # Note that we're assuming the FileNotFoundError relates to the
41        # command missing. It could be raised for some other reason, in
42        # which case this error message would be inaccurate. Still, this
43        # message catches the common case.
44        click.secho(
45            f"You appear not to have the {db_connection.client.executable_name!r} program installed or on your path.",
46            fg="red",
47            err=True,
48        )
49        sys.exit(1)
50    except subprocess.CalledProcessError as e:
51        click.secho(
52            '"{}" returned non-zero exit status {}.'.format(
53                " ".join(e.cmd),
54                e.returncode,
55            ),
56            fg="red",
57            err=True,
58        )
59        sys.exit(e.returncode)
60
61
62@cli.command()
63def wait() -> None:
64    """Wait for the database to be ready"""
65    attempts = 0
66    while True:
67        attempts += 1
68        waiting_for = False
69
70        try:
71            db_connection.ensure_connection()
72        except OperationalError:
73            waiting_for = True
74
75        if waiting_for:
76            if attempts > 1:
77                # After the first attempt, start printing them
78                click.secho(
79                    f"Waiting for database (attempt {attempts})",
80                    fg="yellow",
81                )
82            time.sleep(1.5)
83        else:
84            click.secho("✔ Database ready", fg="green")
85            break