Spaces:
Running
Running
| import os | |
| import tempfile | |
| import sqlite3 | |
| from fastapi import FastAPI, UploadFile, File, Form | |
| from pydantic import BaseModel | |
| from typing import Optional | |
| # Import schema extraction logic | |
| from schema import ( | |
| get_sqlite_schema, | |
| get_uploaded_sqlite_schema, | |
| get_mysql_schema, | |
| get_postgresql_schema, | |
| ) | |
| # Import your AI SQL logic | |
| from sql_analyst import generate_and_execute_sql | |
| app = FastAPI(title="AI SQL Analyst API") | |
| # ---------------------------- | |
| # Health / Root Endpoint | |
| # ---------------------------- | |
| def root(): | |
| return {"status": "AI SQL Analyst API is running"} | |
| # ---------------------------- | |
| # POST /sqlite | |
| # Accept a SQLite database file + user query | |
| # ---------------------------- | |
| async def analyze_sqlite( | |
| file: UploadFile = File(...), | |
| user_query: str = Form(...) | |
| ): | |
| """ | |
| Upload a SQLite file and run a natural‑language query against it. | |
| """ | |
| # Save uploaded file to a temp path | |
| temp_dir = tempfile.gettempdir() | |
| sqlite_path = os.path.join(temp_dir, file.filename) | |
| with open(sqlite_path, "wb") as f: | |
| f.write(await file.read()) | |
| # Extract schema | |
| schema = get_uploaded_sqlite_schema(sqlite_path) | |
| # Build SQLite config | |
| db_info = {"type": "sqlite", "path": sqlite_path} | |
| # Generate + execute SQL | |
| sql, result = generate_and_execute_sql(user_query, schema, db_info) | |
| # Optionally remove the file (cleanup) | |
| try: | |
| os.remove(sqlite_path) | |
| except: | |
| pass | |
| return {"sql": sql, "result": result} | |
| # ---------------------------- | |
| # POST /builtinsqlite | |
| # Run query on built‑in business SQLite | |
| # ---------------------------- | |
| async def analyze_business_sqlite(user_query: str = Form(...)): | |
| """ | |
| Run a query against the built‑in business SQLite database. | |
| """ | |
| # Get business SQLite schema | |
| schema = get_sqlite_schema() | |
| db_info = {"type": "sqlite", "path": "business_database.db"} | |
| sql, result = generate_and_execute_sql(user_query, schema, db_info) | |
| return {"sql": sql, "result": result} | |
| # ---------------------------- | |
| # POST /remote | |
| # Accept remote DB connection details + user query | |
| # ---------------------------- | |
| class RemoteRequest(BaseModel): | |
| host: str | |
| port: str | |
| user: str | |
| password: str | |
| database: str | |
| db_type: str # "mysql" or "postgresql" | |
| user_query: str | |
| def analyze_remote(body: RemoteRequest): | |
| """ | |
| Connect to a remote MySQL/PostgreSQL and run the natural language query. | |
| """ | |
| db_type = body.db_type.lower() | |
| if db_type not in ("mysql", "postgresql"): | |
| return {"error": "db_type must be 'mysql' or 'postgresql'"} | |
| # Extract schema | |
| if db_type == "mysql": | |
| schema = get_mysql_schema( | |
| body.host, body.port, body.user, body.password, body.database | |
| ) | |
| else: # postgresql | |
| schema = get_postgresql_schema( | |
| body.host, body.port, body.user, body.password, body.database | |
| ) | |
| db_info = { | |
| "type": db_type, | |
| "host": body.host, | |
| "port": body.port, | |
| "user": body.user, | |
| "password": body.password, | |
| "database": body.database, | |
| } | |
| sql, result = generate_and_execute_sql(body.user_query, schema, db_info) | |
| return {"sql": sql, "result": result} | |