app.py

import os
import sqlite3
import joblib
import pandas as pd
import numpy as np
import tensorflow as tf
from flask import Flask, render_template, request, jsonify

from src.benchmark_db import PAGILA_DB_PATH, TPCH_DB_PATH, setup_benchmark_databases
from src.feature_extractor import extract_full_feature_vector
from src.dataset_generator import FEATURE_COLUMNS, SCALER_PATH
from src.train import MODEL_PATH, preprocess_features
from src.explainer import generate_query_explanations

app = Flask(__name__)

# Lazy global model and scaler references
model = None
scaler = None

def load_ai_assets():
    global model, scaler
    if scaler is None and os.path.exists(SCALER_PATH):
        scaler = joblib.load(SCALER_PATH)
    if model is None and os.path.exists(MODEL_PATH):
        model = tf.keras.models.load_model(MODEL_PATH)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/analyze', methods=['POST'])
def analyze_queries():
    load_ai_assets()

    data = request.get_json()
    queries = data.get('queries', [])

    # Filter out empty inputs
    valid_queries = [q.strip() for q in queries if q and q.strip()]

    if len(valid_queries) < 2:
        return jsonify({'error': 'Please provide at least 2 candidate SQL queries for comparison.'}), 400

    if len(valid_queries) > 8:
        return jsonify({'error': 'Maximum 8 candidate queries allowed per evaluation.'}), 400

    # Duplicate query validation
    normalized_queries = [q.rstrip(';').strip().lower() for q in valid_queries]
    if len(set(normalized_queries)) < len(normalized_queries):
        return jsonify({'error': 'Duplicate SQL queries detected! Please enter unique candidate query variations.'}), 400

    setup_benchmark_databases()

    candidate_data = []
    feature_rows = []

    for idx, query in enumerate(valid_queries):
        # Select target real-world DB based on query context
        q_upper = query.upper()
        if any(kw in q_upper for kw in ['LINEITEM', 'ORDERS', 'PARTSUPP', 'NATION', 'REGION', 'SUPPLIER']):
            db_target = TPCH_DB_PATH
        else:
            db_target = PAGILA_DB_PATH

        conn = sqlite3.connect(db_target)
        feats = extract_full_feature_vector(conn, query, execute=True)
        conn.close()
        feats['query_id'] = idx
        feats['sql_query'] = query
        candidate_data.append(feats)
        feature_rows.append(feats)

    conn.close()

    df_feats = pd.DataFrame(feature_rows)

    if model is not None and scaler is not None:
        X_scaled = preprocess_features(df_feats, scaler, fit_scaler=False)
        predicted_scores = model.predict(X_scaled).flatten() * 100.0
        predicted_scores = np.clip(predicted_scores, 5.0, 99.0)
    else:
        # Heuristic fallback if model not trained yet
        costs = df_feats['execution_time_ms'] + df_feats['has_select_star']*100 + df_feats['func_in_where']*200
        predicted_scores = 100.0 - (costs / costs.max() * 60.0)

    for i, score in enumerate(predicted_scores):
        candidate_data[i]['performance_score'] = round(float(score), 1)

    # Rank queries by predicted score descending
    best_index = int(np.argmax(predicted_scores))
    explanations = generate_query_explanations(candidate_data, best_index)

    # Prepare detailed output response
    ranked_results = []

    for i, c in enumerate(candidate_data):
        expl = explanations[i]
        ranked_results.append({
            'query_index': i + 1,
            'sql_query': c['sql_query'],
            'performance_score': c['performance_score'],
            'is_recommended': (i == best_index),
            'execution_time_ms': round(c['execution_time_ms'], 2),
            'estimated_cost': round(c['estimated_cost'], 1),
            'num_joins': c['num_joins'],
            'has_select_star': bool(c['has_select_star']),
            'func_in_where': bool(c['func_in_where']),
            'reasons': expl['reasons'],
            'warnings': expl['warnings']
        })

    # Sort results for display by performance score descending
    sorted_ranked_results = sorted(ranked_results, key=lambda x: x['performance_score'], reverse=True)

    # Calculate performance improvement over worst query
    worst_score = min(c['performance_score'] for c in candidate_data)
    best_score = max(c['performance_score'] for c in candidate_data)
    improvement = round(((best_score - worst_score) / max(1.0, worst_score)) * 100, 1)

    return jsonify({
        'status': 'success',
        'recommended_query_index': best_index + 1,
        'performance_improvement_percent': improvement,
        'ranked_results': sorted_ranked_results
    })

if __name__ == '__main__':
    load_ai_assets()
    app.run(host='127.0.0.1', port=5000, debug=False, use_reloader=False)