AI - Python

#!/usr/bin/env python3
"""
Load a CSV file into a PostgreSQL table.

Usage example:
python cpt/load_cpt_csv.py --file cpt/ProcedureSearchList_2026.csv --table procedure_search_list_2026 \
  --host localhost --port 5432 --dbname mydb --user myuser --password secret

The script will create the target table with all columns as text (if it doesn't exist),
sanitize header names into safe column identifiers, and then COPY the CSV into Postgres.
"""
import argparse
import csv
import os
import re
import sys
from typing import List

import psycopg2


def sanitize_column(name: str) -> str:
    name = name.strip()
    name = name.lower()
    # replace non-alphanumeric with underscore
    name = re.sub(r"[^a-z0-9]+", "_", name)
    name = name.strip("_")
    if not name:
        name = "col"
    if re.match(r"^[0-9]", name):
        name = "col_" + name
    return name


def unique_names(names: List[str]) -> List[str]:
    out = []
    seen = {}
    for n in names:
        base = n
        i = 1
        while n in seen:
            n = f"{base}_{i}"
            i += 1
        seen[n] = True
        out.append(n)
    return out


def create_table_if_needed(cur, table: str, columns: List[str]):
    cols_sql = ", ".join(f'"{c}" text' for c in columns)
    sql = f"CREATE TABLE IF NOT EXISTS \"{table}\" ({cols_sql});"
    cur.execute(sql)


def main():
    p = argparse.ArgumentParser(description="Load CSV into Postgres table")
    p.add_argument("--file", required=True, help="Path to CSV file")
    p.add_argument("--table", required=True, help="Target table name")
    p.add_argument("--host", default=os.environ.get("PGHOST", "localhost"))
    p.add_argument("--port", default=os.environ.get("PGPORT", "5432"))
    p.add_argument("--dbname", default=os.environ.get("PGDATABASE"))
    p.add_argument("--user", default=os.environ.get("PGUSER"))
    p.add_argument("--password", default=os.environ.get("PGPASSWORD"))
    p.add_argument("--delimiter", default=None, help="CSV delimiter (auto-detect if omitted)")
    args = p.parse_args()

    if not args.dbname or not args.user or not args.password:
        print("Please provide --dbname, --user and --password (or set PGDATABASE/PGUSER/PGPASSWORD env vars)")
        sys.exit(1)

    if not os.path.exists(args.file):
        print(f"CSV file not found: {args.file}")
        sys.exit(1)

    # Detect delimiter and header
    with open(args.file, "r", newline="", encoding="utf-8") as fh:
        sample = fh.read(8192)
        sniffer = csv.Sniffer()
        try:
            dialect = sniffer.sniff(sample)
            delimiter = dialect.delimiter
        except Exception:
            delimiter = args.delimiter or ","
        fh.seek(0)
        reader = csv.reader(fh, delimiter=delimiter)
        try:
            header = next(reader)
        except StopIteration:
            print("CSV file appears empty")
            sys.exit(1)

    sanitized = [sanitize_column(h or f"col{i+1}") for i, h in enumerate(header)]
    sanitized = unique_names(sanitized)

    # Connect to Postgres
    conn_str = (
        f"host={args.host} port={args.port} dbname={args.dbname} user={args.user} password={args.password}"
    )
    conn = psycopg2.connect(conn_str)
    try:
        cur = conn.cursor()
        create_table_if_needed(cur, args.table, sanitized)
        conn.commit()

        cols_list = ", ".join(f'"{c}"' for c in sanitized)
        copy_sql = f"COPY \"{args.table}\" ({cols_list}) FROM STDIN WITH CSV HEADER DELIMITER '{delimiter}'"

        with open(args.file, "r", encoding="utf-8", newline="") as fh:
            cur.copy_expert(copy_sql, fh)
        conn.commit()
        print(f"Loaded {args.file} into {args.table} ({len(sanitized)} columns)")
    finally:
        conn.close()


if __name__ == "__main__":
    main()