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)