Blame for app.py

import os
import re
import logging
from urllib.parse import quote_plus

import pandas as pd
import streamlit as st
import plotly.express as px

from dotenv import load_dotenv
from sqlalchemy import create_engine, text
from google import genai

# -----------------------------
# Page Config (must be the first Streamlit call)
# -----------------------------

st.set_page_config(page_title="AI SQL Assistant", layout="wide")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ai_sql_assistant")

# -----------------------------
# Config
# -----------------------------

# NOTE: confirm the exact Gemini model id you have access to
# (this is intentionally kept in one place so it's easy to change).
GEMINI_MODEL = "models/gemini-3.5-flash"

MAX_ROWS = 1000  # hard cap injected into generated SELECTs to avoid huge result sets

# Statement-level and keyword-level SQL guardrails.
# App-layer filtering is defense-in-depth only — whichever DB login is used
# below should ALSO be a read-only account (db_datareader), never a login
# with write/DDL permissions.
FORBIDDEN_KEYWORDS = [
    "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE",
    "MERGE", "EXEC", "EXECUTE", "CREATE", "GRANT", "REVOKE",
    "sp_", "xp_",
]

# -----------------------------
# Load Environment Variables (used only as prefill defaults now —
# the sidebar lets you override every one of these and switch databases
# without restarting the app)
# -----------------------------
load_dotenv()

api_key = os.getenv("GOOGLE_API_KEY", "")

default_server = os.getenv("SERVER", "")
default_database = os.getenv("DATABASE", "")
default_driver = os.getenv("DRIVER", "ODBC Driver 17 for SQL Server")

client = None

try:
    if api_key:
        client = genai.Client(api_key=api_key)
except Exception as e:
    st.error(f"Gemini Initialization Error: {e}")

# -----------------------------
# Connection helpers
# -----------------------------


def build_connection_string(srv, db, drv, auth_type, username=None, password=None):
    drv_enc = drv.replace(" ", "+")

    if auth_type == "SQL Login":
        user_enc = quote_plus(username or "")
        pwd_enc = quote_plus(password or "")
        return (
            f"mssql+pyodbc://{user_enc}:{pwd_enc}@{srv}/{db}"
            f"?driver={drv_enc}"
        )

    # Windows / trusted connection
    return (
        f"mssql+pyodbc://@{srv}/{db}"
        f"?driver={drv_enc}"
        "&trusted_connection=yes"
    )


def create_and_test_engine(connection_string):
    """Build an engine and immediately test connectivity.

    Raises on failure so the caller can show a clear error."""
    eng = create_engine(connection_string, pool_pre_ping=True)
    with eng.connect():
        pass
    return eng


def load_schema(engine):
    sql = """

    SELECT

        TABLE_NAME,

        COLUMN_NAME,

        DATA_TYPE

    FROM INFORMATION_SCHEMA.COLUMNS

    ORDER BY TABLE_NAME, ORDINAL_POSITION

    """

    df = pd.read_sql(sql, engine)

    schema = ""

    for table in df["TABLE_NAME"].unique():

        schema += f"\nTable : {table}\n"

        cols = df[df["TABLE_NAME"] == table]

        for _, row in cols.iterrows():
            schema += f"- {row['COLUMN_NAME']} ({row['DATA_TYPE']})\n"

    return schema


def load_tables(engine):
    return pd.read_sql(
        """

        SELECT TABLE_NAME

        FROM INFORMATION_SCHEMA.TABLES

        WHERE TABLE_TYPE='BASE TABLE'

        """,
        engine
    )


