SQL Query Generator (Experimental)

Installation

pip install git+ssh://git@github.com/think41/foundation-sql.git#egg=foundation_sql

Usage

from foundation_sql.query import SQLQueryDecorator
query = SQLQueryDecorator(
 # Required Parameters
 schema=<schema_definition>, # Schema definition DDL as script
 db_url=<database_url>, # Database connection URL
 base_url=<base_url>, # Open AI Compatible base URL
 api_key=<model_api_key>, # API Key
 model=<llm_model_name> # Model to use for generation
 # Optional Parameters
 system_prompt=<prompt>, # Override the default system prompt
 cache_dir=<relative_path>, # where to store the sql file - defaults to __sql__
 name=<name> # name of the sql file - defaults to method name
)

@query
def get_users_with_profile() -> List[UserWithProfile]:
 """
 Gets all users with their profiles.
 """
 pass
@query
def create_user_with_profile(user: UserWithProfile) -> int:
 """
 Creates a new user with a profile.
 """
 pass

import os
from typing import List
import unittest
from foundation_sql import db, query
from pydantic import BaseModel
DB_URL = os.environ.get("DATABASE_URL", "sqlite:///:memory:")
class User(BaseModel):
 id: str
 name: str
 email: str
 role: str
TABLES_SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
 id VARCHAR(36) PRIMARY KEY,
 name VARCHAR(255) NOT NULL,
 email VARCHAR(255) NOT NULL UNIQUE,
 role VARCHAR(50) NOT NULL CHECK (role IN ('admin', 'user', 'guest')),
 created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
 updated_at TIMESTAMP
)
"""
query = query.SQLQueryDecorator(schema=TABLES_SCHEMA, 
 db_url=DB_URL, 
 api_key=os.getenv("OPENAI_API_KEY"),
 base_url=os.getenv("OPENAI_API_BASE_URL"),
 model=os.getenv("OPENAI_MODEL"))
@query
def get_users() -> List[User]:
 """
 Gets all users.
 """
 pass
@query
def create_user(user: User) -> int:
 """
 Creates a new user.
 """
 pass
class TestQuery(unittest.TestCase):
 db_url = DB_URL
 schema_sql = TABLES_SCHEMA
 schema_path = None
 def setUp(self):
 """Create a fresh database connection for each test."""
 # Re-initialize the schema for each test to ensure clean state
 if (self.schema_sql or self.schema_path) and self.db_url:
 db.get_db(self.db_url).init_schema(schema_sql=self.schema_sql, schema_path=self.schema_path)
 else:
 raise ValueError("At least one of schema_sql, schema_path must be provided along with db_url")
 
 def test_users(self):
 users = get_users()
 self.assertEqual(len(users), 0)
 
 user = User(id="xxx", name="John Doe", email="john@example.com", role="user")
 create_user(user=user)
 
 users = get_users()
 self.assertEqual(len(users), 1)
 self.assertEqual(users[0], user)
 def tearDown(self):
 """Close the database connection after each test."""
 for _, connection in db.DATABASES.items():
 connection.get_engine().dispose()
 
 db.DATABASES.clear()

#__sql__/create_user.sql
-- def create_user(user: tests.test_simple_query.User) -> int
-- Creates a new user.
-- Expects user.name, user.email and user.role to be defined
INSERT INTO `users` (
 `id`, 
 `name`, 
 `email`, 
 `role`
)
VALUES (
 {{user.id|default(None)}},
 {{user.name}},
 {{user.email}},
 {{user.role}}
);

#__sql__/get_users.sql
-- def get_users() -> List[tests.test_simple_query.User]
-- Gets all users.
SELECT 
 `id` as `id`,
 `name` as `name`,
 `email` as `email`,
 `role` as `role`
FROM 
 `users`

Development Setup

Project Structure