sqlq.py

#!/usr/bin/env python3
"""
sqlq -- Query SQLite databases with formatted output. Zero deps.

Quick queries, table listing, schema inspection, CSV/JSON export.
Uses Python's built-in sqlite3 -- nothing to install.

Usage:
    py sqlq.py mydb.sqlite "SELECT * FROM users LIMIT 10"
    py sqlq.py mydb.sqlite --tables                          # List all tables
    py sqlq.py mydb.sqlite --schema users                    # Show table schema
    py sqlq.py mydb.sqlite --info                            # Database overview
    py sqlq.py mydb.sqlite --count users                     # Count rows
    py sqlq.py mydb.sqlite "SELECT * FROM users" --csv       # CSV output
    py sqlq.py mydb.sqlite "SELECT * FROM users" --json      # JSON output
    py sqlq.py mydb.sqlite --export users csv > users.csv    # Export table
    py sqlq.py mydb.sqlite --import data.csv mytable         # Import CSV into table
    py sqlq.py mydb.sqlite --repl                            # Interactive mode
"""

import argparse
import csv
import json
import os
import sqlite3
import sys
from pathlib import Path

# Colors
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
RED = "\033[31m"


def color_supported() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    if sys.platform == "win32":
        return bool(os.environ.get("TERM") or os.environ.get("WT_SESSION"))
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


USE_COLOR = color_supported()


def c(code: str, text: str) -> str:
    return f"{code}{text}{RESET}" if USE_COLOR else text


def fmt_size(n: int) -> str:
    if n < 1024:
        return f"{n} B"
    elif n < 1024 ** 2:
        return f"{n / 1024:.1f} KB"
    elif n < 1024 ** 3:
        return f"{n / 1024**2:.1f} MB"
    return f"{n / 1024**3:.2f} GB"


def print_table(headers: list[str], rows: list[tuple], max_width: int = 40):
    """Print a formatted ASCII table."""
    if not headers:
        return

    # Calculate column widths
    widths = [len(str(h)) for h in headers]
    for row in rows:
        for i, val in enumerate(row):
            if i < len(widths):
                widths[i] = max(widths[i], min(len(str(val) if val is not None else "NULL"), max_width))

    # Header
    header_line = "  " + "  ".join(c(BOLD, str(h).ljust(widths[i])) for i, h in enumerate(headers))
    sep_line = "  " + "  ".join("-" * widths[i] for i in range(len(headers)))
    print(header_line)
    print(c(DIM, sep_line))

    # Rows
    for row in rows:
        cells = []
        for i, val in enumerate(row):
            if i >= len(widths):
                break
            s = str(val) if val is not None else c(DIM, "NULL")
            if len(str(val or "")) > max_width:
                s = str(val)[:max_width - 3] + "..."
            cells.append(s.ljust(widths[i]) if val is not None else s.ljust(widths[i] + len(DIM) + len(RESET) if USE_COLOR else widths[i]))
        print("  " + "  ".join(cells))


def get_tables(conn: sqlite3.Connection) -> list[str]:
    """Get list of table names."""
    cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
    return [row[0] for row in cur.fetchall()]


def get_table_info(conn: sqlite3.Connection, table: str) -> list[dict]:
    """Get column info for a table."""
    cur = conn.execute(f"PRAGMA table_info([{table}])")
    cols = []
    for row in cur.fetchall():
        cols.append({
            "cid": row[0], "name": row[1], "type": row[2],
            "notnull": row[3], "default": row[4], "pk": row[5],
        })
    return cols


def get_row_count(conn: sqlite3.Connection, table: str) -> int:
    cur = conn.execute(f"SELECT COUNT(*) FROM [{table}]")
    return cur.fetchone()[0]


# --- Commands ---