def is_query_safe(sql: str):
    """Defense-in-depth check on top of the read-only DB login.



    Rejects multi-statement input and any forbidden keyword, matched on

    word boundaries so it can't be trivially bypassed by e.g. 'SELECTED'.



    Allows queries that start with SELECT *or* WITH (CTE), since the model

    frequently generates `WITH ... AS (...) SELECT ...` for window-function

    questions. Both are still read-only as long as no forbidden keyword

    appears anywhere in the statement.



    Returns (is_safe: bool, reason: str).

    """
    stripped = sql.strip()

    if ";" in stripped.rstrip(";"):
        return False, "Multiple SQL statements are not allowed."

    stripped_upper = stripped.upper()
    if not (stripped_upper.startswith("SELECT") or stripped_upper.startswith("WITH")):
        return False, "Only SELECT (or WITH ... SELECT) statements are allowed."

    upper_sql = stripped_upper
    for kw in FORBIDDEN_KEYWORDS:
        pattern = r"\b" + re.escape(kw.upper()) + r"\b"
        if re.search(pattern, upper_sql):
            return False, f"Query contains a forbidden keyword: {kw}"

    return True, ""


def _find_top_level_select(sql: str):
    """Find the match object for the outermost SELECT keyword (i.e. the

    one that isn't nested inside any parentheses — which is where a

    CTE's own SELECT lives). Returns None if not found.



    SQL Server does NOT allow `WITH ... AS (...)` inside a subquery

    (e.g. `SELECT TOP N * FROM ( WITH x AS (...) SELECT ... ) AS t`

    is a syntax error) — the WITH clause must start the batch. So for

    CTE queries we can't just wrap the whole statement; we have to

    inject TOP N into the final SELECT in place, leaving the WITH

    clause where it is.

    """
    depth = 0
    for m in re.finditer(r"\(|\)|\bSELECT\b", sql, flags=re.IGNORECASE):
        token = m.group()
        if token == "(":
            depth += 1
        elif token == ")":
            depth -= 1
        elif depth == 0:
            return m
    return None


def enforce_row_limit(sql: str, max_rows: int) -> str:
    """Cap the result set size by injecting TOP N into the outermost

    SELECT — the final one, not one nested inside a CTE's own

    parentheses. Works for plain SELECT statements and for

    WITH ... AS (...) SELECT ... (CTE) statements alike.

    """
    upper_sql = sql.upper()

    # If the model already limited the result set itself, leave it alone.
    if re.search(r"\bTOP\s+\d+\b", upper_sql):
        return sql

    match = _find_top_level_select(sql)
    if match is None:
        # Couldn't locate a top-level SELECT (shouldn't normally happen
        # since is_query_safe already required SELECT/WITH) — leave the
        # query as-is rather than risk producing invalid SQL.
        return sql

    insert_at = match.end()

    # T-SQL clause order is SELECT [ALL | DISTINCT] TOP (n) ...,
    # so if DISTINCT/ALL immediately follows SELECT, TOP must go after it.
    modifier_match = re.match(r"\s+(DISTINCT|ALL)\b", sql[insert_at:], flags=re.IGNORECASE)
    if modifier_match:
        insert_at += modifier_match.end()

    return f"{sql[:insert_at]} TOP {max_rows}{sql[insert_at:]}"


def reset_query_results():
    """Clear last query/results — called whenever the DB connection changes."""
    st.session_state.sql = None
    st.session_state.df = None


# -----------------------------
# Session State
# -----------------------------

for key, default in [
    ("engine", None),
    ("schema", None),
    ("conn_label", None),
    ("df", None),
    ("sql", None),
]:
    if key not in st.session_state:
        st.session_state[key] = default

# -----------------------------
# Sidebar — Connection Panel
# -----------------------------

