Spaces:
Runtime error
Runtime error
| 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) |