|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import subprocess |
| 4 | +import sys |
| 5 | +import time |
| 6 | +from collections import defaultdict |
| 7 | + |
| 8 | +import click |
| 9 | +import psycopg |
| 10 | + |
| 11 | +from plain.cli import register_cli |
| 12 | + |
| 13 | +from ..backups.cli import cli as backups_cli |
| 14 | +from ..db import get_connection |
| 15 | +from ..dialect import quote_name |
| 16 | +from ..migrations.recorder import MIGRATION_TABLE_NAME |
| 17 | +from .diagnose import diagnose |
| 18 | + |
| 19 | + |
| 20 | +@register_cli("postgres") |
| 21 | +@click.group() |
| 22 | +def cli() -> None: |
| 23 | + """Postgres operations""" |
| 24 | + |
| 25 | + |
| 26 | +cli.add_command(backups_cli) |
| 27 | +cli.add_command(diagnose) |
| 28 | + |
| 29 | + |
| 30 | +@cli.command() |
| 31 | +@click.argument("parameters", nargs=-1) |
| 32 | +def shell(parameters: tuple[str, ...]) -> None: |
| 33 | + """Open an interactive database shell""" |
| 34 | + conn = get_connection() |
| 35 | + try: |
| 36 | + conn.runshell(list(parameters)) |
| 37 | + except FileNotFoundError: |
| 38 | + # Note that we're assuming the FileNotFoundError relates to the |
| 39 | + # command missing. It could be raised for some other reason, in |
| 40 | + # which case this error message would be inaccurate. Still, this |
| 41 | + # message catches the common case. |
| 42 | + click.secho( |
| 43 | + f"You appear not to have the {conn.executable_name!r} program installed or on your path.", |
| 44 | + fg="red", |
| 45 | + err=True, |
| 46 | + ) |
| 47 | + sys.exit(1) |
| 48 | + except subprocess.CalledProcessError as e: |
| 49 | + click.secho( |
| 50 | + '"{}" returned non-zero exit status {}.'.format( |
| 51 | + " ".join(e.cmd), |
| 52 | + e.returncode, |
| 53 | + ), |
| 54 | + fg="red", |
| 55 | + err=True, |
| 56 | + ) |
| 57 | + sys.exit(e.returncode) |
| 58 | + |
| 59 | + |
| 60 | +@cli.command("drop-unknown-tables") |
| 61 | +@click.option( |
| 62 | + "--yes", |
| 63 | + is_flag=True, |
| 64 | + help="Skip confirmation prompt (for non-interactive use).", |
| 65 | +) |
| 66 | +def drop_unknown_tables(yes: bool) -> None: |
| 67 | + """Drop all tables not associated with a Plain model""" |
| 68 | + conn = get_connection() |
| 69 | + db_tables = set(conn.table_names()) |
| 70 | + model_tables = set(conn.plain_table_names()) |
| 71 | + unknown_tables = sorted(db_tables - model_tables - {MIGRATION_TABLE_NAME}) |
| 72 | + |
| 73 | + if not unknown_tables: |
| 74 | + click.echo("No unknown tables found.") |
| 75 | + return |
| 76 | + |
| 77 | + unknown_set = set(unknown_tables) |
| 78 | + table_count = len(unknown_tables) |
| 79 | + tables_label = f"{table_count} table{'s' if table_count != 1 else ''}" |
| 80 | + |
| 81 | + # Find foreign key constraints from kept tables that reference unknown tables |
| 82 | + cascade_warnings: defaultdict[str, list[tuple[str, str]]] = defaultdict(list) |
| 83 | + with conn.cursor() as cursor: |
| 84 | + for table in unknown_tables: |
| 85 | + cursor.execute( |
| 86 | + """ |
| 87 | + SELECT conname, conrelid::regclass |
| 88 | + FROM pg_constraint |
| 89 | + WHERE confrelid = %s::regclass AND contype = 'f' |
| 90 | + """, |
| 91 | + [table], |
| 92 | + ) |
| 93 | + for constraint_name, referencing_table in cursor.fetchall(): |
| 94 | + if str(referencing_table) not in unknown_set: |
| 95 | + cascade_warnings[table].append( |
| 96 | + (constraint_name, str(referencing_table)) |
| 97 | + ) |
| 98 | + |
| 99 | + click.secho("Unknown tables:", fg="yellow", bold=True) |
| 100 | + for table in unknown_tables: |
| 101 | + click.echo(f" - {table}") |
| 102 | + for constraint_name, referencing_table in cascade_warnings[table]: |
| 103 | + click.secho( |
| 104 | + f" ⚠ CASCADE will drop constraint {constraint_name} on {referencing_table}", |
| 105 | + fg="red", |
| 106 | + ) |
| 107 | + click.echo() |
| 108 | + |
| 109 | + if not yes: |
| 110 | + if not click.confirm(f"Drop {tables_label} (CASCADE)? This cannot be undone."): |
| 111 | + return |
| 112 | + |
| 113 | + with conn.cursor() as cursor: |
| 114 | + for table in unknown_tables: |
| 115 | + click.echo(f" Dropping {table}...", nl=False) |
| 116 | + cursor.execute(f"DROP TABLE IF EXISTS {quote_name(table)} CASCADE") |
| 117 | + click.echo(" OK") |
| 118 | + |
| 119 | + click.secho(f"✓ Dropped {tables_label}.", fg="green") |
| 120 | + |
| 121 | + |
| 122 | +@cli.command() |
| 123 | +def wait() -> None: |
| 124 | + """Wait for the database to be ready""" |
| 125 | + attempts = 0 |
| 126 | + while True: |
| 127 | + attempts += 1 |
| 128 | + waiting_for = False |
| 129 | + |
| 130 | + try: |
| 131 | + get_connection().ensure_connection() |
| 132 | + except psycopg.OperationalError: |
| 133 | + waiting_for = True |
| 134 | + |
| 135 | + if waiting_for: |
| 136 | + if attempts > 1: |
| 137 | + # After the first attempt, start printing them |
| 138 | + click.secho( |
| 139 | + f"Waiting for database (attempt {attempts})", |
| 140 | + fg="yellow", |
| 141 | + ) |
| 142 | + time.sleep(1.5) |
| 143 | + else: |
| 144 | + click.secho("✔ Database ready", fg="green") |
| 145 | + break |
0 commit comments