with st.sidebar:

    st.header("🔌 Database Connection")

    with st.form("connection_form"):

        server_input = st.text_input("Server", value=default_server)
        database_input = st.text_input("Database", value=default_database)
        driver_input = st.text_input("Driver", value=default_driver)

        auth_type = st.radio(
            "Authentication",
            ["Windows (Trusted)", "SQL Login"],
            horizontal=True,
        )

        username_input = ""
        password_input = ""

        if auth_type == "SQL Login":
            username_input = st.text_input("Username")
            password_input = st.text_input("Password", type="password")

        connect_clicked = st.form_submit_button("Connect")

    if connect_clicked:

        if not server_input.strip() or not database_input.strip():
            st.error("Server and Database are required.")

        elif auth_type == "SQL Login" and (not username_input.strip() or not password_input):
            st.error("Username and Password are required for SQL Login.")

        else:

            conn_string = build_connection_string(
                server_input.strip(),
                database_input.strip(),
                driver_input.strip(),
                auth_type,
                username_input.strip(),
                password_input,
            )

            try:
                with st.spinner("Connecting..."):
                    new_engine = create_and_test_engine(conn_string)
                    new_schema = load_schema(new_engine)

                st.session_state.engine = new_engine
                st.session_state.schema = new_schema
                st.session_state.conn_label = f"{server_input} / {database_input}"

                reset_query_results()

                st.success(f"Connected to {database_input} on {server_input}")
                st.rerun()

            except Exception as e:
                logger.exception("Database connection failed")
                st.error("Connection failed. Check Server / Database / Driver / credentials.")
                st.exception(e)

    st.divider()

    if st.session_state.engine is not None:

        st.header("Database")
        st.write(st.session_state.conn_label)

        st.header("Tables")

        try:
            tables = load_tables(st.session_state.engine)
            st.dataframe(tables, use_container_width=True)
        except Exception as e:
            st.warning("Could not load table list.")
            st.exception(e)

    else:
        st.info("Not connected. Fill in the form above and click Connect.")

# -----------------------------
# Main Area
# -----------------------------

st.title("🤖 AI SQL Assistant")

st.caption("Ask questions in natural language.")

if st.session_state.engine is None:
    st.info("👈 Connect to a database in the sidebar to get started.")
    st.stop()

engine = st.session_state.engine
schema = st.session_state.schema

question = st.text_input(
    "Ask your database",
    placeholder="Example : Show department wise salary"
)

# -----------------------------
# Button
# -----------------------------

if st.button("Run"):

    if question.strip() == "":
        st.warning("Please enter a question.")
        st.stop()

    prompt = f"""

You are a SQL Server Expert.



Database Schema



{schema}



Rules



1. Return ONLY SQL.

2. SQL Server syntax only.

3. Use SELECT statements only (a WITH ... SELECT CTE is allowed).

4. Never use INSERT.

5. Never use UPDATE.

6. Never use DELETE.

7. Never use DROP.

8. Never use ALTER.

9. Never explain.

10. Never use markdown.



Question



{question}

"""

    # -----------------------------
    # Generate SQL
    # -----------------------------

    try:

        with st.spinner("Generating SQL..."):

            if client is None:
                st.error("Gemini API key is missing or invalid.")
                st.stop()

            response = client.models.generate_content(
                model=GEMINI_MODEL,
                contents=prompt
            )

            sql = response.text.strip()
            sql = sql.replace("```sql", "")
            sql = sql.replace("```", "")
            sql = sql.strip()

    except Exception as e:

        logger.exception("Gemini generation failed")
        st.error("Gemini Error")
        st.exception(e)
        st.stop()

    # -----------------------------
    # Safety Check
    # -----------------------------

    safe, reason = is_query_safe(sql)

    if not safe:

        st.session_state.sql = sql
        st.session_state.df = None

        st.subheader("Generated SQL")
        st.code(sql, language="sql")

        st.error(reason)

        st.stop()

    sql = enforce_row_limit(sql, MAX_ROWS)

    # -----------------------------
    # Execute SQL
    # -----------------------------

    try:

        with st.spinner("Running query..."):
            df = pd.read_sql(text(sql), engine)

    except Exception as e:

        logger.exception("SQL execution failed: %s", sql)

        st.session_state.sql = sql
        st.session_state.df = None

        st.subheader("Generated SQL")
        st.code(sql, language="sql")

        st.error("SQL Error")

        st.exception(e)

        st.stop()

    # -----------------------------
    # Save results to session state
    # -----------------------------

    st.session_state.sql = sql
    st.session_state.df = df

