Merge pull request #56 from antonbricks/ttl-sql-injection-genie

Implement TTL in SQL connections, mitigate SQL Injections, & expand Genie - Streamlit, Dash

Author: antonbricks

dash/pages/tables_edit.py

@@ -33,9 +33,27 @@ def read_table(table_name: str, conn) -> pd.DataFrame:
 
 def insert_overwrite_table(table_name: str, df: pd.DataFrame, conn):
     with conn.cursor() as cursor:
-        rows = list(df.itertuples(index=False))
-        values = ",".join([f"({','.join(map(repr, row))})" for row in rows])
-        cursor.execute(f"INSERT OVERWRITE {table_name} VALUES {values}")
+        rows = list(df.itertuples(index=False, name=None))
+        if not rows:
+            return
+
+        cols = list(df.columns)
+        num_cols = len(cols)
+        params = {}
+        values_sql_parts = []
+        p = 0
+        for row in rows:
+            ph = []
+            for v in row:
+                key = f"p{p}"
+                ph.append(f":{key}")
+                params[key] = v
+                p += 1
+            values_sql_parts.append("(" + ",".join(ph) + ")")
+
+        values_sql = ",".join(values_sql_parts)
+        col_list_sql = ",".join(cols)
+        cursor.execute(f"INSERT OVERWRITE {table_name} ({col_list_sql}) VALUES {values_sql}", params)
 
 def layout():
     return dbc.Container([
@@ -110,9 +128,27 @@ def read_table(table_name: str, conn) -> pd.DataFrame:
 
 def insert_overwrite_table(table_name: str, df: pd.DataFrame, conn):
     with conn.cursor() as cursor:
-        rows = list(df.itertuples(index=False))
-        values = ",".join([f"({','.join(map(repr, row))})" for row in rows])
-        cursor.execute(f"INSERT OVERWRITE {table_name} VALUES {values}")
+        rows = list(df.itertuples(index=False, name=None))
+        if not rows:
+            return
+
+        cols = list(df.columns)
+        num_cols = len(cols)
+        params = {}
+        values_sql_parts = []
+        p = 0
+        for row in rows:
+            ph = []
+            for v in row:
+                key = f"p{p}"
+                ph.append(f":{key}")
+                params[key] = v
+                p += 1
+            values_sql_parts.append("(" + ",".join(ph) + ")")
+
+        values_sql = ",".join(values_sql_parts)
+        col_list_sql = ",".join(cols)
+        cursor.execute(f"INSERT OVERWRITE {table_name} ({col_list_sql}) VALUES {values_sql}", params)
 
 http_path_input = "/sql/1.0/warehouses/xxxxxx"
 table_name = "catalog.schema.table"

docs/docs/dash/bi/genie_api.mdx

@@ -4,7 +4,7 @@ sidebar_position: 2
 
 # Chat with a Genie Space
 
-This app uses the [AI/BI Genie](https://www.databricks.com/product/ai-bi) [Conversations API](https://docs.databricks.com/api/workspace/genie) to let users ask questions about your data for instant insights.
+This app uses the [AI/BI Genie](https://www.databricks.com/product/ai-bi) [Conversations API](https://docs.databricks.com/api/workspace/genie) to let users ask questions about your data for instant insights (answers and table-like output). Visualizations aren't yet supported in the API.
 
 ## Code snippet
 
@@ -33,7 +33,7 @@ def process_genie_response(response):
             print(f"A: {i.query.description}")
             print(f"Data: {data}")
             print(f"Generated code: {i.query.query}")
-
+            
 
 # Configuration
 w = WorkspaceClient()

docs/docs/dash/tables/tables_edit.mdx

@@ -31,9 +31,27 @@ def read_table(table_name: str, conn) -> pd.DataFrame:
 
 def insert_overwrite_table(table_name: str, df: pd.DataFrame, conn):
     with conn.cursor() as cursor:
-        rows = list(df.itertuples(index=False))
-        values = ",".join([f"({','.join(map(repr, row))})" for row in rows])
-        cursor.execute(f"INSERT OVERWRITE {table_name} VALUES {values}")
+        rows = list(df.itertuples(index=False, name=None))
+        if not rows:
+            return
+
+        cols = list(df.columns)
+        num_cols = len(cols)
+        params = {}
+        values_sql_parts = []
+        p = 0
+        for row in rows:
+            ph = []
+            for v in row:
+                key = f"p{p}"
+                ph.append(f":{key}")
+                params[key] = v
+                p += 1
+            values_sql_parts.append("(" + ",".join(ph) + ")")
+
+        values_sql = ",".join(values_sql_parts)
+        col_list_sql = ",".join(cols)
+        cursor.execute(f"INSERT OVERWRITE {table_name} ({col_list_sql}) VALUES {values_sql}", params)
 
 http_path_input = "/sql/1.0/warehouses/xxxxxx"
 table_name = "catalog.schema.table"

docs/docs/streamlit/authentication/users_obo.mdx

@@ -20,7 +20,7 @@ def get_user_token():
     user_token = headers["X-Forwarded-Access-Token"]
     return user_token
 
-@st.cache_resource
+@st.cache_resource(ttl=300, show_spinner=True)
 def connect_with_obo(http_path, user_token):
     return sql.connect(
         server_hostname=cfg.host,

docs/docs/streamlit/bi/genie_api.mdx

@@ -4,7 +4,7 @@ sidebar_position: 1
 
 # Chat with a Genie Space
 
-This app uses the [AI/BI Genie](https://www.databricks.com/product/ai-bi) [Conversations API](https://docs.databricks.com/api/workspace/genie) to let users ask questions about your data for instant insights.
+This app uses the [AI/BI Genie](https://www.databricks.com/product/ai-bi) [Conversations API](https://docs.databricks.com/api/workspace/genie) to let users ask questions about your data for instant insights (answers and table-like output). You are also able to collect their feedback on the responses. Visualizations aren't yet supported in the API.
 
 ## Code snippet
 
@@ -13,6 +13,7 @@ Refer to the Streamlit Cookbook Genie source code for the full implementation.
 ```python title="app.py"
 import streamlit as st
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.dashboards import GenieFeedbackRating
 import pandas as pd
 
 w = WorkspaceClient()
@@ -39,6 +40,15 @@ def get_query_result(statement_id):
     )
 
 
+def collect_feedback(message_id: str):
+    rating = st.feedback("thumbs", key=f"feedback_{message_id}")
+    mapping = {1: GenieFeedbackRating.POSITIVE, 0: GenieFeedbackRating.NEGATIVE}
+    if rating and message["message_id"]:
+        w.genie.send_message_feedback(
+            genie_space_id, st.session_state.conversation_id, message["message_id"], mapping[rating]
+        )
+
+
 def process_genie_response(response):
     for i in response.attachments:
         if i.text:
@@ -50,6 +60,7 @@ def process_genie_response(response):
                 "role": "assistant", "content": i.query.description, "data": data, "code": i.query.query
             }
             display_message(message)
+    collect_feedback(response.message_id)
 
 
 if prompt := st.chat_input("Ask your question..."):

docs/docs/streamlit/tables/tables_edit.mdx

@@ -18,7 +18,7 @@ from databricks.sdk.core import Config
 cfg = Config() # Set the DATABRICKS_HOST environment variable when running locally
 
 
-@st.cache_resource(ttl="1h")
+@st.cache_resource(ttl=300, show_spinner=True)
 def get_connection(http_path):
     return sql.connect(
         server_hostname=cfg.host,
@@ -36,13 +36,30 @@ def read_table(table_name: str, conn) -> pd.DataFrame:
 def insert_overwrite_table(table_name: str, df: pd.DataFrame, conn):
     progress = st.empty()
     with conn.cursor() as cursor:
-        rows = list(df.itertuples(index=False))
-        values = ",".join([f"({','.join(map(repr, row))})" for row in rows])
+        rows = list(df.itertuples(index=False, name=None))
+        if not rows:
+            return
+
+        cols = list(df.columns)
+        num_cols = len(cols)
+        params = {}
+        values_sql_parts = []
+        p = 0
+        for row in rows:
+            ph = []
+            for v in row:
+                key = f"p{p}"
+                ph.append(f":{key}")
+                params[key] = v
+                p += 1
+            values_sql_parts.append("(" + ",".join(ph) + ")")
+
+        values_sql = ",".join(values_sql_parts)
+        col_list_sql = ",".join(cols)
+                
         with progress:
             st.info("Calling Databricks SQL...")
-        cursor.execute(f"INSERT OVERWRITE {table_name} VALUES {values}")
-    progress.empty()
-    st.success("Changes saved")
+        cursor.execute(f"INSERT OVERWRITE {table_name} ({col_list_sql}) VALUES {values_sql}", params)
 
 
 http_path_input = st.text_input(
@@ -69,7 +86,7 @@ else:
 
 :::info
 
-This sample uses Streamlit's [st.cache_resource](https://docs.streamlit.io/develop/concepts/architecture/caching#stcache_resource) with a 1-hour TTL (time-to-live) to cache the database connection across users, sessions, and reruns. The cached connection will automatically expire after 1 hour, ensuring connections don't become stale. Use Streamlit's caching decorators and TTL parameter to implement a caching strategy that works for your use case.
+This sample uses Streamlit's [st.cache_resource](https://docs.streamlit.io/develop/concepts/architecture/caching#stcache_resource) with a 300-second TTL (time-to-live) to cache the database connection across users, sessions, and reruns. The cached connection will automatically expire after 1 hour, ensuring connections don't become stale. Use Streamlit's caching decorators and TTL parameter to implement a caching strategy that works for your use case.
 
 :::
 

docs/docs/streamlit/tables/tables_read.mdx

@@ -17,7 +17,7 @@ from databricks.sdk.core import Config
 cfg = Config()  # Set the DATABRICKS_HOST environment variable when running locally
 
 
-@st.cache_resource(ttl="1h") # connection is cached
+@st.cache_resource(ttl=300, show_spinner=True) # connection is cached
 def get_connection(http_path):
     return sql.connect(
         server_hostname=cfg.host,
@@ -47,7 +47,7 @@ if http_path_input and table_name:
 
 :::info
 
-This sample uses Streamlit's [st.cache_resource](https://docs.streamlit.io/develop/concepts/architecture/caching#stcache_resource) with a 1-hour TTL (time-to-live) to cache the database connection across users, sessions, and reruns. The cached connection will automatically expire after 1 hour, ensuring connections don't become stale. Use Streamlit's caching decorators and TTL parameter to implement a caching strategy that works for your use case.
+This sample uses Streamlit's [st.cache_resource](https://docs.streamlit.io/develop/concepts/architecture/caching#stcache_resource) with a 300-second TTL (time-to-live) to cache the database connection across users, sessions, and reruns. The cached connection will automatically expire after 1 hour, ensuring connections don't become stale. Use Streamlit's caching decorators and TTL parameter to implement a caching strategy that works for your use case.
 
 :::
 

docs/docs/streamlit/visualizations/visualizations_charts.mdx

@@ -25,7 +25,7 @@ warehouses = w.warehouses.list()
 warehouse_paths = {wh.name: wh.odbc_params.path for wh in warehouses}
 
 # Connect to SQL warehouse
-@st.cache_resource
+@st.cache_resource(ttl=300, show_spinner=True)
 def get_connection(http_path):
     return sql.connect(
         server_hostname=cfg.host,

docs/docs/streamlit/visualizations/visualizations_map.mdx

@@ -25,6 +25,7 @@ warehouses = w.warehouses.list()
 warehouse_paths = {wh.name: wh.odbc_params.path for wh in warehouses}
 
 # Connect to SQL warehouse
+@st.cache_resource(ttl=300, show_spinner=True)
 def get_connection(http_path):
     return sql.connect(
         server_hostname=cfg.host,

streamlit/views/external_connections.py

@@ -4,7 +4,7 @@
 import json
 
 
-@st.cache_resource
+@st.cache_resource(ttl=300, show_spinner=True)
 def get_client_obo() -> WorkspaceClient:
     user_token = st.context.headers.get("x-forwarded-access-token")
     if not user_token: