Merge pull request #21 from antonbricks/lakebase-streamlit

Update Streamlit OLTP docs

Author: antonbricks

docs/docs/streamlit/tables/oltp_database_connect.mdx

@@ -4,7 +4,7 @@ sidebar_position: 1
 
 # Connect an OLTP database
 
-This app connects to a [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp/) OLTP database instance. Provide the instance name, database, schema, and table.
+This app connects to a [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp/) OLTP database instance for reads and writes, e.g., of an App state. Provide the instance name, database, schema, and state table.
 
 ## Code snippet
 
@@ -14,26 +14,20 @@ import streamlit as st
 import pandas as pd
 
 from databricks.sdk import WorkspaceClient
-
 import psycopg
 from psycopg_pool import ConnectionPool
 
 
 w = WorkspaceClient()
 
 
-def _generate_token(instance_name: str) -> str:
-    cred = w.database.generate_database_credential(
-        request_id=str(uuid.uuid4()), instance_names=[instance_name]
-    )
-    return cred.token
-
-    
 class RotatingTokenConnection(psycopg.Connection):
     @classmethod
     def connect(cls, conninfo: str = "", **kwargs):
-        instance_name = kwargs.pop("_instance_name")
-        kwargs["password"] = _generate_token(instance_name)
+        kwargs["password"] = w.database.generate_database_credential(
+            request_id=str(uuid.uuid4()),
+            instance_names=[kwargs.pop("_instance_name")]
+        ).token
         kwargs.setdefault("sslmode", "require")
         return super().connect(conninfo, **kwargs)
 
@@ -45,29 +39,57 @@ def build_pool(instance_name: str, host: str, user: str, database: str) -> Conne
         connection_class=RotatingTokenConnection,
         kwargs={"_instance_name": instance_name},
         min_size=1,
-        max_size=10,
+        max_size=5,
         open=True,
     )
 
-    
+
 def query_df(pool: ConnectionPool, sql: str) -> pd.DataFrame:
     with pool.connection() as conn:
         with conn.cursor() as cur:
             cur.execute(sql)
+            if cur.description is None:
+                return pd.DataFrame()
             cols = [d.name for d in cur.description]
             rows = cur.fetchall()
-
     return pd.DataFrame(rows, columns=cols)
 
+
+session_id = str(uuid.uuid4())
+if "session_id" not in st.session_state:
+    st.session_state["session_id"] = session_id
+
 
 instance_name = "dbase_instance"
-database = "customer_database"
-table = "customer_core.customers_oltp"
+database = "databricks_postgres"
+schema = "public"
+table = "app_state"
 user = w.current_user.me().user_name
 host = w.database.get_database_instance(name=instance_name).read_write_dns
 
 pool = build_pool(instance_name, host, user, database)
-df = query_df(pool, f'SELECT * FROM {table} LIMIT 100')
+
+with pool.connection() as conn:
+    with conn.cursor() as cur:
+        cur.execute(f"""
+        CREATE TABLE IF NOT EXISTS {schema}.{table} (
+            session_id TEXT,
+            key TEXT,
+            value TEXT,
+            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            PRIMARY KEY (session_id, key)
+        )
+        """)
+
+        cur.execute(f"""
+            INSERT INTO app_state (session_id, key, value, updated_at)
+            VALUES ('{session_id}', 'feedback_message', 'true', CURRENT_TIMESTAMP)
+            ON CONFLICT (session_id, key) DO UPDATE
+            SET value = EXCLUDED.value,
+                updated_at = CURRENT_TIMESTAMP
+        """)
+
+df = query_df(pool, f"SELECT * FROM {schema}.{table} WHERE session_id = '{session_id}'")
 st.dataframe(df)
 ```
 
@@ -90,7 +112,7 @@ Then, your [app service principal](https://docs.databricks.com/aws/en/dev-tools/
 
 GRANT CONNECT ON DATABASE databricks_postgres TO "099f0306-9e29-4a87-84c0-3046e4bcea02";
 GRANT USAGE ON SCHEMA public TO "099f0306-9e29-4a87-84c0-3046e4bcea02";
-GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE quotes_history TO "099f0306-9e29-4a87-84c0-3046e4bcea02";
+GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE app_state TO "099f0306-9e29-4a87-84c0-3046e4bcea02";
 
 See [this guide](https://docs.databricks.com/aws/en/oltp/pg-roles?language=PostgreSQL#create-postgres-roles-and-grant-privileges-for-databricks-identities) for more information.
 

streamlit/views/oltp_database_connect.py

@@ -11,7 +11,7 @@
 st.header("OLTP Database", divider=True)
 st.subheader("Connect a table")
 st.write(
-    "This app connects to a [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp/) OLTP database instance for reads and writes, e.g., of App state. "
+    "This app connects to a [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp/) OLTP database instance for reads and writes, e.g., of an App state. "
     "Provide the instance name, database, schema, and state table."
 )