# -----------------------------
# Display SQL + Result (persisted via session state
# so switching chart type below doesn't clear it)
# -----------------------------

if st.session_state.sql is not None:

    st.subheader("Generated SQL")

    st.code(st.session_state.sql, language="sql")

if st.session_state.df is not None:

    df = st.session_state.df

    st.subheader("Result")

    if df.empty:

        st.info("No records found.")

    else:

        st.dataframe(df, use_container_width=True)

        # -----------------------------
        # Chart Type Switcher
        # -----------------------------

        numeric = list(df.select_dtypes(include="number").columns)

        text_cols = list(
            df.select_dtypes(include=["object", "category"]).columns
        )

        all_cols = list(df.columns)

        if len(numeric) > 0 and len(all_cols) > 0:

            st.subheader("Chart")

            chart_col1, chart_col2, chart_col3 = st.columns(3)

            with chart_col1:
                chart_type = st.selectbox(
                    "Chart type",
                    ["Bar", "Line", "Area", "Scatter", "Pie", "Histogram", "Box"],
                    key="chart_type"
                )

            # sensible defaults, same logic as before
            default_x = text_cols[0] if text_cols else all_cols[0]
            default_y = numeric[0]
            default_y2 = numeric[1] if len(numeric) >= 2 else numeric[0]

            x_axis = default_x
            y_axis = default_y

            if chart_type in ["Bar", "Line", "Area"]:

                with chart_col2:
                    x_axis = st.selectbox("X-axis", all_cols, index=all_cols.index(default_x))

                with chart_col3:
                    y_axis = st.selectbox("Y-axis", numeric, index=numeric.index(default_y))

            elif chart_type == "Scatter":

                with chart_col2:
                    x_axis = st.selectbox(
                        "X-axis", numeric,
                        index=numeric.index(default_y)
                    )

                with chart_col3:
                    y_axis = st.selectbox(
                        "Y-axis", numeric,
                        index=numeric.index(default_y2)
                    )

            elif chart_type == "Pie":

                with chart_col2:
                    x_axis = st.selectbox("Labels", all_cols, index=all_cols.index(default_x))

                with chart_col3:
                    y_axis = st.selectbox("Values", numeric, index=numeric.index(default_y))

            elif chart_type in ["Histogram", "Box"]:

                with chart_col2:
                    y_axis = st.selectbox("Column", numeric, index=numeric.index(default_y))

            # -----------------------------
            # Build the selected chart
            # -----------------------------

            fig = None

            try:

                if chart_type == "Bar":
                    fig = px.bar(df, x=x_axis, y=y_axis, title=f"{y_axis} by {x_axis}")

                elif chart_type == "Line":
                    fig = px.line(df, x=x_axis, y=y_axis, title=f"{y_axis} by {x_axis}")

                elif chart_type == "Area":
                    fig = px.area(df, x=x_axis, y=y_axis, title=f"{y_axis} by {x_axis}")

                elif chart_type == "Scatter":
                    fig = px.scatter(df, x=x_axis, y=y_axis, title=f"{y_axis} vs {x_axis}")

                elif chart_type == "Pie":
                    fig = px.pie(df, names=x_axis, values=y_axis, title=f"{y_axis} by {x_axis}")

                elif chart_type == "Histogram":
                    fig = px.histogram(df, x=y_axis, title=f"Distribution of {y_axis}")

                elif chart_type == "Box":
                    fig = px.box(df, y=y_axis, title=f"Box plot of {y_axis}")

                if fig is not None:
                    st.plotly_chart(fig, use_container_width=True)

            except Exception as e:
                st.warning("Could not render this chart for the selected columns.")
                st.exception(e)