Last active
August 18, 2023 15:40
-
-
Save Tobeyforce/76917dbe4031fd6e02b7c1e98a7d5db8 to your computer and use it in GitHub Desktop.
gpt-database-context-provider
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
# While it's possible to feed an LLM the context of your database, such as chatgpt with interpreter, you can also manually give GPT the structure of your database. | |
# This script will output the structure of a postgresql database which you can feed an LLM so that it can write | |
# accurate SQL-queries for you, or API-design. | |
# 1. Add this execution rights to the script: | |
# chmod +x myscript.sh | |
# 2. Move it to /usr/local/bin/ | |
# 3. Run it from the terminal. It will prompt for a database and a password to the database, and attempt to connect. | |
# It will then output all the database tables and their relationships, | |
# so that you can add this as a context to gpt for future queries. | |
import psycopg2 | |
import getpass | |
# Prompt the user for the database name and password | |
db_name = input("Enter the database name: ") | |
db_password = getpass.getpass("Enter the database password: ") | |
# Connect to the database | |
conn = psycopg2.connect( | |
host="localhost", | |
database=db_name, | |
user="postgres", | |
password=db_password | |
) | |
# Open a cursor to perform database operations | |
cur = conn.cursor() | |
# Get the list of tables in the database | |
cur.execute(""" | |
SELECT table_name | |
FROM information_schema.tables | |
WHERE table_schema = 'public' | |
""") | |
tables = [row[0] for row in cur.fetchall()] | |
# Loop through each table and get its schema information | |
for table in tables: | |
print(f"Table: {table}") | |
print("Columns:") | |
cur.execute(""" | |
SELECT column_name, data_type | |
FROM information_schema.columns | |
WHERE table_name = %s | |
""", (table,)) | |
columns = cur.fetchall() | |
for column in columns: | |
print(f" {column[0]} ({column[1]})") | |
print("Constraints:") | |
cur.execute(""" | |
SELECT constraint_name, constraint_type | |
FROM information_schema.table_constraints | |
WHERE table_name = %s | |
""", (table,)) | |
constraints = cur.fetchall() | |
for constraint in constraints: | |
print(f" {constraint[0]} ({constraint[1]})") | |
# Close the cursor and connection | |
cur.close() | |
conn.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment