Skip to content

utsavll0/postgres-code-gen

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

psql-code-gen

A PostgreSQL extension that uses AI/LLM to generate and execute SQL queries from natural language descriptions. By default, it operates in read-only mode for safety - allowing only SELECT queries to execute while returning generated queries for modifications without executing them.

Installation

Prerequisites

Environment Variables

You need to set these 3 environment variables and then restart your postgres server:

  1. PSQL_AI_KEY - Your LLM API key

    • For OpenAI: Your OpenAI API key
    • For Anthropic: Your Anthropic API key
  2. PSQL_AI_URL - The OpenAI spec API endpoint URL

    • For OpenAI: https://api.openai.com/v1
    • For Anthropic: https://api.anthropic.com
  3. PSQL_AI_MODEL - The model identifier to use

    • For OpenAI: gpt-4, gpt-4-turbo, gpt-3.5-turbo, etc.
    • For Anthropic: claude-3-opus-20240229, claude-3-sonnet-20240229, etc.

Build & Install

make clean
make
make install

Then create the extension in your database:

CREATE EXTENSION psql_code_gen;

Quick Start

Get up and running in 5 minutes with the sample ecommerce database.

1. Load the Sample Database

The extension includes a complete sample ecommerce schema (sample/ecom.sql) with realistic marketplace data.

# Connect to your PostgreSQL database and load the sample schema
psql -U postgres -d your_database_name -f sample/ecom.sql

Or from within psql:

\i sample/ecom.sql

What Gets Created:

  • Schema: marketplace - Contains all tables and functions
  • Enums: User roles, order statuses, payment statuses, shipment statuses
  • Core Tables:
    • Users (customers, sellers, admins)
    • Products with variants and inventory
    • Orders, order items, payments
    • Shipments, reviews, audit logs
    • Carts and cart items
    • Sellers, addresses, categories, coupons
  • Sample Data:
    • 5 users (2 customers, 2 sellers, 1 admin)
    • 2 products with variants and inventory
    • 1 complete order with payment and shipment records
    • 2 reviews
    • Sample coupons

2. Verify the Installation

Check that the extension and sample data are loaded:

-- List all tables in the marketplace schema
\dt marketplace.*

-- Count records in key tables
SELECT COUNT(*) as user_count FROM marketplace.users;
SELECT COUNT(*) as product_count FROM marketplace.products;
SELECT COUNT(*) as order_count FROM marketplace.orders;

3. Try Your First Query

Start with a simple SELECT query in read-only mode (default):

-- Get all products with their seller info
SELECT * FROM psql_code_gen('show me all products in the marketplace with seller names');

Expected Output:

 command_type | rows_affected |                           result_json
--------------+---------------+---------------------------------------------------
 SELECT       |             1 | {"id":1,"title":"Air Runner Tee",...}
 SELECT       |             2 | {"id":2,"title":"ACME Anvil",...}

4. Extract Specific Information

Query for orders placed in the last 30 days:

SELECT
  result_json->>'order_number' as order_id,
  result_json->>'total' as total,
  result_json->>'status' as status
FROM psql_code_gen('find all orders from the last 30 days')
WHERE command_type = 'SELECT'
LIMIT 10;

5. Try Write Operations (Preview Only)

By default, write operations show the generated SQL without executing:

-- See what SQL would be generated for an insert
SELECT result_json FROM psql_code_gen('add a new product called "Running Shoes" with price 89.99');

Output shows the generated SQL:

 result_json
----------------------------------------------
 "INSERT INTO products (title, price) ..."

6. Enable Write Mode (Optional)

To actually execute write operations:

-- Enable write mode
SELECT psql_code_gen_set_readonly(false);

-- Now this will actually INSERT
SELECT * FROM psql_code_gen('add a new user with email newuser@example.com');

-- Check the result
SELECT COUNT(*) as user_count FROM marketplace.users;

-- Switch back to read-only mode
SELECT psql_code_gen_set_readonly(true);

Example Queries

Try these natural language queries against the sample database:

Users & Customers

-- List all customers and their email addresses
SELECT result_json FROM psql_code_gen('show me all customer users with their email addresses');

-- Find premium sellers with high ratings
SELECT result_json FROM psql_code_gen('find all sellers with average rating above 4.0');

Products & Inventory

-- Get products by category
SELECT result_json FROM psql_code_gen('list all products in the apparel category');

-- Check inventory levels
SELECT result_json FROM psql_code_gen('show product variants with low inventory less than 10 items');

-- Find discontinued products
SELECT result_json FROM psql_code_gen('find all products where is_active is false');

Orders & Payments

-- Recent orders
SELECT
  result_json->>'order_number' as order_number,
  result_json->>'total' as total
