LLM Based SQL Writing bot

schema_extractor.py

import psycopg2
from psycopg2 import sql


def get_table_structure(db_name, user, password, host='localhost', port='5432'):
    try:
        # Connect to the PostgreSQL database
        connection = psycopg2.connect(
            dbname=db_name,
            user=user,
            password=password,
            host=host,
            port=port
        )
        cursor = connection.cursor()

        # Query to get all table names in the public schema
        cursor.execute("""
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public';
        """)
        tables = cursor.fetchall()

        # Iterate over tables and fetch column details
        for table in tables:
            table_name = table[0]
            print(f"Table: {table_name}")

            # Query to get column details
            cursor.execute(sql.SQL("""
                SELECT
                    column_name,
                    data_type,
                    is_nullable,
                    column_default
                FROM
                    information_schema.columns
                WHERE
                    table_name = %s
                    AND table_schema = 'public';
            """), [table_name])

            columns = cursor.fetchall()

            # Print column details
            for column in columns:
                col_name, data_type, is_nullable, col_default = column
                nullable = "NULLABLE" if is_nullable == 'YES' else "NOT NULL"
                default_val = f"DEFAULT {col_default}" if col_default else ""
                print(f" - {col_name} ({data_type}) {nullable} {default_val}")

            print("\n\n")

    except (Exception, psycopg2.Error) as error:
        print(f"Error while connecting to PostgreSQL: {error}")
    finally:
        # Close the database connection
        if connection:
            cursor.close()
            connection.close()


# Example usage:
get_table_structure(
    db_name='<db_name>',
    user='<db_name>',
    password='<password>',
    host='localhost',
    port='5432'
)

Once you print the table schema, train your LLM for your choice by giving follow instruction:

You are SQL writing bot who writes SQL for Postgres with Database schema given below:

{DB Schema}

I prefer Google AI studio or ChatGPT for this