db_schema_extractor.py

#!/usr/bin/env python3
"""
DB Schema Extractor

This script connects to a PostgreSQL or CockroachDB database using credentials from a config file,
extracts the complete database schema (tables, columns, indexes, views, constraints, etc.),
and saves it to a file for use with LLMs like Claude or ChatGPT.
"""

import argparse
from config_manager import load_config, connect_to_db
from schema_generator import generate_schema_definition, schema_to_text


def main():
    parser = argparse.ArgumentParser(
        description="Extract database schema for LLM context."
    )
    parser.add_argument(
        "-c", "--config", required=True, help="Path to the configuration file"
    )
    parser.add_argument("-o", "--output", required=True, help="Output file path")
    parser.add_argument(
        "-s", "--schema", default="public", help="Database schema (default: public)"
    )
    parser.add_argument(
        "-f",
        "--format",
        choices=["json", "text"],
        default="text",
        help="Output format (default: text)",
    )
    parser.add_argument(
        "--db-type",
        choices=["postgres", "cockroachdb"],
        help="Database type (auto-detected by default)",
    )
    args = parser.parse_args()

    # Load configuration
    try:
        config = load_config(args.config)
    except Exception as e:
        print(f"Error loading config file: {e}")
        return 1

    # Connect to the database
    conn = connect_to_db(config)
    if not conn:
        return 1

    try:
        # Generate schema definition
        schema_info = generate_schema_definition(conn, args.schema)

        # Save to file
        with open(args.output, "w") as f:
            if args.format == "json":
                import json

                json.dump(schema_info, f, indent=2)
            else:  # text format
                f.write(schema_to_text(schema_info))

        print(f"Schema extracted successfully to {args.output}")
    except Exception as e:
        print(f"Error generating schema: {e}")
        return 1
    finally:
        conn.close()

    return 0


if __name__ == "__main__":
    exit(main())