FROM psql_code_gen('get the 10 most recent orders')
WHERE command_type = 'SELECT';

-- Orders by status
SELECT result_json FROM psql_code_gen('count how many orders have status shipped');

-- Payment status
SELECT result_json FROM psql_code_gen('find all failed payment transactions');

Sales Analytics

-- Top sellers by sales
SELECT result_json FROM psql_code_gen('which sellers have the most orders');

-- Revenue by category
SELECT result_json FROM psql_code_gen('calculate total sales by product category');

-- Customer activity
SELECT result_json FROM psql_code_gen('show me the top 5 customers by order count');

Usage

Basic Query - SELECT (Read-Only Mode)

SELECT * FROM psql_code_gen('get all users from the users table');

Output:

 command_type | rows_affected |                           result_json
--------------+---------------+--------------------------------------------------------
 SELECT       |             1 | {"id":1,"email":"user@example.com","name":"John Doe"}
 SELECT       |             2 | {"id":2,"email":"jane@example.com","name":"Jane Smith"}

Query with Field Extraction

Extract specific fields from JSON results:

SELECT
  result_json->>'email' as email,
  result_json->>'name' as name
FROM psql_code_gen('get all users from the users table')
WHERE command_type = 'SELECT';

Write Operations in Read-Only Mode (Default)

When trying to modify data in read-only mode, the extension returns the generated query without executing it:

SELECT command_type, result_json FROM psql_code_gen('create a new user with email test@example.com');

Output:

 command_type |                              result_json
--------------+--------------------------------------------------
 INSERT       | INSERT INTO users (email, created_at) VALUES ...

The query is shown in result_json for your review. Logs also show a WARNING with the generated query.

Enable Write Access

To actually execute INSERT, UPDATE, DELETE, CREATE, DROP, ALTER statements:

SELECT psql_code_gen_set_readonly(false);

-- Now write operations will execute
SELECT * FROM psql_code_gen('insert a new product with name "Widget" and price 29.99');

Output:

 command_type | rows_affected | result_json
--------------+---------------+-------------
 INSERT       |             1 |

Re-enable Read-Only Mode

SELECT psql_code_gen_set_readonly(true);

Function Signatures

psql_code_gen(query text)

Generates and executes (or previews) SQL based on natural language input.

Parameters:

  • query (text): Natural language description of the SQL operation

Returns: SETOF psql_code_gen_result composite type with:

  • command_type (text): Type of SQL command (SELECT, INSERT, UPDATE, DELETE, CREATE, DROP, ALTER)
  • rows_affected (integer): Number of rows affected (for SELECT: row number; for DML/DDL in read-only mode: 0)
  • result_json (json): For SELECT queries, the row data in JSON; for write queries in read-only mode, the generated SQL query

psql_code_gen_set_readonly(boolean)

Control read-only mode.

Parameters:

  • readonly (boolean): true for read-only (default), false to allow write operations

Example:

-- Enable read-only mode (safe)
SELECT psql_code_gen_set_readonly(true);

-- Disable read-only mode (allow writes)
SELECT psql_code_gen_set_readonly(false);

Security Considerations

  1. Read-Only Mode (Default): Write operations are never executed by default, only returned for preview
  2. LLM Injection: The extension passes user queries through an LLM. Ensure your LLM API is properly secured
  3. Schema Information: The extension has access to table definitions to provide context to the LLM
  4. User Permissions: Even in write-enabled mode, standard PostgreSQL permissions apply

Examples

Getting Recent Orders

SELECT
  result_json->>'order_id' as order_id,
  result_json->>'customer_name' as customer,
  result_json->>'total' as total
FROM psql_code_gen('show me recent orders from the last 30 days')
LIMIT 10;

Reviewing Generated Queries

SELECT psql_code_gen_set_readonly(true);

-- Review what query would be generated
SELECT result_json FROM psql_code_gen('update all inactive users to active status');

-- If satisfied, enable writes and re-run
SELECT psql_code_gen_set_readonly(false);
SELECT * FROM psql_code_gen('update all inactive users to active status');

Combining with Other Operations

WITH generated_data AS (
  SELECT
    result_json->>'id' as id,
    result_json->>'email' as email,
    result_json->>'created_at' as created_at
  FROM psql_code_gen('get all premium users')
)
SELECT * FROM generated_data WHERE created_at > NOW() - INTERVAL '7 days';

Logging

The extension logs important events using PostgreSQL's logging system:

  • LOG: LLM response and detected SQL command type
  • WARNING: When write operations are blocked in read-only mode

View logs:

# View PostgreSQL server logs
tail -f /var/log/postgresql/postgresql.log

# Or check PostgreSQL logs in psql
SELECT * FROM pg_read_file('postmaster.log', 0, -1) AS content LIMIT 50;

References

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published