| """Validate queries on disposable native MySQL 8/PostgreSQL 16 temporary tables.""" |
| import argparse |
| import json |
| import re |
| import time |
| from pathlib import Path |
| import sqlglot |
| from sqlglot import exp |
| from tinyquery.data import ddls,fixture,SQL_OPS |
|
|
|
|
| class NativeDB: |
| def __init__(self,backend): |
| self.backend=backend |
| if backend=='mysql': |
| import pymysql |
| self.conn=pymysql.connect(unix_socket='/run/mysqld/mysqld.sock',user='root',database='tinyquery',autocommit=True) |
| else: |
| import psycopg |
| self.conn=psycopg.connect(host='/var/run/postgresql',user='root',dbname='postgres',autocommit=True) |
| self.cursor=self.conn.cursor(); self.tables=[] |
| self.cursor.execute('SET SESSION max_execution_time=2000' if backend=='mysql' else "SET statement_timeout='2s'") |
| def setup(self,s,seed): |
| for table in reversed(self.tables): self.cursor.execute(f'DROP TABLE IF EXISTS {table}') |
| self.tables=[] |
| for ddl in ddls(s): |
| |
| ddl=re.sub(r' REFERENCES [a-zA-Z_]+\(id\)','',ddl) |
| self.cursor.execute(ddl.replace('CREATE TABLE','CREATE TEMPORARY TABLE',1)) |
| self.tables=[s['parent'],s['table']] |
| source=fixture(s,seed) |
| for table in self.tables: |
| rows=source.execute(f'SELECT * FROM {table}').fetchall() |
| self.cursor.executemany(f"INSERT INTO {table} VALUES ({','.join(['%s']*len(rows[0]))})",rows) |
| source.close() |
| def execute(self,sql): |
| statements=sqlglot.parse(sql,read='mysql' if self.backend=='mysql' else 'postgres') |
| if len(statements)!=1 or not isinstance(statements[0],exp.Select): raise ValueError('Expected one SELECT') |
| tree=statements[0] |
| if tree.find(exp.Into): raise ValueError('SELECT INTO is not permitted') |
| if any(isinstance(node,(exp.DML,exp.DDL)) for node in tree.walk()): |
| raise ValueError('Data-changing statements are not permitted inside a query') |
| if any(t.name not in self.tables for t in tree.find_all(exp.Table)): raise ValueError('Query references an unknown table') |
| for f in tree.find_all(exp.Anonymous): |
| if f.name.upper() not in ['YEAR','MONTH','LOWER','UPPER']: |
| raise ValueError('Unsupported function '+f.name) |
| self.cursor.execute(sql) |
| rows=self.cursor.fetchmany(1001) |
| if len(rows)>1000: raise ValueError('Result limit exceeded') |
| return [tuple(str(x) if not isinstance(x,(str,int,float,type(None))) else x for x in row) for row in rows] |
| def close(self): self.cursor.close(); self.conn.close() |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('--data',required=True); p.add_argument('--predictions') |
| p.add_argument('--out',required=True); p.add_argument('--limit',type=int,default=0) |
| args=p.parse_args() |
| rows=[json.loads(l) for l in Path(args.data).read_text().splitlines()] |
| predictions={} |
| if args.predictions: predictions={r['id']:r for r in map(json.loads,Path(args.predictions).read_text().splitlines())} |
| dbs={b:NativeDB(b) for b in ['mysql','supabase']}; results=[]; seen=set(); start=time.time() |
| for r in rows: |
| if r['operation'] not in SQL_OPS: continue |
| if not args.predictions: |
| key=(r['operation'],r['backend']) |
| if key in seen: continue |
| seen.add(key) |
| elif r['id'] not in predictions: continue |
| if args.limit and len(results)>=args.limit: break |
| db=dbs[r['backend']] |
| item={'id':r['id'],'backend':r['backend'],'operation':r['operation'],'gold_executes':False,'prediction_equivalent':None} |
| try: |
| gold_key=next(k for k in ('query','sql') if k in r['target']['arguments']) |
| gold=r['target']['arguments'][gold_key] |
| prediction=None |
| if args.predictions: |
| item['prediction_equivalent']=False |
| from tinyquery.evaluate import check_action,parse_action |
| action=parse_action(predictions[r['id']]['output']) |
| check_action(action,r['context']) |
| if action.get('name')!=r['target']['name']: raise ValueError('Wrong tool') |
| expected_other={k:v for k,v in r['target']['arguments'].items() if k!=gold_key} |
| actual_other={k:v for k,v in action['arguments'].items() if k!=gold_key} |
| if expected_other!=actual_other: raise ValueError('Wrong non-SQL arguments') |
| prediction=action['arguments'][gold_key] |
| equal=True |
| for seed in [1,7]: |
| db.setup(r['slots'],seed); expected=db.execute(gold) |
| if prediction is not None: |
| actual=db.execute(prediction) |
| ordered=r['operation'] in ['sort_asc','sort_desc','top','bottom'] |
| if not ordered: expected=sorted(expected,key=repr); actual=sorted(actual,key=repr) |
| equal=equal and actual==expected |
| item['gold_executes']=True |
| if prediction is not None: item['prediction_equivalent']=equal |
| except Exception as exc: item['error']=str(exc)[:300] |
| results.append(item) |
| for db in dbs.values(): db.close() |
| summary={'seconds':time.time()-start,'cases':len(results),'gold_executes':sum(r['gold_executes'] for r in results), |
| 'prediction_equivalent':sum(r['prediction_equivalent'] is True for r in results) if args.predictions else None, |
| 'engine_versions':{'mysql':'8.0.46','postgresql':'16.15'},'results':results} |
| Path(args.out).write_text(json.dumps(summary,indent=2)) |
| print(json.dumps({k:v for k,v in summary.items() if k!='results'}),flush=True) |
| for r in results: |
| if 'error' in r: print(json.dumps(r),flush=True) |
|
|
|
|
| if __name__=='__main__': main() |
|
|
