2026-06-27 15:42:37 +00:00
"""Live SQL console — PostgreSQL, MySQL, MongoDB, Cassandra, Neo4j + Trino."""
2026-06-25 01:10:57 +00:00
from __future__ import annotations
import os
2026-06-25 01:33:48 +00:00
import re
2026-06-25 01:10:57 +00:00
import time
from typing import Any
import httpx
import psycopg2
2026-06-25 01:33:48 +00:00
import pymysql
2026-06-27 15:42:37 +00:00
from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
from fastapi import APIRouter , Query
2026-06-25 01:10:57 +00:00
from fastapi.responses import JSONResponse
2026-06-27 15:42:37 +00:00
from neo4j import GraphDatabase
2026-06-25 01:10:57 +00:00
from pydantic import BaseModel , Field
2026-06-25 01:33:48 +00:00
from pymongo import MongoClient
2026-06-25 01:10:57 +00:00
2026-06-25 01:33:48 +00:00
DB_HOST = os . getenv ( "DB_VAULT_HOST" , "10.0.21.51" )
2026-06-25 01:10:57 +00:00
PG_PORT = int ( os . getenv ( "PG_PORT" , "5432" ))
PG_USER = os . getenv ( "PG_USER" , "mo" )
PG_PASS = os . getenv ( "PG_PASSWORD" , "Dell2026!" )
PG_DB = os . getenv ( "PG_DATABASE" , "postgres" )
2026-06-25 01:33:48 +00:00
MYSQL_PORT = int ( os . getenv ( "MYSQL_PORT" , "3306" ))
MYSQL_USER = os . getenv ( "MYSQL_USER" , "mo" )
MYSQL_PASS = os . getenv ( "MYSQL_PASSWORD" , "Dell2026!" )
MYSQL_DB = os . getenv ( "MYSQL_DATABASE" , "hr" )
MONGO_HOST = os . getenv ( "MONGO_HOST" , DB_HOST )
MONGO_PORT = int ( os . getenv ( "MONGO_PORT" , "27017" ))
MONGO_DB = os . getenv ( "MONGO_DATABASE" , "supplychain" )
2026-06-27 15:42:37 +00:00
CASS_PORT = int ( os . getenv ( "CASSANDRA_PORT" , "9042" ))
CASS_USER = os . getenv ( "CASSANDRA_USER" , "cassandra" )
CASS_PASS = os . getenv ( "CASSANDRA_PASSWORD" , "cassandra" )
CASS_KS = os . getenv ( "CASSANDRA_KEYSPACE" , "telemetry" )
NEO4J_URI = os . getenv ( "NEO4J_URI" , f "bolt:// { DB_HOST } :7687" )
NEO4J_USER = os . getenv ( "NEO4J_USER" , "neo4j" )
NEO4J_PASS = os . getenv ( "NEO4J_PASSWORD" , "testpwd" )
2026-06-25 01:10:57 +00:00
TRINO_URL = os . getenv ( "TRINO_URL" , "http://10.0.21.50:8089" ) . rstrip ( "/" )
TRINO_USER = os . getenv ( "TRINO_USER" , "atc" )
2026-06-27 15:42:37 +00:00
SOURCE_ENGINES = ( "postgres" , "mysql" , "mongodb" , "cassandra" , "neo4j" )
ENGINES = SOURCE_ENGINES + ( "trino" ,)
2026-06-27 19:37:50 +00:00
from hadoop_sql import (
HADOOP_SAMPLES ,
catalog_hadoop ,
connection_info as hadoop_connection_info ,
health_hadoop ,
sample_hadoop ,
table_row_count_hadoop ,
)
LAKE_ENGINES = ( "hadoop" ,)
2026-06-25 01:10:57 +00:00
router = APIRouter ( prefix = "/api/sql" , tags = [ "sql" ])
SAMPLES : dict [ str , list [ dict [ str , str ]]] = {
"postgres" : [
{ "id" : "pg1" , "label" : "Server version" , "sql" : "SELECT version();" },
{ "id" : "pg2" , "label" : "Current session" , "sql" : "SELECT current_database(), current_user, now();" },
{ "id" : "pg3" , "label" : "Public tables" , "sql" : "SELECT tablename FROM pg_tables WHERE schemaname='public' ORDER BY 1 LIMIT 20;" },
{ "id" : "pg4" , "label" : "Table count" , "sql" : "SELECT count(*) AS public_tables FROM information_schema.tables WHERE table_schema='public';" },
{ "id" : "pg5" , "label" : "Active queries" , "sql" : "SELECT pid, usename, state, left(query, 100) AS query FROM pg_stat_activity WHERE state='active' AND pid <> pg_backend_pid() LIMIT 10;" },
{ "id" : "pg6" , "label" : "Database sizes" , "sql" : "SELECT datname, pg_size_pretty(pg_database_size(datname)) AS size FROM pg_database ORDER BY pg_database_size(datname) DESC LIMIT 10;" },
{ "id" : "pg7" , "label" : "Row counts (stats)" , "sql" : "SELECT relname, n_live_tup FROM pg_stat_user_tables ORDER BY n_live_tup DESC NULLS LAST LIMIT 10;" },
{ "id" : "pg8" , "label" : "Connections" , "sql" : "SELECT count(*) AS connections, state FROM pg_stat_activity GROUP BY state;" },
2026-06-27 15:42:37 +00:00
{ "id" : "pg9" , "label" : "Sales orders sample" , "sql" : "SELECT order_id, customer_id, region, total_amount FROM sales_orders ORDER BY order_id DESC LIMIT 10;" },
2026-06-25 01:10:57 +00:00
{ "id" : "pg10" , "label" : "Memory settings" , "sql" : "SELECT name, setting, unit FROM pg_settings WHERE name IN ('max_connections','shared_buffers','work_mem','effective_cache_size');" },
],
2026-06-25 01:33:48 +00:00
"mysql" : [
{ "id" : "my1" , "label" : "Server version" , "sql" : "SELECT VERSION() AS version;" },
{ "id" : "my2" , "label" : "Current session" , "sql" : "SELECT DATABASE() AS db, USER() AS user, NOW() AS now_ts;" },
{ "id" : "my3" , "label" : "HR tables" , "sql" : "SHOW TABLES;" },
{ "id" : "my4" , "label" : "Table sizes" , "sql" : "SELECT table_name, table_rows, ROUND((data_length+index_length)/1024/1024,1) AS mb FROM information_schema.tables WHERE table_schema='hr' ORDER BY (data_length+index_length) DESC;" },
{ "id" : "my5" , "label" : "Employee events sample" , "sql" : "SELECT event_type, COUNT(*) AS cnt FROM employee_events GROUP BY event_type ORDER BY cnt DESC LIMIT 10;" },
{ "id" : "my6" , "label" : "Recent events" , "sql" : "SELECT employee_id, event_type, event_date FROM employee_events ORDER BY event_date DESC LIMIT 10;" },
{ "id" : "my7" , "label" : "Active connections" , "sql" : "SHOW STATUS LIKE 'Threads_connected';" },
{ "id" : "my8" , "label" : "InnoDB buffer pool" , "sql" : "SHOW STATUS LIKE 'Innodb_buffer_pool%';" },
{ "id" : "my9" , "label" : "Explain count scan" , "sql" : "EXPLAIN SELECT COUNT(*) FROM employee_events;" },
2026-06-25 01:34:07 +00:00
{ "id" : "my10" , "label" : "Events per employee" , "sql" : "SELECT employee_id, COUNT(*) AS events FROM employee_events GROUP BY employee_id ORDER BY events DESC LIMIT 10;" },
2026-06-25 01:33:48 +00:00
],
"mongodb" : [
{ "id" : "mg1" , "label" : "List databases" , "sql" : "SHOW DATABASES" },
{ "id" : "mg2" , "label" : "Supplychain collections" , "sql" : "SHOW COLLECTIONS supplychain" },
{ "id" : "mg3" , "label" : "Count events" , "sql" : "COUNT supplychain.events" },
{ "id" : "mg4" , "label" : "Sample document" , "sql" : "FIND supplychain.events LIMIT 3" },
{ "id" : "mg5" , "label" : "Events by type" , "sql" : "AGGREGATE supplychain.events GROUP type TOP 10" },
{ "id" : "mg6" , "label" : "Events by region" , "sql" : "AGGREGATE supplychain.events GROUP region TOP 10" },
{ "id" : "mg7" , "label" : "Distinct event types" , "sql" : "DISTINCT supplychain.events type" },
{ "id" : "mg8" , "label" : "Server status" , "sql" : "SERVER STATUS" },
{ "id" : "mg9" , "label" : "Indexes on events" , "sql" : "INDEXES supplychain.events" },
{ "id" : "mg10" , "label" : "Recent events" , "sql" : "FIND supplychain.events SORT timestamp DESC LIMIT 5" },
],
2026-06-27 15:42:37 +00:00
"cassandra" : [
{ "id" : "cs1" , "label" : "Cluster release" , "sql" : "SELECT release_version FROM system.local" },
{ "id" : "cs2" , "label" : "Keyspaces" , "sql" : "SELECT keyspace_name FROM system_schema.keyspaces" },
{ "id" : "cs3" , "label" : "Telemetry tables" , "sql" : "SELECT table_name FROM system_schema.tables WHERE keyspace_name='telemetry'" },
{ "id" : "cs4" , "label" : "Device metrics sample" , "sql" : "SELECT device_id, metric_name, metric_value, recorded_at FROM telemetry.device_metrics LIMIT 10" },
{ "id" : "cs5" , "label" : "Metrics by device" , "sql" : "SELECT device_id, count(*) AS cnt FROM telemetry.device_metrics GROUP BY device_id LIMIT 10" },
{ "id" : "cs6" , "label" : "Distinct metric names" , "sql" : "SELECT DISTINCT metric_name FROM telemetry.device_metrics LIMIT 20" },
{ "id" : "cs7" , "label" : "Table columns" , "sql" : "SELECT column_name, type FROM system_schema.columns WHERE keyspace_name='telemetry' AND table_name='device_metrics'" },
{ "id" : "cs8" , "label" : "Recent metrics" , "sql" : "SELECT device_id, metric_name, metric_value FROM telemetry.device_metrics LIMIT 15" },
{ "id" : "cs9" , "label" : "Partition token sample" , "sql" : "SELECT token(device_id) AS tok, device_id FROM telemetry.device_metrics LIMIT 5" },
{ "id" : "cs10" , "label" : "Cluster peers" , "sql" : "SELECT peer, data_center, rack FROM system.peers" },
],
"neo4j" : [
{ "id" : "nj1" , "label" : "Node labels" , "sql" : "CALL db.labels() YIELD label RETURN label" },
{ "id" : "nj2" , "label" : "Relationship types" , "sql" : "CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType" },
{ "id" : "nj3" , "label" : "Total nodes" , "sql" : "MATCH (n) RETURN count(n) AS node_count" },
{ "id" : "nj4" , "label" : "Total relationships" , "sql" : "MATCH ()-[r]->() RETURN count(r) AS rel_count" },
{ "id" : "nj5" , "label" : "Products sample" , "sql" : "MATCH (p:Product) RETURN p.product_id AS id, p.name AS name, p.category AS category LIMIT 10" },
{ "id" : "nj6" , "label" : "Suppliers sample" , "sql" : "MATCH (s:Supplier) RETURN s.supplier_id AS id, s.name AS name, s.country AS country LIMIT 10" },
{ "id" : "nj7" , "label" : "Supply chain links" , "sql" : "MATCH (s:Supplier)-[r:SUPPLIES]->(p:Product) RETURN s.name AS supplier, p.name AS product LIMIT 10" },
{ "id" : "nj8" , "label" : "Products by category" , "sql" : "MATCH (p:Product) RETURN p.category AS category, count(*) AS cnt ORDER BY cnt DESC LIMIT 10" },
{ "id" : "nj9" , "label" : "Schema visualization" , "sql" : "CALL db.schema.visualization()" },
{ "id" : "nj10" , "label" : "Constraint info" , "sql" : "SHOW CONSTRAINTS" },
],
2026-06-27 19:37:50 +00:00
"hadoop" : HADOOP_SAMPLES ,
2026-06-25 01:10:57 +00:00
"trino" : [
{ "id" : "tq1" , "label" : "Trino version" , "sql" : "SELECT version()" },
{ "id" : "tq2" , "label" : "Catalogs" , "sql" : "SHOW CATALOGS" },
{ "id" : "tq3" , "label" : "Iceberg schemas" , "sql" : "SHOW SCHEMAS FROM iceberg" },
2026-06-25 01:24:35 +00:00
{ "id" : "tq4" , "label" : "Sales tables (federated PG)" , "sql" : "SHOW TABLES FROM postgres_sales.public" },
2026-06-25 01:10:57 +00:00
{ "id" : "tq5" , "label" : "Federated PG catalog" , "sql" : "SHOW SCHEMAS FROM postgres_sales" },
{ "id" : "tq6" , "label" : "Cluster nodes" , "sql" : "SELECT node_id, http_uri, state FROM system.runtime.nodes" },
{ "id" : "tq7" , "label" : "Running queries" , "sql" : "SELECT query_id, state, user FROM system.runtime.queries WHERE state != 'FINISHED' LIMIT 10" },
2026-06-25 01:24:35 +00:00
{ "id" : "tq8" , "label" : "Count 1M rows (parallel)" , "sql" : "SELECT count(*) AS rows FROM unnest(sequence(1,1000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)" },
{ "id" : "tq9" , "label" : "Federated sales by region" , "sql" : "SELECT region, count(*) AS orders FROM postgres_sales.public.sales_orders WHERE order_id < 500000 GROUP BY region ORDER BY orders DESC" },
{ "id" : "tq10" , "label" : "Parallel sin() 1M rows" , "sql" : "SELECT sum(sin(a.x)) AS total FROM unnest(sequence(1,1000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)" },
2026-06-25 01:10:57 +00:00
],
}
BENCHMARK_SQL = {
2026-06-25 01:24:35 +00:00
"postgres" : "SELECT sum(sin(i)) AS total FROM generate_series(1, 5000000) i" ,
2026-06-25 01:24:58 +00:00
"trino" : "SELECT sum(sin(x + (y - 1) * 5000)) AS total FROM unnest(sequence(1,5000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)" ,
2026-06-25 01:10:57 +00:00
}
2026-06-25 01:33:48 +00:00
def _tabular ( columns : list [ str ], rows : list [ list [ Any ]], elapsed_ms : int , ** extra : Any ) -> dict [ str , Any ]:
return {
"ok" : True ,
"columns" : columns ,
"rows" : rows ,
"row_count" : len ( rows ),
"elapsed_ms" : elapsed_ms ,
"truncated" : extra . pop ( "truncated" , False ),
** extra ,
}
2026-06-27 15:42:37 +00:00
def _fmt ( val : Any ) -> Any :
if val is None or isinstance ( val , ( str , int , float , bool )):
return val
return str ( val )[: 200 ]
2026-06-25 01:10:57 +00:00
def _run_postgres ( sql : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
conn = psycopg2 . connect (
2026-06-25 01:33:48 +00:00
host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 ,
2026-06-25 01:10:57 +00:00
)
try :
conn . set_session ( readonly = True , autocommit = True )
cur = conn . cursor ()
cur . execute ( sql )
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if cur . description :
columns = [ d [ 0 ] for d in cur . description ]
2026-06-27 15:42:37 +00:00
rows = [[ _fmt ( r ) for r in row ] for row in cur . fetchmany ( limit )]
2026-06-25 01:33:48 +00:00
return _tabular ( columns , rows , elapsed_ms , truncated = len ( rows ) >= limit )
return _tabular ([], [], elapsed_ms , message = "OK" )
finally :
conn . close ()
def _run_mysql ( sql : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
conn = pymysql . connect (
host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS ,
database = MYSQL_DB , connect_timeout = 8 , read_timeout = 30 ,
)
try :
cur = conn . cursor ()
cur . execute ( sql )
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if cur . description :
columns = [ d [ 0 ] for d in cur . description ]
2026-06-27 15:42:37 +00:00
rows = [[ _fmt ( r ) for r in row ] for row in cur . fetchmany ( limit )]
2026-06-25 01:33:48 +00:00
return _tabular ( columns , rows , elapsed_ms , truncated = len ( rows ) >= limit )
return _tabular ([], [], elapsed_ms , message = "OK" )
2026-06-25 01:10:57 +00:00
finally :
conn . close ()
2026-06-25 01:33:48 +00:00
def _mongo_client () -> MongoClient :
return MongoClient ( f "mongodb:// { MONGO_HOST } : { MONGO_PORT } /" , serverSelectionTimeoutMS = 8000 )
def _run_mongo ( query : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
q = query . strip ()
client = _mongo_client ()
try :
columns : list [ str ] = []
rows : list [ list [ Any ]] = []
if q == "SHOW DATABASES" :
columns = [ "database" ]
rows = [[ name ] for name in client . list_database_names ()]
elif m := re . match ( r "^SHOW COLLECTIONS(?:\s+(\w+))?$" , q , re . I ):
db_name = m . group ( 1 ) or MONGO_DB
columns = [ "collection" ]
rows = [[ name ] for name in client [ db_name ] . list_collection_names ()]
elif m := re . match ( r "^COUNT\s+(\w+)\.(\w+)$" , q , re . I ):
db_name , coll = m . group ( 1 ), m . group ( 2 )
columns = [ "count" ]
rows = [[ client [ db_name ][ coll ] . estimated_document_count ()]]
elif m := re . match ( r "^FIND\s+(\w+)\.(\w+)(?:\s+SORT\s+(\w+)\s+DESC)?(?:\s+LIMIT\s+(\d+))?$" , q , re . I ):
db_name , coll , sort_field , lim = m . group ( 1 ), m . group ( 2 ), m . group ( 3 ), m . group ( 4 )
n = min ( int ( lim ) if lim else 10 , limit )
cursor = client [ db_name ][ coll ] . find ({})
if sort_field :
cursor = cursor . sort ( sort_field , - 1 )
docs = list ( cursor . limit ( n ))
if not docs :
columns = [ "result" ]
rows = [[ "(empty)" ]]
else :
columns = sorted ({ k for d in docs for k in d if k != "_id" })
2026-06-27 15:42:37 +00:00
rows = [[ _fmt ( d . get ( c )) for c in columns ] for d in docs ]
2026-06-25 01:33:48 +00:00
elif m := re . match ( r "^DISTINCT\s+(\w+)\.(\w+)\s+(\w+)$" , q , re . I ):
db_name , coll , field = m . group ( 1 ), m . group ( 2 ), m . group ( 3 )
columns = [ field ]
vals = client [ db_name ][ coll ] . distinct ( field )[: limit ]
2026-06-27 15:42:37 +00:00
rows = [[ _fmt ( v )] for v in vals ]
2026-06-25 01:33:48 +00:00
elif m := re . match ( r "^AGGREGATE\s+(\w+)\.(\w+)\s+GROUP\s+(\w+)\s+TOP\s+(\d+)$" , q , re . I ):
db_name , coll , field , top_n = m . group ( 1 ), m . group ( 2 ), m . group ( 3 ), int ( m . group ( 4 ))
pipe = [
{ "$group" : { "_id" : f "$ { field } " , "count" : { "$sum" : 1 }}},
{ "$sort" : { "count" : - 1 }},
{ "$limit" : min ( top_n , limit )},
]
columns = [ field , "count" ]
rows = [[ r . get ( "_id" ), r . get ( "count" )] for r in client [ db_name ][ coll ] . aggregate ( pipe , maxTimeMS = 15000 )]
elif m := re . match ( r "^INDEXES\s+(\w+)\.(\w+)$" , q , re . I ):
db_name , coll = m . group ( 1 ), m . group ( 2 )
columns = [ "name" , "keys" ]
rows = [[ i . get ( "name" ), str ( i . get ( "key" ))] for i in client [ db_name ][ coll ] . list_indexes ()]
elif q . upper () == "SERVER STATUS" :
status = client . admin . command ( "serverStatus" )
columns = [ "metric" , "value" ]
rows = [
[ "version" , status . get ( "version" )],
[ "uptime_secs" , status . get ( "uptime" )],
[ "connections_current" , status . get ( "connections" , {}) . get ( "current" )],
[ "mem_resident_mb" , status . get ( "mem" , {}) . get ( "resident" )],
]
else :
return { "ok" : False , "error" : f "Unknown MongoDB command: { q [: 120 ] } " }
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
return _tabular ( columns , rows [: limit ], elapsed_ms , truncated = len ( rows ) > limit )
finally :
client . close ()
2026-06-27 15:42:37 +00:00
def _cass_cluster () -> Cluster :
auth = PlainTextAuthProvider ( CASS_USER , CASS_PASS ) if CASS_USER else None
return Cluster ([ DB_HOST ], port = CASS_PORT , auth_provider = auth , connect_timeout = 8 )
def _run_cassandra ( cql : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
cluster = _cass_cluster ()
session = cluster . connect ()
try :
rows_raw = list ( session . execute ( cql , timeout = 30 ))
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if not rows_raw :
return _tabular ([], [], elapsed_ms , message = "OK" )
columns = list ( rows_raw [ 0 ] . _fields )
rows = [[ _fmt ( getattr ( r , c )) for c in columns ] for r in rows_raw [: limit ]]
return _tabular ( columns , rows , elapsed_ms , truncated = len ( rows_raw ) > limit )
finally :
cluster . shutdown ()
def _neo4j_driver ():
return GraphDatabase . driver ( NEO4J_URI , auth = ( NEO4J_USER , NEO4J_PASS ))
def _run_neo4j ( cypher : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
driver = _neo4j_driver ()
try :
with driver . session () as session :
result = session . run ( cypher )
keys = result . keys ()
columns = list ( keys )
rows_raw = [ record . values () for record in result ]
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
rows = [[ _fmt ( v ) for v in row ] for row in rows_raw [: limit ]]
return _tabular ( columns , rows , elapsed_ms , truncated = len ( rows_raw ) > limit )
finally :
driver . close ()
2026-06-25 01:10:57 +00:00
def _run_trino ( sql : str , limit : int = 200 ) -> dict [ str , Any ]:
t0 = time . perf_counter ()
headers = { "X-Trino-User" : TRINO_USER , "Content-Type" : "text/plain" }
with httpx . Client ( timeout = 120.0 ) as client :
resp = client . post ( f " { TRINO_URL } /v1/statement" , content = sql , headers = headers )
if resp . status_code >= 400 :
return { "ok" : False , "error" : resp . text [: 500 ]}
data = resp . json ()
if data . get ( "error" ):
return { "ok" : False , "error" : str ( data [ "error" ])[: 500 ]}
columns : list [ str ] = []
rows : list [ list [ Any ]] = []
while True :
if data . get ( "columns" ) and not columns :
columns = [ c [ "name" ] for c in data [ "columns" ]]
if data . get ( "data" ):
rows . extend ( data [ "data" ])
if len ( rows ) >= limit :
rows = rows [: limit ]
break
if data . get ( "error" ):
return { "ok" : False , "error" : str ( data [ "error" ])[: 500 ]}
nxt = data . get ( "nextUri" )
if not nxt :
break
data = client . get ( nxt , headers = headers ) . json ()
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
stats = data . get ( "stats" ) or {}
2026-06-25 01:33:48 +00:00
return _tabular ( columns , rows , elapsed_ms , engine_stats_ms = stats . get ( "elapsedTimeMillis" ), truncated = len ( rows ) >= limit )
2026-06-25 01:10:57 +00:00
class SqlRequest ( BaseModel ):
2026-06-27 19:37:50 +00:00
engine : str = Field ( ... , pattern = "^(postgres|mysql|mongodb|trino|cassandra|neo4j|hadoop)$" )
2026-06-25 01:10:57 +00:00
sql : str = Field ( ... , min_length = 1 , max_length = 8000 )
2026-06-27 19:37:50 +00:00
class RowUpdateRequest ( BaseModel ):
engine : str = Field ( ... , pattern = "^(postgres|mysql|mongodb|cassandra|neo4j)$" )
object : str = Field ( ... , min_length = 1 , max_length = 256 )
pk : dict [ str , Any ] = Field ( ... , min_length = 1 )
changes : dict [ str , Any ] = Field ( ... , min_length = 1 )
2026-06-28 17:12:29 +00:00
class RowInsertRequest ( BaseModel ):
engine : str = Field ( ... , pattern = "^(postgres|mysql|mongodb|cassandra|neo4j)$" )
object : str = Field ( ... , min_length = 1 , max_length = 256 )
values : dict [ str , Any ] = Field ( ... , min_length = 1 )
2026-06-27 15:42:37 +00:00
def _connection_info ( engine : str ) -> dict [ str , str ]:
if engine == "postgres" :
return { "host" : DB_HOST , "port" : str ( PG_PORT ), "database" : PG_DB , "user" : PG_USER }
if engine == "mysql" :
return { "host" : DB_HOST , "port" : str ( MYSQL_PORT ), "database" : MYSQL_DB , "user" : MYSQL_USER }
if engine == "mongodb" :
return { "host" : MONGO_HOST , "port" : str ( MONGO_PORT ), "database" : MONGO_DB }
if engine == "cassandra" :
return { "host" : DB_HOST , "port" : str ( CASS_PORT ), "keyspace" : CASS_KS , "user" : CASS_USER }
if engine == "neo4j" :
return { "uri" : NEO4J_URI , "user" : NEO4J_USER }
return { "url" : TRINO_URL , "user" : TRINO_USER }
def _dispatch ( engine : str , sql : str , limit : int = 200 ) -> dict [ str , Any ]:
if engine == "postgres" :
return _run_postgres ( sql , limit )
if engine == "mysql" :
return _run_mysql ( sql , limit )
if engine == "mongodb" :
return _run_mongo ( sql , limit )
if engine == "cassandra" :
return _run_cassandra ( sql , limit )
if engine == "neo4j" :
return _run_neo4j ( sql , limit )
2026-06-27 19:37:50 +00:00
if engine == "hadoop" :
return _run_trino ( sql , limit )
2026-06-27 15:42:37 +00:00
return _run_trino ( sql , limit )
def _catalog_postgres () -> dict [ str , Any ]:
conn = psycopg2 . connect ( host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 )
try :
conn . set_session ( readonly = True , autocommit = True )
cur = conn . cursor ()
cur . execute (
"SELECT schemaname, relname, n_live_tup FROM pg_stat_user_tables ORDER BY n_live_tup DESC NULLS LAST LIMIT 50"
)
objects = []
for schema , table , rows in cur . fetchall ():
objects . append ({ "type" : "table" , "schema" : schema , "name" : table , "fqn" : f " { schema } . { table } " , "row_count" : int ( rows or 0 )})
cur . execute ( "SELECT version()" )
version = cur . fetchone ()[ 0 ]
return { "engine" : "postgres" , "version" : version , "objects" : objects }
finally :
conn . close ()
def _catalog_mysql () -> dict [ str , Any ]:
conn = pymysql . connect ( host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS , database = MYSQL_DB , connect_timeout = 8 )
try :
cur = conn . cursor ()
cur . execute (
"SELECT table_name, table_rows, ROUND((data_length+index_length)/1024/1024,1) AS mb "
"FROM information_schema.tables WHERE table_schema= %s ORDER BY table_rows DESC" ,
( MYSQL_DB ,),
)
objects = [{ "type" : "table" , "schema" : MYSQL_DB , "name" : r [ 0 ], "fqn" : f " { MYSQL_DB } . { r [ 0 ] } " , "row_count" : int ( r [ 1 ] or 0 ), "size_mb" : float ( r [ 2 ] or 0 )} for r in cur . fetchall ()]
cur . execute ( "SELECT VERSION()" )
version = cur . fetchone ()[ 0 ]
return { "engine" : "mysql" , "version" : version , "objects" : objects }
finally :
conn . close ()
def _catalog_mongodb () -> dict [ str , Any ]:
client = _mongo_client ()
try :
objects = []
for db_name in client . list_database_names ():
if db_name in ( "admin" , "config" , "local" ):
continue
for coll in client [ db_name ] . list_collection_names ():
try :
cnt = client [ db_name ][ coll ] . estimated_document_count ()
except Exception :
cnt = None
objects . append ({ "type" : "collection" , "schema" : db_name , "name" : coll , "fqn" : f " { db_name } . { coll } " , "row_count" : cnt })
build = client . admin . command ( "buildInfo" )
return { "engine" : "mongodb" , "version" : build . get ( "version" ), "objects" : objects }
finally :
client . close ()
def _catalog_cassandra () -> dict [ str , Any ]:
cluster = _cass_cluster ()
session = cluster . connect ()
try :
objects = []
kss = [
r . keyspace_name for r in session . execute ( "SELECT keyspace_name FROM system_schema.keyspaces" )
if not r . keyspace_name . startswith ( "system" )
]
for ks in kss :
tables = session . execute (
"SELECT table_name FROM system_schema.tables WHERE keyspace_name= %s " , ( ks ,)
)
for t in tables :
objects . append ({ "type" : "table" , "schema" : ks , "name" : t . table_name , "fqn" : f " { ks } . { t . table_name } " , "row_count" : None })
ver = list ( session . execute ( "SELECT release_version FROM system.local" ))[ 0 ] . release_version
return { "engine" : "cassandra" , "version" : ver , "objects" : objects }
finally :
cluster . shutdown ()
def _catalog_neo4j () -> dict [ str , Any ]:
driver = _neo4j_driver ()
try :
with driver . session () as session :
labels = [ r [ "label" ] for r in session . run ( "CALL db.labels() YIELD label RETURN label" )]
rels = [ r [ "relationshipType" ] for r in session . run ( "CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType" )]
objects = []
for label in labels :
cnt = session . run ( f "MATCH (n:` { label } `) RETURN count(n) AS c" ) . single ()[ "c" ]
objects . append ({ "type" : "node_label" , "schema" : "graph" , "name" : label , "fqn" : label , "row_count" : cnt })
for rel in rels :
cnt = session . run ( f "MATCH ()-[r:` { rel } `]->() RETURN count(r) AS c" ) . single ()[ "c" ]
objects . append ({ "type" : "relationship" , "schema" : "graph" , "name" : rel , "fqn" : rel , "row_count" : cnt })
total = session . run ( "MATCH (n) RETURN count(n) AS c" ) . single ()[ "c" ]
return { "engine" : "neo4j" , "version" : "4.4" , "total_nodes" : total , "objects" : objects }
finally :
driver . close ()
2026-06-27 19:37:50 +00:00
def _catalog_hadoop () -> dict [ str , Any ]:
return catalog_hadoop ( _run_trino )
def _sample_hadoop ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
return sample_hadoop ( object_name , limit , offset , _run_trino , _tabular )
2026-06-27 15:42:37 +00:00
def _catalog ( engine : str ) -> dict [ str , Any ]:
if engine == "postgres" :
return _catalog_postgres ()
if engine == "mysql" :
return _catalog_mysql ()
if engine == "mongodb" :
return _catalog_mongodb ()
if engine == "cassandra" :
return _catalog_cassandra ()
if engine == "neo4j" :
return _catalog_neo4j ()
2026-06-27 19:37:50 +00:00
if engine == "hadoop" :
return _catalog_hadoop ()
2026-06-27 15:42:37 +00:00
raise ValueError ( f "Catalog not supported for { engine } " )
2026-06-27 19:37:50 +00:00
def _sample_postgres ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
if "." in object_name :
schema , table = object_name . split ( "." , 1 )
2026-06-27 19:37:50 +00:00
sql = f 'SELECT * FROM " { schema } "." { table } " OFFSET { int ( offset ) } LIMIT { int ( limit ) } '
2026-06-27 15:42:37 +00:00
else :
2026-06-27 19:37:50 +00:00
sql = f 'SELECT * FROM public." { object_name } " OFFSET { int ( offset ) } LIMIT { int ( limit ) } '
2026-06-27 15:42:37 +00:00
return _run_postgres ( sql , limit )
2026-06-27 19:37:50 +00:00
def _sample_mysql ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
table = object_name . split ( "." )[ - 1 ]
2026-06-27 19:37:50 +00:00
return _run_mysql ( f "SELECT * FROM ` { table } ` LIMIT { int ( limit ) } OFFSET { int ( offset ) } " , limit )
2026-06-27 15:42:37 +00:00
2026-06-27 19:37:50 +00:00
def _sample_mongodb ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
if "." in object_name :
db_name , coll = object_name . split ( "." , 1 )
else :
db_name , coll = MONGO_DB , object_name
2026-06-27 19:37:50 +00:00
t0 = time . perf_counter ()
client = _mongo_client ()
try :
docs = list ( client [ db_name ][ coll ] . find ({}) . skip ( int ( offset )) . limit ( int ( limit )))
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if not docs :
return _tabular ([ "result" ], [[ "(empty)" ]], elapsed_ms )
columns = sorted ({ k for d in docs for k in d })
rows = [[ _fmt ( d . get ( c )) for c in columns ] for d in docs ]
return _tabular ( columns , rows , elapsed_ms , object = object_name )
finally :
client . close ()
2026-06-27 15:42:37 +00:00
2026-06-27 19:37:50 +00:00
def _sample_cassandra ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
if "." in object_name :
ks , table = object_name . split ( "." , 1 )
else :
ks , table = CASS_KS , object_name
2026-06-27 19:37:50 +00:00
t0 = time . perf_counter ()
cluster = _cass_cluster ()
session = cluster . connect ()
try :
stmt = f "SELECT * FROM { ks } . { table } "
result = session . execute ( stmt , timeout = 60 )
skipped = 0
picked : list [ Any ] = []
columns : list [ str ] = []
for row in result :
if skipped < offset :
skipped += 1
continue
if not columns :
columns = list ( row . _fields )
picked . append ( row )
if len ( picked ) >= limit :
break
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if not picked :
return _tabular ([ "result" ], [[ "(empty)" ]], elapsed_ms )
rows = [[ _fmt ( getattr ( r , c )) for c in columns ] for r in picked ]
return _tabular ( columns , rows , elapsed_ms , truncated = len ( picked ) >= limit )
finally :
cluster . shutdown ()
2026-06-27 15:42:37 +00:00
2026-06-27 15:48:56 +00:00
def _graph_neo4j ( edge_limit : int = 60 , rel_type : str | None = None ) -> dict [ str , Any ]:
"""Return a sampled subgraph (nodes + edges) for interactive graph visualization."""
t0 = time . perf_counter ()
driver = _neo4j_driver ()
try :
with driver . session () as session :
total = session . run ( "MATCH (n) RETURN count(n) AS c" ) . single ()[ "c" ]
rel_types = [
r [ "relationshipType" ]
for r in session . run (
"CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType"
)
]
if rel_type :
cypher = """
MATCH (a)-[r]->(b)
WHERE type(r) = $rel_type
WITH a, r, b LIMIT $limit
RETURN a, r, b, labels(a)[0] AS la, type(r) AS rt, labels(b)[0] AS lb
"""
params : dict [ str , Any ] = { "rel_type" : rel_type , "limit" : edge_limit }
else :
cypher = """
MATCH (a)-[r]->(b)
WITH a, r, b LIMIT $limit
RETURN a, r, b, labels(a)[0] AS la, type(r) AS rt, labels(b)[0] AS lb
"""
params = { "limit" : edge_limit }
records = session . run ( cypher , ** params ) . data ()
def _node_key ( node : Any , label : str ) -> str :
props = dict ( node )
if label == "Product" and props . get ( "product_id" ) is not None :
return f "Product: { props [ 'product_id' ] } "
if label == "Supplier" and props . get ( "supplier_id" ) is not None :
return f "Supplier: { props [ 'supplier_id' ] } "
nid = getattr ( node , "element_id" , None ) or getattr ( node , "id" , None )
return f " { label } : { nid } "
nodes_map : dict [ str , dict [ str , Any ]] = {}
edges : list [ dict [ str , Any ]] = []
for rec in records :
a , b , rt = rec [ "a" ], rec [ "b" ], rec [ "rt" ]
la , lb = rec [ "la" ], rec [ "lb" ]
aid , bid = _node_key ( a , la ), _node_key ( b , lb )
if aid not in nodes_map :
ap = dict ( a )
nodes_map [ aid ] = {
"id" : aid ,
"label" : la ,
"name" : ap . get ( "name" ) or str ( ap . get ( "product_id" ) or ap . get ( "supplier_id" ) or aid ),
"properties" : { k : _fmt ( v ) for k , v in ap . items ()},
}
if bid not in nodes_map :
bp = dict ( b )
nodes_map [ bid ] = {
"id" : bid ,
"label" : lb ,
"name" : bp . get ( "name" ) or str ( bp . get ( "product_id" ) or bp . get ( "supplier_id" ) or bid ),
"properties" : { k : _fmt ( v ) for k , v in bp . items ()},
}
edges . append ({ "id" : f " { aid } | { rt } | { bid } " , "source" : aid , "target" : bid , "type" : rt })
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
return {
"ok" : True ,
"nodes" : list ( nodes_map . values ()),
"edges" : edges ,
"rel_types" : rel_types ,
"stats" : {
"total_nodes_db" : total ,
"returned_nodes" : len ( nodes_map ),
"returned_edges" : len ( edges ),
},
"elapsed_ms" : elapsed_ms ,
}
finally :
driver . close ()
2026-06-27 19:37:50 +00:00
def _sample_neo4j ( object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
label = object_name . split ( "." )[ - 1 ]
2026-06-27 19:37:50 +00:00
cypher = f "MATCH (n:` { label } `) RETURN n SKIP { int ( offset ) } LIMIT { int ( limit ) } "
2026-06-27 15:42:37 +00:00
t0 = time . perf_counter ()
driver = _neo4j_driver ()
try :
with driver . session () as session :
result = session . run ( cypher )
docs = [ dict ( record [ "n" ]) for record in result ]
elapsed_ms = int (( time . perf_counter () - t0 ) * 1000 )
if not docs :
return _tabular ([ "result" ], [[ "(empty)" ]], elapsed_ms )
columns = sorted ({ k for d in docs for k in d })
rows = [[ _fmt ( d . get ( c )) for c in columns ] for d in docs ]
return _tabular ( columns , rows , elapsed_ms , object = object_name )
finally :
driver . close ()
2026-06-27 19:37:50 +00:00
def _sample ( engine : str , object_name : str , limit : int , offset : int = 0 ) -> dict [ str , Any ]:
2026-06-27 15:42:37 +00:00
if engine == "postgres" :
2026-06-27 19:37:50 +00:00
return _sample_postgres ( object_name , limit , offset )
2026-06-27 15:42:37 +00:00
if engine == "mysql" :
2026-06-27 19:37:50 +00:00
return _sample_mysql ( object_name , limit , offset )
2026-06-27 15:42:37 +00:00
if engine == "mongodb" :
2026-06-27 19:37:50 +00:00
return _sample_mongodb ( object_name , limit , offset )
2026-06-27 15:42:37 +00:00
if engine == "cassandra" :
2026-06-27 19:37:50 +00:00
return _sample_cassandra ( object_name , limit , offset )
2026-06-27 15:42:37 +00:00
if engine == "neo4j" :
2026-06-27 19:37:50 +00:00
return _sample_neo4j ( object_name , limit , offset )
if engine == "hadoop" :
return _sample_hadoop ( object_name , limit , offset )
2026-06-27 15:42:37 +00:00
raise ValueError ( f "Sample not supported for { engine } " )
2026-06-27 19:37:50 +00:00
def _table_row_count ( engine : str , object_name : str ) -> int | None :
try :
if engine == "postgres" :
schema , table = _parse_fqn ( engine , object_name )
conn = psycopg2 . connect (
host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 ,
)
try :
cur = conn . cursor ()
2026-06-28 10:10:17 +00:00
# Planner estimate first; only do a (bounded) exact count for small
# tables so huge tables don't trigger a slow full scan.
cur . execute ( "SELECT reltuples::bigint FROM pg_class WHERE oid = %s ::regclass" , ( f '" { schema } "." { table } "' ,))
row = cur . fetchone ()
est = int ( row [ 0 ]) if row and row [ 0 ] is not None and row [ 0 ] >= 0 else 0
if est <= 50000 :
try :
cur . execute ( "SET LOCAL statement_timeout = 4000" )
cur . execute ( f 'SELECT count(*) FROM " { schema } "." { table } "' )
return int ( cur . fetchone ()[ 0 ])
except Exception :
conn . rollback ()
return est
return est
2026-06-27 19:37:50 +00:00
finally :
conn . close ()
if engine == "mysql" :
schema , table = _parse_fqn ( engine , object_name )
conn = pymysql . connect (
host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS ,
database = MYSQL_DB , connect_timeout = 8 ,
)
try :
cur = conn . cursor ()
2026-06-28 10:10:17 +00:00
# Fast estimate from statistics; only do a (bounded) exact count for
# small tables to avoid full scans on huge ones (e.g. 24M rows).
cur . execute (
"SELECT table_rows FROM information_schema.tables WHERE table_schema= %s AND table_name= %s " ,
( schema , table ),
)
row = cur . fetchone ()
est = int ( row [ 0 ]) if row and row [ 0 ] is not None else 0
if est <= 50000 :
try :
cur . execute ( "SET SESSION MAX_EXECUTION_TIME=4000" )
cur . execute ( f "SELECT count(*) FROM ` { table } `" )
return int ( cur . fetchone ()[ 0 ])
except Exception :
return est
return est
2026-06-27 19:37:50 +00:00
finally :
conn . close ()
if engine == "mongodb" :
db_name , coll = _parse_fqn ( engine , object_name )
client = _mongo_client ()
try :
return int ( client [ db_name ][ coll ] . estimated_document_count ())
finally :
client . close ()
if engine == "neo4j" :
label = object_name . split ( "." )[ - 1 ]
driver = _neo4j_driver ()
try :
with driver . session () as session :
return int ( session . run ( f "MATCH (n:` { label } `) RETURN count(n) AS c" ) . single ()[ "c" ])
finally :
driver . close ()
except Exception :
pass
return None
def _parse_fqn ( engine : str , object_name : str ) -> tuple [ str , str ]:
if engine == "postgres" :
if "." in object_name :
return object_name . split ( "." , 1 )
return "public" , object_name
if engine == "mysql" :
if "." in object_name :
return object_name . split ( "." , 1 )
return MYSQL_DB , object_name
if engine == "mongodb" :
if "." in object_name :
return object_name . split ( "." , 1 )
return MONGO_DB , object_name
if engine == "cassandra" :
if "." in object_name :
return object_name . split ( "." , 1 )
return CASS_KS , object_name
return "graph" , object_name . split ( "." )[ - 1 ]
def _object_primary_keys ( engine : str , object_name : str ) -> list [ str ]:
try :
if engine == "postgres" :
schema , table = _parse_fqn ( engine , object_name )
conn = psycopg2 . connect (
host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 ,
)
try :
cur = conn . cursor ()
cur . execute (
"""
SELECT a.attname
FROM pg_constraint c
JOIN pg_class t ON t.oid = c.conrelid
JOIN pg_namespace n ON n.oid = t.relnamespace
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(c.conkey)
WHERE c.contype = 'p' AND n.nspname = %s AND t.relname = %s
ORDER BY array_position(c.conkey, a.attnum)
""" ,
( schema , table ),
)
return [ r [ 0 ] for r in cur . fetchall ()]
finally :
conn . close ()
if engine == "mysql" :
schema , table = _parse_fqn ( engine , object_name )
conn = pymysql . connect (
host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS ,
database = MYSQL_DB , connect_timeout = 8 ,
)
try :
cur = conn . cursor ()
cur . execute (
"""
SELECT COLUMN_NAME FROM information_schema.KEY_COLUMN_USAGE
WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND CONSTRAINT_NAME = 'PRIMARY'
ORDER BY ORDINAL_POSITION
""" ,
( schema , table ),
)
return [ r [ 0 ] for r in cur . fetchall ()]
finally :
conn . close ()
if engine == "mongodb" :
return [ "_id" ]
if engine == "cassandra" :
ks , table = _parse_fqn ( engine , object_name )
cluster = _cass_cluster ()
session = cluster . connect ()
try :
rows = session . execute (
"SELECT column_name FROM system_schema.columns "
"WHERE keyspace_name= %s AND table_name= %s AND kind IN ('partition_key','clustering') "
"ORDER BY position" ,
( ks , table ),
)
return [ r . column_name for r in rows ]
finally :
cluster . shutdown ()
if engine == "neo4j" :
label = object_name . split ( "." )[ - 1 ]
if label == "Product" :
return [ "product_id" ]
if label == "Supplier" :
return [ "supplier_id" ]
except Exception :
pass
return []
def _coerce_value ( val : Any ) -> Any :
if val is None or val == "" :
return None
if isinstance ( val , ( int , float , bool )):
return val
s = str ( val )
if s . lower () == "null" :
return None
if re . fullmatch ( r "-?\d+" , s ):
try :
return int ( s )
except ValueError :
pass
if re . fullmatch ( r "-?\d+\.\d+" , s ):
try :
return float ( s )
except ValueError :
pass
return s
def _safe_ident ( name : str ) -> bool :
return bool ( re . fullmatch ( r "[a-zA-Z_][a-zA-Z0-9_]*" , name ))
def _update_postgres ( object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> str :
schema , table = _parse_fqn ( "postgres" , object_name )
if not changes or not pk :
raise ValueError ( "Primary key and changes required" )
for c in list ( changes ) + list ( pk ):
if not _safe_ident ( c ):
raise ValueError ( f "Invalid column name: { c } " )
set_sql = ", " . join ( f '" { c } "=%s' for c in changes )
where_sql = " AND " . join ( f '" { c } "=%s' for c in pk )
sql = f 'UPDATE " { schema } "." { table } " SET { set_sql } WHERE { where_sql } '
params = list ( changes . values ()) + list ( pk . values ())
conn = psycopg2 . connect (
host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 ,
)
try :
conn . autocommit = True
cur = conn . cursor ()
cur . execute ( sql , params )
if cur . rowcount == 0 :
raise ValueError ( "No row matched primary key" )
return sql
finally :
conn . close ()
def _update_mysql ( object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> str :
_schema , table = _parse_fqn ( "mysql" , object_name )
if not changes or not pk :
raise ValueError ( "Primary key and changes required" )
for c in list ( changes ) + list ( pk ):
if not _safe_ident ( c ):
raise ValueError ( f "Invalid column name: { c } " )
set_sql = ", " . join ( f "` { c } `=%s" for c in changes )
where_sql = " AND " . join ( f "` { c } `=%s" for c in pk )
sql = f "UPDATE ` { table } ` SET { set_sql } WHERE { where_sql } "
params = list ( changes . values ()) + list ( pk . values ())
conn = pymysql . connect (
host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS ,
database = MYSQL_DB , connect_timeout = 8 ,
)
try :
conn . autocommit = True
cur = conn . cursor ()
cur . execute ( sql , params )
if cur . rowcount == 0 :
raise ValueError ( "No row matched primary key" )
return sql
finally :
conn . close ()
def _update_mongodb ( object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> str :
db_name , coll = _parse_fqn ( "mongodb" , object_name )
if not pk or not changes :
raise ValueError ( "Primary key and changes required" )
client = _mongo_client ()
try :
from bson import ObjectId
filt = dict ( pk )
if "_id" in filt and isinstance ( filt [ "_id" ], str ) and len ( filt [ "_id" ]) == 24 :
try :
filt [ "_id" ] = ObjectId ( filt [ "_id" ])
except Exception :
pass
res = client [ db_name ][ coll ] . update_one ( filt , { "$set" : changes })
if res . matched_count == 0 :
raise ValueError ( "No document matched primary key" )
return f "UPDATE { db_name } . { coll } { filt } "
finally :
client . close ()
def _update_cassandra ( object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> str :
ks , table = _parse_fqn ( "cassandra" , object_name )
if not pk or not changes :
raise ValueError ( "Primary key and changes required" )
set_cql = ", " . join ( f " { c } =%s" for c in changes )
where_cql = " AND " . join ( f " { c } =%s" for c in pk )
cql = f "UPDATE { ks } . { table } SET { set_cql } WHERE { where_cql } "
cluster = _cass_cluster ()
session = cluster . connect ()
try :
session . execute ( cql , list ( changes . values ()) + list ( pk . values ()))
return cql
finally :
cluster . shutdown ()
def _update_neo4j ( object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> str :
label = object_name . split ( "." )[ - 1 ]
if not pk or not changes :
raise ValueError ( "Primary key and changes required" )
pk_col , pk_val = next ( iter ( pk . items ()))
if not _safe_ident ( pk_col ):
raise ValueError ( f "Invalid property name: { pk_col } " )
for c in changes :
if not _safe_ident ( c ):
raise ValueError ( f "Invalid property name: { c } " )
set_frag = ", " . join ( f "n. { c } = $ { c } " for c in changes )
cypher = f "MATCH (n:` { label } ` {{ { pk_col } : $pk }} ) SET { set_frag } RETURN n"
params : dict [ str , Any ] = { "pk" : _coerce_value ( pk_val )}
params . update ({ c : _coerce_value ( v ) for c , v in changes . items ()})
driver = _neo4j_driver ()
try :
with driver . session () as session :
rec = session . run ( cypher , ** params ) . single ()
if not rec :
raise ValueError ( "No node matched primary key" )
return cypher
finally :
driver . close ()
def _update_row ( engine : str , object_name : str , pk : dict [ str , Any ], changes : dict [ str , Any ]) -> dict [ str , Any ]:
pk = { k : _coerce_value ( v ) for k , v in pk . items ()}
changes = { k : _coerce_value ( v ) for k , v in changes . items ()}
if engine == "postgres" :
sql = _update_postgres ( object_name , pk , changes )
elif engine == "mysql" :
sql = _update_mysql ( object_name , pk , changes )
elif engine == "mongodb" :
sql = _update_mongodb ( object_name , pk , changes )
elif engine == "cassandra" :
sql = _update_cassandra ( object_name , pk , changes )
elif engine == "neo4j" :
sql = _update_neo4j ( object_name , pk , changes )
else :
raise ValueError ( f "Update not supported for { engine } " )
return { "ok" : True , "engine" : engine , "object" : object_name , "statement" : sql , "cdc" : engine in ( "postgres" , "mysql" , "mongodb" )}
2026-06-28 17:12:29 +00:00
def _insert_postgres ( object_name : str , values : dict [ str , Any ]) -> str :
schema , table = _parse_fqn ( "postgres" , object_name )
cols = [ c for c in values if values [ c ] is not None ]
if not cols :
raise ValueError ( "No values provided" )
for c in cols :
if not _safe_ident ( c ):
raise ValueError ( f "Invalid column name: { c } " )
col_sql = ", " . join ( f '" { c } "' for c in cols )
ph = ", " . join ([ " %s " ] * len ( cols ))
sql = f 'INSERT INTO " { schema } "." { table } " ( { col_sql } ) VALUES ( { ph } )'
conn = psycopg2 . connect ( host = DB_HOST , port = PG_PORT , user = PG_USER , password = PG_PASS , dbname = PG_DB , connect_timeout = 8 )
try :
conn . autocommit = True
conn . cursor () . execute ( sql , [ values [ c ] for c in cols ])
return sql
finally :
conn . close ()
def _insert_mysql ( object_name : str , values : dict [ str , Any ]) -> str :
_schema , table = _parse_fqn ( "mysql" , object_name )
cols = [ c for c in values if values [ c ] is not None ]
if not cols :
raise ValueError ( "No values provided" )
for c in cols :
if not _safe_ident ( c ):
raise ValueError ( f "Invalid column name: { c } " )
col_sql = ", " . join ( f "` { c } `" for c in cols )
ph = ", " . join ([ " %s " ] * len ( cols ))
sql = f "INSERT INTO ` { table } ` ( { col_sql } ) VALUES ( { ph } )"
conn = pymysql . connect ( host = DB_HOST , port = MYSQL_PORT , user = MYSQL_USER , password = MYSQL_PASS , database = MYSQL_DB , connect_timeout = 8 )
try :
conn . autocommit = True
conn . cursor () . execute ( sql , [ values [ c ] for c in cols ])
return sql
finally :
conn . close ()
def _insert_mongodb ( object_name : str , values : dict [ str , Any ]) -> str :
db_name , coll = _parse_fqn ( "mongodb" , object_name )
doc = { k : v for k , v in values . items () if v is not None }
if not doc :
raise ValueError ( "No values provided" )
client = _mongo_client ()
try :
res = client [ db_name ][ coll ] . insert_one ( doc )
return f "INSERT { db_name } . { coll } _id= { res . inserted_id } "
finally :
client . close ()
def _insert_cassandra ( object_name : str , values : dict [ str , Any ]) -> str :
ks , table = _parse_fqn ( "cassandra" , object_name )
cols = [ c for c in values if values [ c ] is not None ]
if not cols :
raise ValueError ( "No values provided" )
for c in cols :
if not _safe_ident ( c ):
raise ValueError ( f "Invalid column name: { c } " )
col_sql = ", " . join ( cols )
ph = ", " . join ([ " %s " ] * len ( cols ))
cql = f "INSERT INTO { ks } . { table } ( { col_sql } ) VALUES ( { ph } )"
cluster = _cass_cluster ()
session = cluster . connect ()
try :
session . execute ( cql , [ values [ c ] for c in cols ])
return cql
finally :
cluster . shutdown ()
def _insert_neo4j ( object_name : str , values : dict [ str , Any ]) -> str :
label = object_name . split ( "." )[ - 1 ]
props = { k : _coerce_value ( v ) for k , v in values . items () if v is not None }
if not props :
raise ValueError ( "No values provided" )
for c in props :
if not _safe_ident ( c ):
raise ValueError ( f "Invalid property name: { c } " )
set_frag = ", " . join ( f "n. { c } = $ { c } " for c in props )
cypher = f "CREATE (n:` { label } `) SET { set_frag } RETURN n"
driver = _neo4j_driver ()
try :
with driver . session () as session :
session . run ( cypher , ** props ) . consume ()
return cypher
finally :
driver . close ()
def _insert_row ( engine : str , object_name : str , values : dict [ str , Any ]) -> dict [ str , Any ]:
values = { k : _coerce_value ( v ) for k , v in values . items ()}
if engine == "postgres" :
sql = _insert_postgres ( object_name , values )
elif engine == "mysql" :
sql = _insert_mysql ( object_name , values )
elif engine == "mongodb" :
sql = _insert_mongodb ( object_name , values )
elif engine == "cassandra" :
sql = _insert_cassandra ( object_name , values )
elif engine == "neo4j" :
sql = _insert_neo4j ( object_name , values )
else :
raise ValueError ( f "Insert not supported for { engine } " )
return { "ok" : True , "engine" : engine , "object" : object_name , "statement" : sql , "cdc" : engine in ( "postgres" , "mysql" , "mongodb" )}
2026-06-25 01:10:57 +00:00
@router.get ( "/samples/ {engine} " )
async def get_samples ( engine : str ):
2026-06-27 19:37:50 +00:00
if engine not in SAMPLES and engine not in LAKE_ENGINES :
2026-06-25 01:10:57 +00:00
return JSONResponse ({ "error" : "unknown engine" }, status_code = 404 )
2026-06-27 19:37:50 +00:00
if engine == "hadoop" :
return { "engine" : engine , "samples" : SAMPLES [ engine ], "connection" : hadoop_connection_info ()}
2026-06-25 01:33:48 +00:00
return { "engine" : engine , "samples" : SAMPLES [ engine ], "connection" : _connection_info ( engine )}
2026-06-25 01:10:57 +00:00
@router.get ( "/health" )
async def sql_health ():
2026-06-25 01:33:48 +00:00
out : dict [ str , Any ] = {}
2026-06-27 15:42:37 +00:00
for eng in SOURCE_ENGINES :
2026-06-25 01:33:48 +00:00
try :
if eng == "postgres" :
_run_postgres ( "SELECT 1" )
out [ eng ] = { "ok" : True , "host" : DB_HOST , "user" : PG_USER , "error" : None }
elif eng == "mysql" :
_run_mysql ( "SELECT 1" )
out [ eng ] = { "ok" : True , "host" : DB_HOST , "database" : MYSQL_DB , "user" : MYSQL_USER , "error" : None }
elif eng == "mongodb" :
_run_mongo ( "SHOW DATABASES" )
out [ eng ] = { "ok" : True , "host" : MONGO_HOST , "database" : MONGO_DB , "error" : None }
2026-06-27 15:42:37 +00:00
elif eng == "cassandra" :
_run_cassandra ( "SELECT release_version FROM system.local LIMIT 1" )
out [ eng ] = { "ok" : True , "host" : DB_HOST , "keyspace" : CASS_KS , "error" : None }
2026-06-25 01:33:48 +00:00
else :
2026-06-27 15:42:37 +00:00
_run_neo4j ( "RETURN 1 AS n" )
out [ eng ] = { "ok" : True , "uri" : NEO4J_URI , "user" : NEO4J_USER , "error" : None }
2026-06-25 01:33:48 +00:00
except Exception as exc :
out [ eng ] = { "ok" : False , "error" : str ( exc )[: 200 ]}
return out
2026-06-25 01:10:57 +00:00
2026-06-27 19:37:50 +00:00
@router.get ( "/health/ {engine} " )
async def sql_health_engine ( engine : str ):
if engine == "hadoop" :
return health_hadoop ( _run_trino )
if engine not in SOURCE_ENGINES :
return JSONResponse ({ "error" : "unknown engine" }, status_code = 404 )
all_h = await sql_health ()
return all_h . get ( engine ) or { "ok" : False , "error" : "not found" }
2026-06-27 15:42:37 +00:00
@router.get ( "/catalog/ {engine} " )
async def get_catalog ( engine : str ):
2026-06-27 19:37:50 +00:00
if engine not in SOURCE_ENGINES and engine not in LAKE_ENGINES :
2026-06-27 15:42:37 +00:00
return JSONResponse ({ "error" : "unknown engine" }, status_code = 404 )
try :
return _catalog ( engine )
except Exception as exc :
return JSONResponse ({ "error" : str ( exc )[: 500 ]}, status_code = 502 )
2026-06-25 01:10:57 +00:00
2026-06-27 15:42:37 +00:00
@router.get ( "/sample/ {engine} " )
2026-06-27 19:37:50 +00:00
async def get_sample (
engine : str ,
object : str = Query ( ... , min_length = 1 ),
limit : int = Query ( 100 , ge = 1 , le = 500 ),
offset : int = Query ( 0 , ge = 0 ),
):
if engine not in SOURCE_ENGINES and engine not in LAKE_ENGINES :
2026-06-27 15:42:37 +00:00
return JSONResponse ({ "error" : "unknown engine" }, status_code = 404 )
try :
2026-06-27 19:37:50 +00:00
result = _sample ( engine , object , limit , offset )
2026-06-27 15:42:37 +00:00
if not result . get ( "ok" ):
return JSONResponse ( result , status_code = 422 )
2026-06-27 19:37:50 +00:00
pks = _object_primary_keys ( engine , object )
total = _table_row_count ( engine , object )
return {
** result ,
"engine" : engine ,
"object" : object ,
"offset" : offset ,
"limit" : limit ,
"total_count" : total ,
"primary_keys" : pks ,
"cdc" : engine in ( "postgres" , "mysql" , "mongodb" ),
"editable" : bool ( pks ),
}
2026-06-27 15:42:37 +00:00
except Exception as exc :
return JSONResponse ({ "ok" : False , "error" : str ( exc )[: 500 ]}, status_code = 502 )
2026-06-25 01:33:48 +00:00
2026-06-27 15:48:56 +00:00
@router.get ( "/graph/neo4j" )
async def get_neo4j_graph (
limit : int = Query ( 60 , ge = 10 , le = 150 ),
rel_type : str | None = Query ( None ),
):
try :
result = _graph_neo4j ( limit , rel_type or None )
if not result . get ( "ok" ):
return JSONResponse ( result , status_code = 422 )
return result
except Exception as exc :
return JSONResponse ({ "ok" : False , "error" : str ( exc )[: 500 ]}, status_code = 502 )
2026-06-25 01:10:57 +00:00
@router.post ( "/execute" )
async def execute_sql ( body : SqlRequest ):
sql = body . sql . strip () . rstrip ( ";" )
if not sql :
return JSONResponse ({ "ok" : False , "error" : "Empty query" }, status_code = 400 )
try :
2026-06-25 01:33:48 +00:00
result = _dispatch ( body . engine , sql )
2026-06-25 01:10:57 +00:00
if not result . get ( "ok" ):
return JSONResponse ( result , status_code = 422 )
return { ** result , "engine" : body . engine , "sql" : sql }
except Exception as exc :
return JSONResponse ({ "ok" : False , "error" : str ( exc )[: 500 ]}, status_code = 502 )
2026-06-27 19:37:50 +00:00
@router.post ( "/row/update" )
async def update_row ( body : RowUpdateRequest ):
allowed_pks = _object_primary_keys ( body . engine , body . object )
if allowed_pks :
for k in body . pk :
if k not in allowed_pks :
return JSONResponse ({ "ok" : False , "error" : f "Invalid primary key column ' { k } '" }, status_code = 400 )
for k in body . changes :
if k in body . pk :
return JSONResponse ({ "ok" : False , "error" : "Cannot change primary key columns" }, status_code = 400 )
try :
return _update_row ( body . engine , body . object , body . pk , body . changes )
except Exception as exc :
return JSONResponse ({ "ok" : False , "error" : str ( exc )[: 500 ]}, status_code = 422 )
2026-06-28 17:12:29 +00:00
@router.post ( "/row/insert" )
async def insert_row ( body : RowInsertRequest ):
if not body . values :
return JSONResponse ({ "ok" : False , "error" : "No values provided" }, status_code = 400 )
try :
return _insert_row ( body . engine , body . object , body . values )
except Exception as exc :
return JSONResponse ({ "ok" : False , "error" : str ( exc )[: 500 ]}, status_code = 422 )
2026-06-25 01:10:57 +00:00
@router.post ( "/benchmark" )
async def benchmark ():
2026-06-25 01:24:35 +00:00
"""Parallel analytics: PostgreSQL OLTP vs Trino distributed engine (5M sin() rows)."""
2026-06-25 01:10:57 +00:00
results : dict [ str , Any ] = {}
for engine , sql in BENCHMARK_SQL . items ():
try :
2026-06-25 01:33:48 +00:00
results [ engine ] = _dispatch ( engine , sql )
2026-06-25 01:10:57 +00:00
results [ engine ][ "sql" ] = sql
except Exception as exc :
results [ engine ] = { "ok" : False , "error" : str ( exc )[: 300 ], "sql" : sql }
pg_ms = results . get ( "postgres" , {}) . get ( "elapsed_ms" )
tr_ms = results . get ( "trino" , {}) . get ( "elapsed_ms" )
faster = None
2026-06-25 01:33:48 +00:00
speedup = None
2026-06-25 01:10:57 +00:00
if pg_ms and tr_ms :
faster = "trino" if tr_ms < pg_ms else "postgres"
2026-06-25 01:33:48 +00:00
speedup = round ( pg_ms / tr_ms , 2 ) if tr_ms < pg_ms else round ( tr_ms / pg_ms , 2 )
2026-06-25 01:10:57 +00:00
return {
"ok" : True ,
"postgres" : results . get ( "postgres" ),
"trino" : results . get ( "trino" ),
"comparison" : {
"postgres_ms" : pg_ms ,
"trino_ms" : tr_ms ,
"faster" : faster ,
"speedup_factor" : speedup ,
2026-06-25 01:24:35 +00:00
"note" : "5M row parallel compute — Trino distributes across workers; PostgreSQL single-node OLTP" ,
2026-06-25 01:10:57 +00:00
},
}