Upload 3 files

.env

@@ -0,0 +1,4 @@
+GOOGLE_API_KEY=AQ.Ab8RN6LHRAnA-2C98mESLiiFdyruPB60QVrfekf-vYL1HribMQ
+SERVER=DESKTOP-322IVR3\SQLEXPRESS
+DATABASE=test
+DRIVER=ODBC Driver 17 for SQL Server

app.py

@@ -0,0 +1,587 @@
+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")
+
+if not api_key:
+    st.error("Missing required environment variable: GOOGLE_API_KEY")
+    st.stop()
+
+# -----------------------------
+# Gemini Client
+# -----------------------------
+
+client = genai.Client(api_key=api_key)
+
+# -----------------------------
+# 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..."):
+            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)

requirements.txt

@@ -0,0 +1,7 @@
+google-genai
+python-dotenv
+sqlalchemy
+pyodbc
+pandas
+plotly
+streamlit