def cmd_query(conn: sqlite3.Connection, sql: str, fmt: str = "table"):
    """Execute a query and display results."""
    try:
        cur = conn.execute(sql)
    except sqlite3.Error as e:
        print(c(RED, f"  SQL Error: {e}"), file=sys.stderr)
        return False

    if cur.description is None:
        # Non-SELECT statement
        conn.commit()
        print(c(GREEN, f"  OK. Rows affected: {cur.rowcount}"))
        return True

    headers = [desc[0] for desc in cur.description]
    rows = cur.fetchall()

    if fmt == "csv":
        writer = csv.writer(sys.stdout, lineterminator="\n")
        writer.writerow(headers)
        writer.writerows(rows)
    elif fmt == "json":
        result = [dict(zip(headers, row)) for row in rows]
        print(json.dumps(result, indent=2, default=str))
    elif fmt == "jsonl":
        for row in rows:
            print(json.dumps(dict(zip(headers, row)), default=str))
    else:
        if not rows:
            print(c(DIM, "  (no results)"))
        else:
            print()
            print_table(headers, rows)
            print(c(DIM, f"\n  {len(rows)} row(s)"))
            print()
    return True


def cmd_tables(conn: sqlite3.Connection):
    """List all tables."""
    tables = get_tables(conn)
    if not tables:
        print(c(DIM, "  No tables found."))
        return

    print(f"\n  {c(BOLD, 'Tables:')}\n")
    for t in tables:
        count = get_row_count(conn, t)
        cols = get_table_info(conn, t)
        col_names = ", ".join(col["name"] for col in cols[:5])
        if len(cols) > 5:
            col_names += f", ... (+{len(cols) - 5})"
        print(f"  {c(CYAN, t):<30}  {count:>8} rows  {c(DIM, col_names)}")
    print()


def cmd_schema(conn: sqlite3.Connection, table: str):
    """Show table schema."""
    cols = get_table_info(conn, table)
    if not cols:
        print(c(RED, f"  Table '{table}' not found."))
        return

    count = get_row_count(conn, table)
    print(f"\n  {c(BOLD, f'Table: {table}')} ({count} rows)\n")
    headers = ["#", "Column", "Type", "Nullable", "Default", "PK"]
    rows = []
    for col in cols:
        nullable = "" if col["notnull"] else "YES"
        pk = "PK" if col["pk"] else ""
        default = str(col["default"]) if col["default"] is not None else ""
        rows.append((col["cid"], col["name"], col["type"] or "ANY", nullable, default, pk))

    print_table(headers, rows)

    # Show indexes
    cur = conn.execute(f"PRAGMA index_list([{table}])")
    indexes = cur.fetchall()
    if indexes:
        print(f"\n  {c(DIM, 'Indexes:')}")
        for idx in indexes:
            unique = "UNIQUE " if idx[2] else ""
            print(f"    {unique}{idx[1]}")

    # Show CREATE statement
    cur = conn.execute(f"SELECT sql FROM sqlite_master WHERE name=?", (table,))
    row = cur.fetchone()
    if row and row[0]:
        print(f"\n  {c(DIM, 'CREATE:')}")
        for line in row[0].splitlines():
            print(f"    {c(DIM, line)}")
    print()


def cmd_info(conn: sqlite3.Connection, db_path: str):
    """Show database overview."""
    tables = get_tables(conn)
    db_size = os.path.getsize(db_path)

    print(f"\n  {c(BOLD, 'Database Info')}")
    print(f"  Path:    {db_path}")
    print(f"  Size:    {fmt_size(db_size)}")
    print(f"  Tables:  {len(tables)}")

    total_rows = 0
    print(f"\n  {'Table':<25}  {'Rows':>10}  {'Columns':>8}")
    print(f"  {'-'*25}  {'-'*10}  {'-'*8}")
    for t in tables:
        count = get_row_count(conn, t)
        cols = len(get_table_info(conn, t))
        total_rows += count
        print(f"  {c(CYAN, t):<{25 + (len(CYAN) + len(RESET) if USE_COLOR else 0)}}  {count:>10,}  {cols:>8}")

    print(f"\n  Total rows: {total_rows:,}\n")


def cmd_export(conn: sqlite3.Connection, table: str, fmt: str):
    """Export a table."""
    cmd_query(conn, f"SELECT * FROM [{table}]", fmt=fmt)


