explore Postgres db via CLI

explore.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "psycopg2-binary",
#     "tabulate",
# ]
# ///

import os
import psycopg2
from psycopg2 import Error, sql
from tabulate import tabulate
from conf import DB_CONFIG

# DB_CONFIG = {
#     "user": "...",
#     "password": "...",
#     "host": "...",
#     "port": "...",
#     "database": "...",
# }

PAGE_SIZE = 20

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def clear():
    os.system("clear" if os.name == "posix" else "cls")


def print_header(title: str):
    bar = "=" * 64
    print(bar)
    print(f"  {title}")
    print(bar)


# ---------------------------------------------------------------------------
# DB queries
# ---------------------------------------------------------------------------


def list_tables(cursor) -> list[str]:
    cursor.execute(
        """
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'public'
        ORDER BY table_name;
        """
    )
    return [row[0] for row in cursor.fetchall()]


def get_row_count(cursor, table: str) -> int:
    cursor.execute(sql.SQL("SELECT COUNT(*) FROM {}").format(sql.Identifier(table)))
    return cursor.fetchone()[0]


def get_table_data(cursor, table: str, offset: int = 0, limit: int = PAGE_SIZE):
    cursor.execute(
        sql.SQL("SELECT * FROM {} LIMIT %s OFFSET %s").format(sql.Identifier(table)),
        (limit, offset),
    )
    rows = cursor.fetchall()
    headers = [desc[0] for desc in cursor.description] if cursor.description else []
    return headers, rows


def get_table_schema(cursor, table: str):
    cursor.execute(
        """
        SELECT column_name, data_type, character_maximum_length, is_nullable
        FROM information_schema.columns
        WHERE table_schema = 'public' AND table_name = %s
        ORDER BY ordinal_position;
        """,
        (table,),
    )
    rows = cursor.fetchall()
    headers = ["Column", "Type", "Max Length", "Nullable"]
    return headers, rows


# ---------------------------------------------------------------------------
# Views
# ---------------------------------------------------------------------------


def view_data(cursor, table: str, row_count: int):
    page = 0
    total_pages = max(1, (row_count + PAGE_SIZE - 1) // PAGE_SIZE)

    while True:
        clear()
        offset = page * PAGE_SIZE
        headers, rows = get_table_data(cursor, table, offset=offset, limit=PAGE_SIZE)

        print_header(
            f"Table: {table.upper()}  |  Page {page + 1}/{total_pages}  ({row_count} rows)"
        )

        if rows:
            print(tabulate(rows, headers=headers, tablefmt="psql"))
        else:
            print("\n  (Table is empty)")

        nav = []
        if page > 0:
            nav.append("[p] Prev")
        if page < total_pages - 1:
            nav.append("[n] Next")
        nav.append("[b] Back")
        print("\n  " + "   ".join(nav))

        choice = input("\nChoice: ").strip().lower()
        if choice == "b":
            return
        elif choice == "n" and page < total_pages - 1:
            page += 1
        elif choice == "p" and page > 0:
            page -= 1


def view_schema(cursor, table: str):
    clear()
    headers, rows = get_table_schema(cursor, table)
    print_header(f"Schema: {table.upper()}")
    print()
    if rows:
        print(tabulate(rows, headers=headers, tablefmt="psql"))
    else:
        print("  (No columns found)")
    input("\nPress Enter to go back...")


def table_menu(cursor, table: str):
    while True:
        clear()
        row_count = get_row_count(cursor, table)
        print_header(
            f"Table: {table.upper()}  |  {row_count} row{'s' if row_count != 1 else ''}"
        )
        print("\n  [1] View data")
        print("  [2] View schema")
        print("  [b] Back")

        choice = input("\nChoice: ").strip().lower()
        if choice == "b":
            return
        elif choice == "1":
            view_data(cursor, table, row_count)
        elif choice == "2":
            view_schema(cursor, table)


def tables_menu(cursor):
    while True:
        clear()
        tables = list_tables(cursor)
        print_header(
            f"PostgreSQL Browser — {DB_CONFIG['database']}@{DB_CONFIG['host']}"
        )

        if not tables:
            print("\n  No public tables found.")
            input("\nPress Enter to exit...")
            return

        print(f"\n  {'#':<5} Table")
        print("  " + "-" * 35)
        for i, table in enumerate(tables, 1):
            print(f"  {i:<5} {table}")

        print("\n  [q] Quit")
        choice = input("\nSelect a table number: ").strip().lower()

        if choice == "q":
            return

        try:
            idx = int(choice) - 1
            if 0 <= idx < len(tables):
                table_menu(cursor, tables[idx])
        except ValueError:
            pass


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main():
    conn = None
    cursor = None
    try:
        conn = psycopg2.connect(
            user=DB_CONFIG["user"],
            password=DB_CONFIG["password"],
            host=DB_CONFIG["host"],
            port=DB_CONFIG["port"],
            database=DB_CONFIG["database"],
        )
        cursor = conn.cursor()
        tables_menu(cursor)
    except (Exception, Error) as e:
        print(f"Error connecting to PostgreSQL: {e}")
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


if __name__ == "__main__":
    main()