def cmd_import_csv(conn: sqlite3.Connection, csv_path: str, table: str):
    """Import CSV into a table."""
    with open(csv_path, "r", encoding="utf-8", errors="replace", newline="") as f:
        reader = csv.reader(f)
        headers = next(reader)
        rows = list(reader)

    if not headers:
        print(c(RED, "  Empty CSV."))
        return

    # Create table
    cols_def = ", ".join(f"[{h}] TEXT" for h in headers)
    conn.execute(f"CREATE TABLE IF NOT EXISTS [{table}] ({cols_def})")

    # Insert rows
    placeholders = ", ".join(["?"] * len(headers))
    conn.executemany(f"INSERT INTO [{table}] VALUES ({placeholders})", rows)
    conn.commit()

    print(c(GREEN, f"  Imported {len(rows)} rows into '{table}' ({len(headers)} columns)"))


def cmd_repl(conn: sqlite3.Connection):
    """Interactive SQL REPL."""
    print(f"\n  {c(BOLD, 'sqlq')} interactive mode. Type SQL or:")
    print(c(DIM, "  .tables  .schema <table>  .quit"))
    print()

    while True:
        try:
            sql = input(c(CYAN, "  sql> ")).strip()
        except (EOFError, KeyboardInterrupt):
            print(c(DIM, "\n  Bye."))
            break

        if not sql:
            continue
        if sql in (".quit", ".exit", "quit", "exit"):
            print(c(DIM, "  Bye."))
            break
        if sql == ".tables":
            cmd_tables(conn)
            continue
        if sql.startswith(".schema"):
            parts = sql.split(None, 1)
            if len(parts) > 1:
                cmd_schema(conn, parts[1])
            else:
                print(c(DIM, "  Usage: .schema <table>"))
            continue

        cmd_query(conn, sql)


def main():
    parser = argparse.ArgumentParser(
        description="sqlq -- query SQLite databases with formatted output",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("database", help="Path to SQLite database file")
    parser.add_argument("sql", nargs="?", help="SQL query to execute")
    parser.add_argument("--tables", action="store_true", help="List all tables")
    parser.add_argument("--schema", metavar="TABLE", help="Show table schema")
    parser.add_argument("--info", action="store_true", help="Database overview")
    parser.add_argument("--count", metavar="TABLE", help="Count rows in table")
    parser.add_argument("--export", nargs=2, metavar=("TABLE", "FMT"), help="Export table (csv/json/jsonl)")
    parser.add_argument("--import", nargs=2, metavar=("CSV", "TABLE"), dest="import_csv",
                        help="Import CSV file into table")
    parser.add_argument("--csv", action="store_true", help="Output as CSV")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    parser.add_argument("--jsonl", action="store_true", help="Output as JSON Lines")
    parser.add_argument("--repl", action="store_true", help="Interactive mode")

    args = parser.parse_args()

    db_path = args.database
    if not Path(db_path).exists() and not args.import_csv:
        print(f"Error: database not found: {db_path}", file=sys.stderr)
        sys.exit(1)

    try:
        conn = sqlite3.connect(db_path)
    except sqlite3.Error as e:
        print(f"Error: cannot open database: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        if args.tables:
            cmd_tables(conn)
        elif args.schema:
            cmd_schema(conn, args.schema)
        elif args.info:
            cmd_info(conn, db_path)
        elif args.count:
            count = get_row_count(conn, args.count)
            print(count)
        elif args.export:
            cmd_export(conn, args.export[0], args.export[1])
        elif args.import_csv:
            cmd_import_csv(conn, args.import_csv[0], args.import_csv[1])
        elif args.repl:
            cmd_repl(conn)
        elif args.sql:
            fmt = "csv" if args.csv else "json" if args.json else "jsonl" if args.jsonl else "table"
            cmd_query(conn, args.sql, fmt=fmt)
        else:
            # Default: show info
            cmd_info(conn, db_path)
    finally:
        conn.close()


if __name__ == "__main__":
    main()