From 51c7d9d6122c3520114b065a65703b34d7eddeba Mon Sep 17 00:00:00 2001 From: Steven Yu Date: Mon, 13 Jul 2026 13:29:11 -0700 Subject: [PATCH] Add realtime-clickstream-serving demo (Real-Time Mode to Lakebase, transformWithState) --- .../00-provision-lakebase.py | 237 ++++++ .../01-realtime-sessionize.py | 434 ++++++++++ .../02-deploy-personalization-console.py | 113 +++ .../_resources/00-setup.py | 102 +++ .../_resources/bundle_config.py | 85 ++ .../realtime-clickstream-serving/app/app.py | 755 ++++++++++++++++++ .../realtime-clickstream-serving/app/app.yaml | 4 + .../app/banner.html | 31 + .../app/preflight.py | 52 ++ .../app/requirements.txt | 6 + 10 files changed, 1819 insertions(+) create mode 100644 product_demos/realtime-clickstream-serving/00-provision-lakebase.py create mode 100644 product_demos/realtime-clickstream-serving/01-realtime-sessionize.py create mode 100644 product_demos/realtime-clickstream-serving/02-deploy-personalization-console.py create mode 100644 product_demos/realtime-clickstream-serving/_resources/00-setup.py create mode 100644 product_demos/realtime-clickstream-serving/_resources/bundle_config.py create mode 100644 product_demos/realtime-clickstream-serving/app/app.py create mode 100644 product_demos/realtime-clickstream-serving/app/app.yaml create mode 100644 product_demos/realtime-clickstream-serving/app/banner.html create mode 100644 product_demos/realtime-clickstream-serving/app/preflight.py create mode 100644 product_demos/realtime-clickstream-serving/app/requirements.txt diff --git a/product_demos/realtime-clickstream-serving/00-provision-lakebase.py b/product_demos/realtime-clickstream-serving/00-provision-lakebase.py new file mode 100644 index 00000000..c3467603 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/00-provision-lakebase.py @@ -0,0 +1,237 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC ## Live Session Serving with Lakebase +# MAGIC +# MAGIC Traditional architectures rely on a separate operational data store synchronized with the lakehouse to serve live session state. **Lakebase consolidates this stack** by enabling the streaming pipeline to write session data directly into a managed PostgreSQL instance, which downstream applications consume via standard SQL. +# MAGIC +# MAGIC > **Provisioning Note:** This notebook performs a one-time initialization; subsequent executions automatically target and reuse the existing instance. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Real-time personalization on live session state + +# COMMAND ---------- + +# DBTITLE 1,Demo Configuration +dbutils.widgets.text("lakebase_instance", "clickstream-sessions", "Lakebase instance name") +dbutils.widgets.text("lakebase_db", "databricks_postgres", "Lakebase database name") +dbutils.widgets.text("lakebase_schema", "live", "Lakebase schema name") +dbutils.widgets.text("lakebase_table", "sessions", "Lakebase table name") +dbutils.widgets.text("app_service_principal", "", "App SP id for read grant (set after deploying the app)") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Sizing and Cost +# MAGIC +# MAGIC We use the smallest tier (**CU_1**) - plenty for demo loads. For production, scale up or switch to Lakebase Autoscaling (scale-to-zero). +# MAGIC +# MAGIC | Tier | RAM | When to use | +# MAGIC |-------|-------|-------------| +# MAGIC | CU_1 | 16 GB | Demo, development, small applications (this notebook's default) | +# MAGIC | CU_2 | 32 GB | Light production | +# MAGIC | CU_4 | 64 GB | Production with mid-size working sets | +# MAGIC | CU_8 | 128 GB | High concurrency or large state | +# MAGIC +# MAGIC The instance persists indefinitely until explicitly stopped. Execute `w.database.update_database_instance(stopped=True)` post-demo to prevent unnecessary compute charges. + +# COMMAND ---------- + +# MAGIC %run ./_resources/00-setup + +# COMMAND ---------- + +# MAGIC %pip install --quiet --upgrade "databricks-sdk>=0.85.0" "protobuf==5.29.5" "psycopg[binary]>=3.0" + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_resources/00-setup + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 1/ Create (or Reuse) the Lakebase Instance + +# COMMAND ---------- + +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.database import DatabaseInstance +from databricks.sdk.errors.platform import NotFound, ResourceConflict, ResourceAlreadyExists +import time + +w = WorkspaceClient() + +try: + inst = w.database.get_database_instance(name=lakebase_instance) + print(f"Reusing existing Lakebase instance: {inst.name} (state={inst.state})") + # If a previous run stopped the instance, resume it so the wait loop below can reach AVAILABLE. + if getattr(inst, "effective_stopped", False) or str(inst.state).upper().endswith("STOPPED"): + print("Instance is stopped, resuming it...") + w.database.update_database_instance( + name=lakebase_instance, + database_instance=DatabaseInstance(name=lakebase_instance, stopped=False), + update_mask="stopped", + ) +except NotFound: + print(f"Creating new Lakebase instance: {lakebase_instance}") + # create_database_instance hands back a long-running-operation handle, not the instance itself, + # so we do not read .state off it here. The wait loop just below re-fetches the instance and polls. + w.database.create_database_instance( + database_instance=DatabaseInstance( + name=lakebase_instance, + capacity="CU_1", + stopped=False, + ) + ) + print("Submitted. Waiting for it to become available...") + +# Wait until ready +for _ in range(60): + inst = w.database.get_database_instance(name=lakebase_instance) + if str(inst.state).upper().endswith("AVAILABLE") or str(inst.state).upper() == "AVAILABLE": + break + print(f" state={inst.state} - waiting...") + time.sleep(15) + +print(f"\nReady: {inst.name}") +print(f" DNS: {inst.read_write_dns}") +print(f" Capacity: {inst.capacity}") +print(f" State: {inst.state}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 2/ Provision the Target State Table +# MAGIC +# MAGIC The streaming sink executes low-latency upserts into this target table, maintaining a single mutable record per active user that updates continuously as the session evolves. This table serves as the real-time feature store queried by the personalization layer during user interactions. +# MAGIC +# MAGIC ### Materialized State Attributes: +# MAGIC * **Location context:** Current application surface. +# MAGIC * **Funnel progression:** Derived user intent state. +# MAGIC * **Engagement score:** Quantified real-time activity metric. +# MAGIC * **Assistance status:** Active flag for support surface interactions. + +# COMMAND ---------- + +import psycopg, uuid + +def _credential(): + cred = w.database.generate_database_credential( + request_id=str(uuid.uuid4()), + instance_names=[lakebase_instance], + ) + return cred.token + +def _connect(dbname: str): + return psycopg.connect( + f"host={inst.read_write_dns} " + f"dbname={dbname} " + f"user={w.current_user.me().user_name} " + f"password={_credential()} " + f"sslmode=require" + ) + +# Targets the default 'databricks_postgres' database; no CREATE DATABASE statement required. +# The 'AVAILABLE' state status slightly precedes network socket readiness. +# Implement connection retries to prevent premature initialization failures. +for _attempt in range(40): + try: + _probe = _connect(lakebase_db); _probe.close(); break + except Exception as _e: + print(f" waiting for the Postgres endpoint to accept connections... ({str(_e)[:70]})") + time.sleep(15) + +with _connect(lakebase_db) as app_conn: + with app_conn.cursor() as cur: + cur.execute(f"CREATE SCHEMA IF NOT EXISTS {lakebase_schema}") + cur.execute(f""" + CREATE TABLE IF NOT EXISTS {lakebase_schema}.{lakebase_table} ( + user_id text PRIMARY KEY, + click_count bigint, -- events seen this session + start_time bigint, + end_time bigint, + status text, -- online / offline + current_surface text, -- the app surface the user is on right now + funnel_stage text, -- browsing / engaged / converting + engagement_score bigint, -- 0-100, derived from activity and funnel progress + needs_help boolean, -- has the user hit a support surface + last_updated timestamptz NOT NULL DEFAULT now() + ) + """) + app_conn.commit() +print(f" schema + table ready: {lakebase_db}.{lakebase_schema}.{lakebase_table}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 3/ Initial State Verification + +# COMMAND ---------- + +with _connect(lakebase_db) as app_conn: + with app_conn.cursor() as cur: + cur.execute(f"SELECT COUNT(*) FROM {lakebase_schema}.{lakebase_table}") + n = cur.fetchone()[0] + cur.execute(f"""SELECT column_name, data_type + FROM information_schema.columns + WHERE table_schema=%s AND table_name=%s + ORDER BY ordinal_position""", (lakebase_schema, lakebase_table)) + cols = cur.fetchall() +print(f"Rows in {lakebase_schema}.{lakebase_table}: {n}") +print("Schema:") +for c, t in cols: + print(f" {c}: {t}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## 4/ Grant Service Principal Read Access +# MAGIC +# MAGIC The downstream Personalization Console application authenticates via a dedicated Service Principal (SP). This cell configures the required role-based access control (RBAC). +# MAGIC +# MAGIC ### Infrastructure Dependencies +# MAGIC * **Execution Order:** This step must be run **after** application deployment. The SP's PostgreSQL role is only provisioned when the application's `database` resource attaches. +# MAGIC * **Required Privileges:** Grants `USAGE` on the target schema and `SELECT` on the table. +# MAGIC +# MAGIC ### Execution Instructions +# MAGIC 1. Input the application SP's Client ID into the `app_service_principal` widget. +# MAGIC 2. Execute the cell. +# MAGIC +# MAGIC > **Note:** This operation is idempotent and can be safely re-executed to reassert or remediate access control states. + +# COMMAND ---------- + +# Provisions read privileges for the application Service Principal using DBDemos.grant_app_sp_read. +# Automatically skips if the widget is unpopulated. Post-deployment, populate with the SP client ID and re-run. +app_sp = dbutils.widgets.get("app_service_principal").strip() +if not app_sp: + print("app_service_principal widget is empty, skipping the grant. Set it to your app SP's client id " + "and re-run this cell after deploying the Personalization Console app.") +else: + DBDemos.grant_app_sp_read(w, lakebase_instance, lakebase_db, lakebase_schema, lakebase_table, app_sp) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC **Connection info recorded for the streaming notebook:** +# MAGIC - Instance name: `{lakebase_instance}` (widgets carry this through) +# MAGIC - Target table: `{lakebase_db}.{lakebase_schema}.{lakebase_table}` +# MAGIC +# MAGIC ### Next: [01-realtime-sessionize]($./01-realtime-sessionize) - the rate source feeds Real-Time Mode sessionization, which writes to Lakebase. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Production Optimization: Lakebase Autoscaling +# MAGIC +# MAGIC For production workloads, leverage **Lakebase Autoscaling** to enable scale-to-zero cost efficiencies. This mechanism automatically suspends the PostgreSQL compute instance during idle periods and wakes it transparently upon the first incoming query. +# MAGIC +# MAGIC * **Code Impact:** The core streaming pipeline logic remains completely unchanged. +# MAGIC * **Provisioning Adjustment:** Update the initialization notebook to invoke `w.postgres.create_project()` instead of `w.database.create_database_instance()`. +# MAGIC +# MAGIC For comprehensive implementation details, reference the [Lakebase Autoscaling Documentation](https://docs.databricks.com/aws/en/oltp/projects/autoscaling). \ No newline at end of file diff --git a/product_demos/realtime-clickstream-serving/01-realtime-sessionize.py b/product_demos/realtime-clickstream-serving/01-realtime-sessionize.py new file mode 100644 index 00000000..a9cb6c5e --- /dev/null +++ b/product_demos/realtime-clickstream-serving/01-realtime-sessionize.py @@ -0,0 +1,434 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # Real-Time Personalization on Live Session State +# MAGIC +# MAGIC +# MAGIC Spark Real-Time Mode sessionizes live clickstreams, continuously updating active user state in Lakebase. A personalization service queries this fresh state via single-point lookups to serve next-best actions with low-millisecond latency. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Real-time personalization on live session state + +# COMMAND ---------- + +# DBTITLE 1,Demo Configuration (Set These, Then Run All) +# Set widgets below before running all cells. Sliding-cohort overlap +# approximates active concurrency at ~1.5–2x the base cohort size. +dbutils.widgets.text("lakebase_instance", "clickstream-sessions", "Lakebase instance name") +dbutils.widgets.text("lakebase_db", "databricks_postgres", "Lakebase database name") +dbutils.widgets.text("lakebase_schema", "live", "Lakebase schema name") +dbutils.widgets.text("lakebase_table", "sessions", "Lakebase table name") +dbutils.widgets.text("total_users", "200", "Total users (universe)") +dbutils.widgets.text("concurrent_users", "120", "Concurrent users (live, approx)") +dbutils.widgets.text("events_per_second", "50", "Events per second") +dbutils.widgets.dropdown("reset_data", "true", ["true", "false"], "Reset data on run") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Real-Time Mode & Lakebase +# MAGIC +# MAGIC **Ideal Use Cases:** +# MAGIC +# MAGIC - Sub-Second SLAs: Applications requiring real-time response to user behavior (e.g., instant personalization, inline fraud mitigation, immediate churn intervention). +# MAGIC - Unified Pipelines: Architectures that consolidate cold-path warehouse ingestion and hot-path processing from a single event stream (Kafka, Kinesis, Event Hubs). +# MAGIC - Point Lookups: Read patterns limited to low-latency key-value queries for specific entity states rather than full-table scans. +# MAGIC + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Live Clickstream Ingestion +# MAGIC +# MAGIC This architecture utilizes Spark's built-in rate source to synthesize live event data. Because Real-Time Mode natively supports the rate source, the pipeline executes over a true low-latency continuous path rather than a simulated micro-batch framework. +# MAGIC +# MAGIC - Event Schema: Each payload captures two core behavioral dimensions: surface (application boundary context) and action (user interaction). +# MAGIC - Session Simulation: The pipeline implements a sliding user cohort to simulate realistic session lifecycles (concurrently opening, deepening, and closing sessions) to accurately test downstream personalization logic. +# MAGIC - Production Transition: To migrate to production, swap the rate source with the Kafka configuration lines provided at the bottom of the cell. + +# COMMAND ---------- + +# MAGIC %pip install --quiet --upgrade "databricks-sdk>=0.102.0" "protobuf==5.29.5" "psycopg[binary]==3.2.9" + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_resources/00-setup + +# COMMAND ---------- + +# DBTITLE 1,Reset for a Clean Run (Clear Checkpoint + Truncate the Live Table) +# Clean initialization requires purging the checkpoint directory and target table. +# Set 'reset_data=false' via the configuration widget to preserve state across runs. + +if dbutils.widgets.get("reset_data") == "true": + import uuid as _uuid, psycopg as _pg + from databricks.sdk import WorkspaceClient as _WC + _chkpt = f"/Volumes/{catalog}/{db}/_checkpoints/realtime_sessions" + try: + dbutils.fs.rm(_chkpt, recurse=True) + print(f"[reset] cleared checkpoint {_chkpt}") + except Exception as _e: + print(f"[reset] checkpoint clear skipped ({_e})") + _w0 = _WC() + _i0 = _w0.database.get_database_instance(name=lakebase_instance) + _c0 = _w0.database.generate_database_credential( + request_id=str(_uuid.uuid4()), instance_names=[lakebase_instance]) + _cn = _pg.connect(host=_i0.read_write_dns, dbname=lakebase_db, + user=_w0.current_user.me().user_name, password=_c0.token, + sslmode="require", connect_timeout=30) + _cn.autocommit = True + with _cn.cursor() as _cur: + _cur.execute(f"TRUNCATE TABLE {lakebase_schema}.{lakebase_table}") + _cn.close() + print(f"[reset] truncated {lakebase_schema}.{lakebase_table}") + +# COMMAND ---------- + +import pyspark.sql.functions as F + +# Parse configuration widgets. Concurrency is modeled via sliding-window overlap: +# Concurrency ~= ACTIVE_COHORT * (session_gap / cohort_slide). +# With a fixed 30s gap and 15s slide (2x multiplier), ACTIVE_COHORT is sized at 50% of the target. +USER_POOL_SIZE = int(dbutils.widgets.get("total_users")) +ROWS_PER_SECOND = int(dbutils.widgets.get("events_per_second")) +COHORT_SLIDE_SEC = 15 # How often the live cohort rotates forward (inactivity gap below is 30s) +NUM_PARTITIONS = 4 # Rate-source partitions - raise with worker count for scale tests +_target_concurrent = int(dbutils.widgets.get("concurrent_users")) +ACTIVE_COHORT = max(2, min(_target_concurrent // 2, USER_POOL_SIZE)) +if ACTIVE_COHORT >= USER_POOL_SIZE: + print(f"[config] concurrent_users is high relative to total_users - cohort clamped to " + f"{ACTIVE_COHORT} of {USER_POOL_SIZE}. Raise total_users for real session turnover.") +print(f"[config] total_users={USER_POOL_SIZE} concurrent~{_target_concurrent} " + f"(cohort={ACTIVE_COHORT}, slide={COHORT_SLIDE_SEC}s, gap=30s) events/s={ROWS_PER_SECOND}") + +# Abstract, domain-agnostic taxonomy featuring a funnel-shaped probability distribution. +# High upper-funnel frequency ensures a realistic, multi-state user distribution (browsing/engaged/converting) +# and prevents artificial bottom-of-funnel saturation. +SURFACES = (["home"] * 5 + ["search"] * 4 + ["catalog"] * 4 + ["feature_x"] * 3 + + ["account"] * 2 + ["pricing"] * 2 + ["support"] * 1 + ["checkout"] * 1) +ACTIONS = ["view", "click", "search", "submit"] + +source_events = ( + spark.readStream + .format("rate") + .option("rowsPerSecond", ROWS_PER_SECOND) + .option("numPartitions", NUM_PARTITIONS) + .load() + # Shifts the active user ID pool forward every COHORT_SLIDE_SEC to rotate the cohort. + # Users remain active for ~30s before dropping out and exceeding the inactivity timeout. + # This deterministic churn simulates a realistic distribution of concurrent session opens and closes. + .withColumn("win", (F.unix_timestamp(F.col("timestamp")) / F.lit(COHORT_SLIDE_SEC)).cast("long")) + .withColumn("cohort_base", (F.col("win") * F.lit(ACTIVE_COHORT // 2)) % F.lit(USER_POOL_SIZE)) + .withColumn("user_id", F.concat(F.lit("user-"), + (((F.col("cohort_base") + (F.col("value") % F.lit(ACTIVE_COHORT))) % F.lit(USER_POOL_SIZE)).cast("string")))) + .withColumn("event_id", F.expr("uuid()")) + .withColumn("event_date", F.unix_timestamp(F.col("timestamp"))) + # surface = where the user is, action = what they did. Decorrelated so a user moves across surfaces + # rather than tracking action one-to-one. + .withColumn("surface", F.element_at(F.array(*[F.lit(s) for s in SURFACES]), ((F.col("value") % F.lit(len(SURFACES))) + 1).cast("int"))) + .withColumn("action", F.element_at(F.array(*[F.lit(a) for a in ACTIONS]), (((F.col("value") / F.lit(3)).cast("long") % F.lit(len(ACTIONS))) + 1).cast("int"))) + .select("user_id", "event_id", "event_date", "surface", "action") +) + + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## State Architecture: transformWithState Sessionization +# MAGIC +# MAGIC The stream ingestion layer leverages transformWithState to materialize a continuous, per-user state object consumed by low-latency personalization services. +# MAGIC +# MAGIC Session State Schema +# MAGIC - current_surface (String): Current application runtime boundary context. +# MAGIC - funnel_stage (String): Derived behavioral phase (browsing | engaged | converting). +# MAGIC - engagement_score (Integer): Weighted activity metric bounded from 0 to 100. +# MAGIC - needs_help (Boolean): Flag triggered by active support surface interaction. +# MAGIC +# MAGIC To meet sub-second SLAs, Spark Real-Time Mode introduces specific runtime behaviors distinct from micro-batching: +# MAGIC +# MAGIC - Non-Vectorized Ingestion: Data is processed row-by-row rather than via batch pandas structures. +# MAGIC - Wall-Clock Eviction: Session boundaries are managed via processing-time timers. Sessions undergo state eviction after 30 seconds of absolute wall-clock inactivity. +# MAGIC +# MAGIC This architecture encapsulates state maintenance entirely within the data pipeline, decoupling it from the target application's decision engine. + +# COMMAND ---------- + +from typing import Iterator +from datetime import datetime, timezone +from pyspark.sql import Row +from pyspark.sql.streaming import StatefulProcessor, StatefulProcessorHandle +from pyspark.sql.types import StructType, StructField, LongType, StringType, TimestampType, BooleanType + +# A session closes after this much inactivity. +SESSION_GAP_MS = 30 * 1000 + + +def _derive(click_count, current_surface): + # Live snapshot: stage and needs_help come from the user's current surface, not history. A model + # could slot in here in place of these rules. + if current_surface == "checkout": + stage = "converting" + elif current_surface in ("pricing", "account"): + stage = "engaged" + else: + stage = "browsing" + needs_help = (current_surface == "support") + bonus = 40 if stage == "converting" else (20 if stage == "engaged" else 0) + score = min(100, int(click_count) * 2 + bonus) + return stage, int(score), needs_help + + +class SessionProcessor(StatefulProcessor): + def init(self, handle: StatefulProcessorHandle) -> None: + # One ValueState struct per user: running counts plus the surface they are on right now. + self.handle = handle + state_schema = StructType([ + StructField("click_count", LongType(), True), + StructField("start_time", LongType(), True), + StructField("end_time", LongType(), True), + StructField("current_surface", StringType(), True), + ]) + self.session = handle.getValueState("session", state_schema) + + def handleInputRows(self, key, rows: Iterator[Row], timerValues) -> Iterator[Row]: + # Called per user per micro-batch. Under Real-Time Mode the rows arrive one at a time, but folding + # them into the running session looks the same either way. + (user_id,) = key + + if self.session.exists(): + click_count, start_time, end_time, current_surface = self.session.get() + else: + click_count, start_time, end_time, current_surface = 0, None, None, None + + for ev in rows: + ts = ev["event_date"] + start_time = ts if start_time is None else min(start_time, ts) + end_time = ts if end_time is None else max(end_time, ts) + click_count += 1 + current_surface = ev["surface"] # where they are right now + + if end_time is None: + return + + self.session.update((click_count, start_time, end_time, current_surface)) + + # Re-arm the processing-time inactivity timer. Clear the old one first so a user with frequent events + # does not stack timers. If nothing arrives for SESSION_GAP_MS, handleExpiredTimer fires and closes the session. + for t in self.handle.listTimers(): + self.handle.deleteTimer(t) + self.handle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + SESSION_GAP_MS) + + stage, score, needs_help = _derive(click_count, current_surface) + # last_updated is emitted on every update, so the Postgres row always reflects current freshness. + yield Row( + user_id=user_id, + click_count=int(click_count), + start_time=int(start_time), + end_time=int(end_time), + status="online", + current_surface=current_surface, + funnel_stage=stage, + engagement_score=score, + needs_help=needs_help, + last_updated=datetime.fromtimestamp(int(end_time), tz=timezone.utc), + ) + + def handleExpiredTimer(self, key, timerValues, expiredTimerInfo) -> Iterator[Row]: + # The close path: fires SESSION_GAP_MS of wall-clock after the user's last event. Emit the final + # "offline" row and drop the state for that user. + (user_id,) = key + if not self.session.exists(): + return + click_count, start_time, end_time, current_surface = self.session.get() + self.session.clear() + stage, score, needs_help = _derive(click_count, current_surface) + yield Row( + user_id=user_id, + click_count=int(click_count), + start_time=int(start_time), + end_time=int(end_time), + status="offline", + current_surface=current_surface, + funnel_stage=stage, + engagement_score=score, + needs_help=needs_help, + last_updated=datetime.fromtimestamp(int(end_time), tz=timezone.utc), + ) + + def close(self) -> None: + pass + + +output_schema = StructType([ + StructField("user_id", StringType(), False), + StructField("click_count", LongType(), False), + StructField("start_time", LongType(), False), + StructField("end_time", LongType(), False), + StructField("status", StringType(), False), + StructField("current_surface", StringType(), True), + StructField("funnel_stage", StringType(), False), + StructField("engagement_score", LongType(), False), + StructField("needs_help", BooleanType(), False), + StructField("last_updated", TimestampType(), False), +]) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Pipeline Egress: Real-Time Mode & Lakebase Sink +# MAGIC +# MAGIC The pipeline leverages the native PostgreSQL streaming sink (.writeStream.format("postgresql")), available in Databricks Runtime 18.3 and above. +# MAGIC +# MAGIC - Upsert Semantics: Executes an INSERT ... ON CONFLICT (user_id) DO UPDATE statement to apply atomic state updates directly to Lakebase. +# MAGIC - Target Routing: Resolves the Lakebase compute endpoint using the project.branch.endpoint resource path, mapping the upsertkey config to the table's user_id primary key. +# MAGIC - Stream Management: Natively handles connection credentials and fully supports low-latency Real-Time Mode execution triggers. + +# COMMAND ---------- + +sessions = ( + source_events + .groupBy("user_id") + .transformWithState( + statefulProcessor=SessionProcessor(), + outputStructType=output_schema, + outputMode="Update", # emit only the changed users each batch (the sink upserts them) + timeMode="ProcessingTime", # Real-Time Mode supports processing-time timers only - the key RTM constraint + ) +) + +# COMMAND ---------- + +# Lakebase endpoints resolve via the 'project.branch.endpoint' canonical naming scheme. +# For provisioned instances, 'project' maps to the instance name, defaulting to 'production.primary'. +sink_endpoint = f"{lakebase_instance}.production.primary" + +# Enables Real-Time Mode for continuous, sub-second streaming instead of micro-batching. +# PySpark requires a duration string; this governs the checkpoint/metadata interval +# rather than processing latency. Use trigger(processingTime="5 seconds") for micro-batch fallback. +q = ( + sessions.writeStream + .format("postgresql") + .outputMode("update") + .option("endpoint", sink_endpoint) + .option("database", lakebase_db) + .option("dbtable", f"{lakebase_schema}.{lakebase_table}") + .option("upsertkey", "user_id") + .option("checkpointLocation", f"/Volumes/{catalog}/{db}/_checkpoints/realtime_sessions") + .trigger(realTime="5 minutes") + .start() +) + +print(f"Streaming query started: {q.id}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Read Live Session State +# MAGIC +# MAGIC Give the stream a few seconds to warm up (the first batch lands ~20-30 seconds after the query starts), then run, and re-run, the cell below. It connects straight to Lakebase Postgres and reads current per-user session state, exactly the point lookup an application would issue. +# MAGIC + +# COMMAND ---------- + +import psycopg, uuid +from psycopg.rows import dict_row +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient() +inst = w.database.get_database_instance(name=lakebase_instance) +cred = w.database.generate_database_credential( + request_id=str(uuid.uuid4()), + instance_names=[lakebase_instance], +) +conn_string = ( + f"host={inst.read_write_dns} " + f"dbname={lakebase_db} " + f"user={w.current_user.me().user_name} " + f"password={cred.token} " + f"sslmode=require" +) + + +def next_best_action(row): + # The personalization decision, computed from one Postgres row. + if row["needs_help"]: + return "Offer live support" + if row["funnel_stage"] == "converting": + return "Nudge to complete (assist or incentive)" + if row["funnel_stage"] == "engaged": + return f"Surface tailored content for '{row['current_surface']}'" + return "Highlight popular feature / onboarding" + + +with psycopg.connect(conn_string, row_factory=dict_row) as conn: + with conn.cursor() as cur: + cur.execute(f""" + SELECT + COUNT(*) FILTER (WHERE status='online') AS active_users, + COUNT(*) FILTER (WHERE funnel_stage='converting') AS converting, + COUNT(*) FILTER (WHERE needs_help) AS needs_help, + MAX(last_updated) AS most_recent, + now() AS server_now + FROM {lakebase_schema}.{lakebase_table} + """) + summary = cur.fetchone() + cur.execute(f""" + SELECT user_id, current_surface, funnel_stage, engagement_score, needs_help, last_updated + FROM {lakebase_schema}.{lakebase_table} + WHERE status='online' + ORDER BY engagement_score DESC, last_updated DESC + LIMIT 8 + """) + top = cur.fetchall() + +if not summary["most_recent"]: + print("no rows yet - give the stream another few seconds, then re-run") +else: + fresh = (summary["server_now"] - summary["most_recent"]).total_seconds() + print(f"active_users={summary['active_users']} converting={summary['converting']} needs_help={summary['needs_help']}") + print(f"freshest session updated {fresh:.1f}s ago") + print() + print("highest-engagement live sessions, and the action the engine would serve:") + for r in top: + print(f" {r['user_id']:>9} on={r['current_surface']:<9} {r['funnel_stage']:<11} score={r['engagement_score']:>3} -> {next_best_action(r)}") + +# COMMAND ---------- + +# DBTITLE 1,Keep the Stream Alive When Run as a Job +# As a job task this blocks so the stream keeps serving the app until the run is cancelled. +# Interactively it returns immediately - the cluster keeps the stream alive after Run All. +if dbutils.notebook.entry_point.getDbutils().notebook().getContext().jobId().isDefined(): + q.awaitTermination() + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Validating Live Session State +# MAGIC +# MAGIC Allow the streaming query approximately 20–30 seconds to initialize and commit the first batch before executing the cell below. +# MAGIC +# MAGIC ### Execution Details +# MAGIC * **Direct Integration:** Establishes a direct connection to the Lakebase PostgreSQL instance. +# MAGIC * **Production Simulation:** Queries active, per-user session records, mirroring the exact low-latency point lookup pattern executed by downstream applications. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Architectural Extensibility: Cross-Domain Adaptation +# MAGIC +# MAGIC The underlying stateful architecture is domain-agnostic. By remapping the primary entity boundaries and context signals, this streaming pattern seamlessly adapts to other industry verticals: +# MAGIC +# MAGIC | Vertical | Primary Entity | Operational Context ("Surface") | Targeted Next-Best Action | +# MAGIC | :--- | :--- | :--- | :--- | +# MAGIC | **Banking & Fintech** | Customer | Application Screen | Real-Time Promotion / Risk Hold | +# MAGIC | **Telecommunications** | Subscriber | Self-Service Workflow | Retention Nudge / Upsell Trigger | +# MAGIC | **Digital Gaming** | Player | Active Gameplay Mode | Matchmaking Queue / Anti-Cheat Flag | +# MAGIC | **Industrial IoT** | Device | Core Operating State | Telemetry Alert / Control Command | +# MAGIC +# MAGIC ### Infrastructure Invariance +# MAGIC The foundational storage and compute components—including the stateful processor, wall-clock timer eviction logic, and PostgreSQL target sink—remain identical. Scaling to a new business domain requires modifying only the schema projection layer that maps raw input events into the per-entity state vector. \ No newline at end of file diff --git a/product_demos/realtime-clickstream-serving/02-deploy-personalization-console.py b/product_demos/realtime-clickstream-serving/02-deploy-personalization-console.py new file mode 100644 index 00000000..a897a395 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/02-deploy-personalization-console.py @@ -0,0 +1,113 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC # Deploy the Personalization Console +# MAGIC +# MAGIC +# MAGIC The Console is a Databricks App that serves a next-best action from each user's live session. It reads the sessions table in Lakebase as its own service principal. + +# COMMAND ---------- + +# MAGIC %md +# MAGIC Real-time personalization on live session state + +# COMMAND ---------- + +# MAGIC %pip install --quiet --upgrade "databricks-sdk>=0.85.0" "psycopg[binary]>=3.0" + +# COMMAND ---------- + +dbutils.library.restartPython() + +# COMMAND ---------- + +# MAGIC %run ./_resources/00-setup + +# COMMAND ---------- + +# DBTITLE 1,Demo Configuration +dbutils.widgets.text("app_name", "dbdemos-rtm-clickstream-app", "App name") +app_name = dbutils.widgets.get("app_name") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Application Provisioning +# MAGIC +# MAGIC The `database` resource binds the target Lakebase instance with `CAN_CONNECT_AND_CREATE` privileges, automatically provisioning a dedicated PostgreSQL role for the application's Service Principal. +# MAGIC +# MAGIC * **Deployment Source:** Source code is evaluated from the adjacent `app/` directory. + +# COMMAND ---------- + +import os +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.apps import ( + App, + AppResource, + AppResourceDatabase, + AppResourceDatabaseDatabasePermission, + AppDeployment, +) + +w = WorkspaceClient() +app_source_path = os.path.join(os.getcwd(), "app") + +database_resource = AppResource( + name="database", + database=AppResourceDatabase( + instance_name=lakebase_instance, + database_name=lakebase_db, + permission=AppResourceDatabaseDatabasePermission.CAN_CONNECT_AND_CREATE, + ), +) + +console_app = App( + name=app_name, + description="Personalization Console - live session state from Lakebase", + default_source_code_path=app_source_path, + resources=[database_resource], +) + +try: + w.apps.create_and_wait(app=console_app) +except Exception as e: + if "already exists" in str(e): + print(f"App '{app_name}' already exists, reusing it.") + else: + raise + +app_sp = w.apps.get(name=app_name).service_principal_client_id +print(f"App '{app_name}' ready. Service principal client id: {app_sp}") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Deploy the App Code + +# COMMAND ---------- + +w.apps.deploy_and_wait(app_name=app_name, app_deployment=AppDeployment(source_code_path=app_source_path)) +print("Deployed.") + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Service Principal Authorization +# MAGIC +# MAGIC The Personalization Console application queries the `live.sessions` table using its assigned Service Principal. This step provisions the required role-based privileges. +# MAGIC +# MAGIC * **Permissions:** Assigns `USAGE` on the schema and `SELECT` on the target table. +# MAGIC * **Implementation:** Executes shared access-control logic encapsulated in `DBDemos.grant_app_sp_read`. + +# COMMAND ---------- + +DBDemos.grant_app_sp_read(w, lakebase_instance, lakebase_db, lakebase_schema, lakebase_table, app_sp) + +# COMMAND ---------- + +# MAGIC %md +# MAGIC ## Open the Console + +# COMMAND ---------- + +print(w.apps.get(name=app_name).url) \ No newline at end of file diff --git a/product_demos/realtime-clickstream-serving/_resources/00-setup.py b/product_demos/realtime-clickstream-serving/_resources/00-setup.py new file mode 100644 index 00000000..3cec37c0 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/_resources/00-setup.py @@ -0,0 +1,102 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC ## Setup +# MAGIC Catalog/schema for tracking demo state. The actual session data lives in **Lakebase** +# MAGIC (Postgres), not the Lakehouse - that's the whole point of this demo. We use a Unity Catalog schema +# MAGIC only as a place to record the Lakebase instance name + a status log. + +# COMMAND ---------- + +dbutils.widgets.text("lakebase_instance", "clickstream-sessions", "Lakebase Provisioned instance name") +dbutils.widgets.text("lakebase_db", "databricks_postgres", "Lakebase database name (default Lakebase DB)") +dbutils.widgets.text("lakebase_schema", "live", "Lakebase schema name") +dbutils.widgets.text("lakebase_table", "sessions", "Lakebase table name") + +# COMMAND ---------- + +catalog = "main__build" +schema = db = "dbdemos_realtime_clickstream" +lakebase_instance = dbutils.widgets.get("lakebase_instance") +lakebase_db = dbutils.widgets.get("lakebase_db") +lakebase_schema = dbutils.widgets.get("lakebase_schema") +lakebase_table = dbutils.widgets.get("lakebase_table") + +# COMMAND ---------- + +# Inlined DBDemos helpers (same as the micro-batch sibling clickstream-direct-to-lakehouse). +import time + +class DBDemos: + @staticmethod + def setup_schema(catalog, db): + assert catalog not in ("hive_metastore", "spark_catalog"), "Demo requires Unity Catalog." + current = spark.sql("SELECT current_catalog() c").collect()[0]["c"] + if current != catalog: + catalogs = [r["catalog"] for r in spark.sql("SHOW CATALOGS").collect()] + if catalog not in catalogs: + try: + spark.sql(f"CREATE CATALOG IF NOT EXISTS `{catalog}`") + except Exception as e: + raise RuntimeError( + f"Could not create catalog `{catalog}`: {str(e).splitlines()[0]}\n" + f"This workspace may not allow creating catalogs (eg governed prod, or Default Storage). " + f"Set the `catalog` variable in this setup notebook to an existing Unity Catalog catalog you can create schemas in, " + f"and do not point this demo at a production catalog." + ) from None + spark.sql(f"USE CATALOG `{catalog}`") + spark.sql(f"CREATE DATABASE IF NOT EXISTS `{db}`") + spark.sql(f"USE `{catalog}`.`{db}`") + print(f"Using {catalog}.{db}") + + @staticmethod + def stop_all_streams(sleep_time=0): + if sleep_time: time.sleep(sleep_time) + for q in spark.streams.active: + try: q.stop() + except Exception as e: print(f" failed to stop stream {q.id}: {e}") + + @staticmethod + def grant_app_sp_read(w, lakebase_instance, lakebase_db, lakebase_schema, lakebase_table, app_sp): + # Grant the app service principal USAGE on the schema and SELECT on the sessions table. The SP's + # Postgres role appears once the app's database resource attaches, so poll for it before granting. + import time, uuid, psycopg + inst = w.database.get_database_instance(name=lakebase_instance) + def _connect(): + cred = w.database.generate_database_credential( + request_id=str(uuid.uuid4()), instance_names=[lakebase_instance]) + return psycopg.connect( + f"host={inst.read_write_dns} dbname={lakebase_db} " + f"user={w.current_user.me().user_name} password={cred.token} sslmode=require") + with _connect() as conn: + conn.autocommit = True + with conn.cursor() as cur: + for _ in range(20): + cur.execute("SELECT 1 FROM pg_roles WHERE rolname = %s", (app_sp,)) + if cur.fetchone() is not None: + break + print(f" waiting for the app service principal role '{app_sp}' to appear...") + time.sleep(15) + else: + raise RuntimeError( + f"Service principal role '{app_sp}' did not appear on {lakebase_instance}. " + f"Confirm the app deployed, then re-run this step.") + cur.execute(f'GRANT USAGE ON SCHEMA {lakebase_schema} TO "{app_sp}"') + cur.execute(f'GRANT SELECT ON {lakebase_schema}.{lakebase_table} TO "{app_sp}"') + cur.execute( + """SELECT privilege_type FROM information_schema.role_table_grants + WHERE table_schema=%s AND table_name=%s AND grantee=%s""", + (lakebase_schema, lakebase_table, app_sp)) + privs = sorted(r[0] for r in cur.fetchall()) + cur.execute("SELECT has_schema_privilege(%s, %s, 'USAGE')", (app_sp, lakebase_schema)) + usage_ok = cur.fetchone()[0] + print(f"Granted. '{app_sp}' now holds {privs} on {lakebase_schema}.{lakebase_table} (schema USAGE={usage_ok}).") + +# COMMAND ---------- + +DBDemos.setup_schema(catalog, db) + +# Managed volume for streaming checkpoints +spark.sql(f"CREATE VOLUME IF NOT EXISTS `{catalog}`.`{db}`.`_checkpoints`") + +import pyspark.sql.functions as F +from pyspark.sql.functions import col diff --git a/product_demos/realtime-clickstream-serving/_resources/bundle_config.py b/product_demos/realtime-clickstream-serving/_resources/bundle_config.py new file mode 100644 index 00000000..d67dfef8 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/_resources/bundle_config.py @@ -0,0 +1,85 @@ +# Databricks notebook source +# MAGIC %md +# MAGIC ## Demo bundle configuration +# MAGIC Please ignore / do not delete, only used to prep and bundle the demo. + +# COMMAND ---------- + +{ + "name": "realtime-clickstream-serving", + "category": "data-engineering", + "title": "Real-Time Clickstream Sessions Served from Lakebase (RTM + Postgres)", + "serverless_supported": False, # Spark Real-Time Mode requires classic DBR 16.2+ + "custom_schema_supported": True, + "default_catalog": "main", + "default_schema": "dbdemos_realtime_clickstream", + "description": "Sub-second sessionization with Spark Real-Time Mode, output to Lakebase (managed Postgres) so your application reads live session state via SQL.", + "custom_message": "This demo provisions a Lakebase Postgres instance and deploys a Databricks App - make sure you have permission to create both in your workspace. The streaming pipeline requires classic compute on DBR 18.3 or above.", + "fullDescription": "

Some teams need sessions in under a second - fraud holds, real-time personalization, cart-abandonment alerts, anti-cheat. The usual answer is to stand up a second streaming framework (Flink) and a separate KV store (Redis or Cassandra), and operate two pipelines.

This demo shows the Databricks-native alternative: Spark Real-Time Mode with transformWithState sessionizes events with sub-second latency, and the streaming sink writes per-user state directly into Lakebase (Databricks-managed Postgres). Your application reads sessions via a normal SELECT on a Postgres table.

You'll see:

When to use this: hot-path features that need sub-second latency and are read-heavy on the serving side.
When to use something else: if seconds-class latency is fine, see clickstream-direct-to-lakehouse (Zerobus + micro-batch).

", + "usecase": "Data Engineering", + "products": ["Real-Time Mode", "Lakebase", "Spark", "Unity Catalog"], + "related_links": [ + {"title": "Real-time mode in Structured Streaming", "url": "https://docs.databricks.com/aws/en/structured-streaming/real-time/"}, + {"title": "Lakebase Provisioned", "url": "https://docs.databricks.com/aws/en/oltp/instances/"}], + "recommended_items": ["streaming-sessionization", "clickstream-direct-to-lakehouse"], + "demo_assets": [], + "bundle": True, + "tags": [{"rtm": "Real-Time Mode"}, {"lakebase": "Lakebase"}, {"streaming": "Streaming"}], + "notebooks": [ + { + "path": "_resources/00-setup", + "pre_run": False, + "publish_on_website": False, + "add_cluster_setup_cell": False, + "title": "Setup", + "description": "Catalog + schema for demo state tracking", + "depends_on_previous": False + }, + { + "path": "00-provision-lakebase", + "pre_run": False, + "publish_on_website": True, + "add_cluster_setup_cell": True, + "title": "Provision Lakebase (Postgres)", + "description": "Stand up the Lakebase instance + sessions table that the streaming pipeline writes into.", + "depends_on_previous": False + }, + { + "path": "01-realtime-sessionize", + "pre_run": False, + "publish_on_website": True, + "add_cluster_setup_cell": True, + "title": "Real-Time Sessions in Lakebase", + "description": "Rate source feeding Real-Time Mode transformWithState (Row API, processing-time timers), with the native Postgres streaming sink that upserts per-user sessions into Lakebase. Reads back from Postgres at the end.", + "depends_on_previous": True + }, + { + "path": "02-deploy-personalization-console", + "pre_run": False, + "publish_on_website": True, + "add_cluster_setup_cell": True, + "title": "Deploy the Personalization Console", + "description": "Create and deploy the Personalization Console app with a Lakebase database resource, then grant its service principal read access to the sessions table.", + "depends_on_previous": True + }, + { + "path": "app", + "pre_run": False, + "publish_on_website": False, + "add_cluster_setup_cell": False, + "title": "Personalization Console App", + "description": "Streamlit application folder deployed by the 02 notebook.", + "depends_on_previous": False + } + ], + "cluster": { + "spark_version": "18.x-scala2.13", + "spark_conf": { + "spark.databricks.streaming.realTimeMode.enabled": "true", + "spark.sql.shuffle.partitions": "4" + }, + "num_workers": 2, + "single_user_name": "{{CURRENT_USER}}", + "data_security_mode": "SINGLE_USER" + } +} diff --git a/product_demos/realtime-clickstream-serving/app/app.py b/product_demos/realtime-clickstream-serving/app/app.py new file mode 100644 index 00000000..c83a4735 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/app/app.py @@ -0,0 +1,755 @@ +"""Personalization Console - serves live session state out of Lakebase. + +Each row in live.sessions is a user's current session, continuously upserted by +the Real-Time Mode pipeline. Serving a next-best action is then just a point +lookup against Postgres (clickstream-sessions / databricks_postgres / +live.sessions) - no cache, no feature-store round trip. + +The connection uses a short-lived credential minted by the running app's service +principal via WorkspaceClient(), so there is no static password. Tokens last +~1 hour, so we re-mint on a fixed cadence and reconnect on any connection error. + +Local preview: set CONSOLE_MOCK=1 to render the full UI against synthetic +session data with no Lakebase connection. The production path (CONSOLE_MOCK +unset) is unchanged - it always talks to Lakebase. +""" + +import os +import time +import uuid +from datetime import datetime +from pathlib import Path + +import pandas as pd +import streamlit as st + +INSTANCE_NAME = "clickstream-sessions" +DB_NAME = "databricks_postgres" +SCHEMA = "live" +TABLE = "sessions" +CRED_TTL_SECONDS = 45 * 60 # re-mint well before the ~1h expiry +REFRESH_SECONDS = 3 + +MOCK = os.getenv("CONSOLE_MOCK") == "1" + +# Stage palette - faint tint background plus an accent color, so it reads on +# both the light and dark Databricks themes (no hardcoded page background). +STAGE_COLORS = { + "converting": ("rgba(52, 168, 83, 0.16)", "#1e8e3e"), + "engaged": ("rgba(26, 115, 232, 0.16)", "#1a73e8"), + "browsing": ("rgba(140, 140, 150, 0.16)", "#80868b"), + "needs help": ("rgba(234, 67, 53, 0.16)", "#d93025"), +} +ACCENT = "#FF3621" # Databricks orange + +st.set_page_config( + page_title="Personalization Console", + page_icon="*", + layout="wide", +) + + +# -------------------------------------------------------------------------- +# Connection: mint a fresh Postgres credential as the app service principal. +# -------------------------------------------------------------------------- +@st.cache_resource +def _workspace_client(): + from databricks.sdk import WorkspaceClient + + return WorkspaceClient() + + +def _pg_role(w) -> str: + """Resolve the Postgres role to authenticate as. + + When the Lakebase `database` app resource is attached, Databricks injects + PGUSER = the service principal's client id, which is the role it creates on + the instance. Prefer that. Fall back to the caller's identity for local / + user-driven runs (e.g. steven.yu). + """ + pg_user = os.getenv("PGUSER") + if pg_user: + return pg_user + return w.current_user.me().user_name + + +def _build_connection(): + """Mint a credential and open a psycopg connection. Caller owns the lifecycle.""" + import psycopg + from psycopg.rows import dict_row + + w = _workspace_client() + inst = w.database.get_database_instance(name=INSTANCE_NAME) + cred = w.database.generate_database_credential( + request_id=str(uuid.uuid4()), instance_names=[INSTANCE_NAME] + ) + user = _pg_role(w) + conn_string = ( + f"host={inst.read_write_dns} dbname={DB_NAME} " + f"user={user} password={cred.token} sslmode=require" + ) + conn = psycopg.connect(conn_string, row_factory=dict_row, connect_timeout=10) + conn.autocommit = True + return conn, user + + +def get_connection(): + """Return a live connection, (re)minting the credential past its TTL.""" + state = st.session_state + now = time.time() + needs_new = ( + "pg_conn" not in state + or state.get("pg_conn") is None + or (now - state.get("pg_conn_minted_at", 0)) > CRED_TTL_SECONDS + or state.get("pg_conn").closed + ) + if needs_new: + old = state.get("pg_conn") + if old is not None: + try: + old.close() + except Exception: + pass + conn, user = _build_connection() + state["pg_conn"] = conn + state["pg_conn_minted_at"] = now + state["pg_role"] = user + return state["pg_conn"] + + +def run_query(sql: str, params=None) -> pd.DataFrame: + """Execute a query, reconnecting once on a dropped/expired connection.""" + import psycopg + + for attempt in range(2): + try: + conn = get_connection() + with conn.cursor() as cur: + cur.execute(sql, params or ()) + rows = cur.fetchall() + return pd.DataFrame(rows) + except (psycopg.OperationalError, psycopg.InterfaceError): + st.session_state["pg_conn"] = None + if attempt == 1: + raise + return pd.DataFrame() + + +# -------------------------------------------------------------------------- +# Data access. Each loader hits Lakebase in production, or returns synthetic +# rows in mock mode so the layout can be previewed locally. The SQL is the +# real serving query - unchanged from the validated pipeline. +# -------------------------------------------------------------------------- +FQ_TABLE = f"{SCHEMA}.{TABLE}" + + +def load_metrics() -> pd.DataFrame: + if MOCK: + return _mock_metrics() + return run_query( + f""" + SELECT + COUNT(*) FILTER (WHERE status = 'online') AS active_users, + COUNT(*) FILTER (WHERE funnel_stage = 'converting') AS converting, + COUNT(*) FILTER (WHERE needs_help) AS needs_help, + EXTRACT(EPOCH FROM (now() - MAX(last_updated))) AS freshest_secs, + COUNT(*) AS total_rows + FROM {FQ_TABLE} + """ + ) + + +def load_grid() -> pd.DataFrame: + if MOCK: + return _mock_grid() + return run_query( + f""" + SELECT user_id, current_surface, funnel_stage, engagement_score, + EXTRACT(EPOCH FROM (now() - last_updated)) AS last_active_secs, + status, needs_help + FROM {FQ_TABLE} + WHERE status = 'online' + ORDER BY engagement_score DESC + LIMIT 500 + """ + ) + + +def load_detail(user_id: str) -> pd.DataFrame: + if MOCK: + return _mock_detail(user_id) + return run_query(f"SELECT * FROM {FQ_TABLE} WHERE user_id = %s", (user_id,)) + + +# -------------------------------------------------------------------------- +# Mock data for local preview. Deterministic base population, lightly jittered +# each refresh so the console feels live without a backing stream. +# -------------------------------------------------------------------------- +SURFACE_WEIGHTS = { + "home": 5, "search": 4, "catalog": 4, "feature_x": 3, + "account": 2, "pricing": 2, "support": 1, "checkout": 1, +} + + +def _surface_to_stage(surface: str): + if surface == "support": + return "browsing", True + if surface == "checkout": + return "converting", False + if surface in ("pricing", "account"): + return "engaged", False + return "browsing", False + + +@st.cache_data +def _mock_base() -> pd.DataFrame: + import random + + rng = random.Random(42) + weighted = [s for s, w in SURFACE_WEIGHTS.items() for _ in range(w)] + ids, seen = [], set() + while len(ids) < 200: + uid = f"user_{rng.randint(1000, 9999)}" + if uid not in seen: + seen.add(uid) + ids.append(uid) + rows = [] + for uid in ids: + surface = rng.choice(weighted) + stage, needs = _surface_to_stage(surface) + clicks = rng.randint(1, 40) + bonus = {"converting": 30, "engaged": 15, "browsing": 0}[stage] + score = min(100, clicks * 2 + bonus) + rows.append({ + "user_id": uid, "current_surface": surface, "funnel_stage": stage, + "needs_help": needs, "click_count": clicks, "engagement_score": score, + }) + return pd.DataFrame(rows) + + +def _mock_snapshot() -> pd.DataFrame: + import random + + df = _mock_base().copy() + # New seed roughly every refresh tick so counts and freshness drift. + r = random.Random(int(time.time()) // REFRESH_SECONDS) + df["status"] = ["online" if r.random() < 0.72 else "offline" for _ in range(len(df))] + df["last_active_secs"] = [ + r.uniform(0.2, 18.0) if s == "online" else r.uniform(40.0, 320.0) + for s in df["status"] + ] + now = pd.Timestamp.now(tz="UTC") + df["last_updated"] = [now - pd.Timedelta(seconds=float(x)) for x in df["last_active_secs"]] + df["start_time"] = df["last_updated"] - pd.to_timedelta(df["click_count"] * 4, unit="s") + df["end_time"] = df["last_updated"] + return df + + +def _mock_metrics() -> pd.DataFrame: + df = _mock_snapshot() + online = df[df["status"] == "online"] + freshest = float(online["last_active_secs"].min()) if not online.empty else None + return pd.DataFrame([{ + "active_users": int((df["status"] == "online").sum()), + "converting": int((df["funnel_stage"] == "converting").sum()), + "needs_help": int(df["needs_help"].sum()), + "freshest_secs": freshest, + "total_rows": int(len(df)), + }]) + + +def _mock_grid() -> pd.DataFrame: + df = _mock_snapshot() + cols = ["user_id", "current_surface", "funnel_stage", "engagement_score", + "last_active_secs", "status", "needs_help"] + return ( + df[df["status"] == "online"][cols] + .sort_values("engagement_score", ascending=False) + .reset_index(drop=True) + ) + + +def _mock_detail(user_id: str) -> pd.DataFrame: + df = _mock_snapshot() + return df[df["user_id"] == user_id].reset_index(drop=True) + + +# -------------------------------------------------------------------------- +# Domain logic +# -------------------------------------------------------------------------- +def _acted() -> dict: + """user_id -> action kind taken this session ('support' | 'incentive'). + Client-side demo state, never written back to Lakebase.""" + return st.session_state.setdefault("acted", {}) + + +def _available_action(r): + """The operator action offered for this session's next-best-action, or None + when the NBA is automated (no human button). Returns (kind, button_label, + taken_label).""" + if r["needs_help"]: + return ("support", "Route to support", "Support in progress, live agent connected") + if r["funnel_stage"] == "converting": + return ("incentive", "Send assist / incentive", "Assist offer sent") + return None + + +def _support_acted_ids() -> set: + return {uid for uid, kind in _acted().items() if kind == "support"} + + +def next_best_action(r) -> str: + """The personalization decision computed from this user's live session state.""" + act = _available_action(r) + if act and _acted().get(r["user_id"]) == act[0]: + return act[2] # the 'action taken' state + if r["needs_help"]: + return "Offer live support" + if r["funnel_stage"] == "converting": + return "Nudge to complete (assist or incentive)" + if r["funnel_stage"] == "engaged": + return f"Surface tailored content for '{r['current_surface']}'" + return "Highlight popular feature / onboarding" + + +def nba_why(r) -> str: + done = _acted().get(r["user_id"]) + if done == "support": + return "you routed live support to this session" + if done == "incentive": + return "you sent an assist / incentive to this session" + if r["needs_help"]: + return "needs_help flag is set on this session" + if r["funnel_stage"] == "converting": + return "user is in the converting funnel stage, close to a decision" + if r["funnel_stage"] == "engaged": + return f"user is engaged and actively on the '{r['current_surface']}' surface" + return "user is browsing, not yet engaged" + + +def stage_label(r) -> str: + """The label shown for a session - needs-help outranks the funnel stage. A + routed (support-acted) session is no longer 'needs help' - it's being handled.""" + if r["needs_help"] and r["user_id"] not in _support_acted_ids(): + return "needs help" + return r["funnel_stage"] + + +# -------------------------------------------------------------------------- +# Presentation helpers +# -------------------------------------------------------------------------- +def _inject_css(): + st.markdown( + f""" + + """, + unsafe_allow_html=True, + ) + + +def _pill(label: str) -> str: + bg, fg = STAGE_COLORS.get(label, STAGE_COLORS["browsing"]) + return f'{label}' + + +def _metric_card(col, label, value, accent=None, foot=None): + color = f"color:{accent};" if accent else "" + foot_html = f'
{foot}
' if foot else "" + col.markdown( + f'
{label}
' + f'
{value}
{foot_html}
', + unsafe_allow_html=True, + ) + + +def _freshness_badge(freshest): + if freshest is None: + return "#80868b", "no data" + f = float(freshest) + color = "#34a853" if f < 2 else ("#f9ab00" if f < 5 else "#d93025") + return color, f"live · {f:.1f}s" + + +# -------------------------------------------------------------------------- +# UI - rendered inside an auto-refreshing fragment so the whole page does not +# flash on every tick. +# -------------------------------------------------------------------------- +def _picked_user(resp): + """Pull the selected user id out of an AgGrid selection response. AgGrid + returns the rendered ROW DATA, so the key is the displayed column name + "user" (we rename user_id -> user for the grid), not "user_id". Handles both + the list-of-dicts and DataFrame selection shapes across st_aggrid versions.""" + picked = getattr(resp, "selected_rows", None) + if picked is None: + try: + picked = resp["selected_rows"] + except Exception: + picked = None + if isinstance(picked, pd.DataFrame): + return picked.iloc[0].get("user") if not picked.empty else None + if picked: + return picked[0].get("user") + return None + + +SEGMENTS = ["All", "Converting", "Needs help", "Engaged"] + + +def _apply_segment(grid: pd.DataFrame, seg: str) -> pd.DataFrame: + if grid.empty or seg == "All": + return grid + if seg == "Converting": + return grid[grid["funnel_stage"] == "converting"] + if seg == "Needs help": + return grid[grid["needs_help"] & ~grid["user_id"].isin(_support_acted_ids())] + if seg == "Engaged": + return grid[grid["funnel_stage"] == "engaged"] + return grid + + +def _detect_theme() -> str: + """Active (toggle-aware) Streamlit theme - 'light' or 'dark', else 'base'. + Folded into the AgGrid key so the grid remounts and re-themes on a toggle.""" + try: + return st.context.theme.type or "base" + except Exception: + return "base" + + +try: + _BANNER_TEMPLATE = (Path(__file__).parent / "banner.html").read_text(encoding="utf-8") +except Exception: + _BANNER_TEMPLATE = "" + + +def _render_banner(theme: str, active: int): + """Architecture banner at the top of the console. Theme-aware (dark/light), and + the data-rail flow speed tracks the live active-session count - a busier stream + visibly flows faster (0.6s when busy, up to 1.8s when quiet).""" + if not _BANNER_TEMPLATE: + return + scheme = "light" if theme == "light" else "dark" + flow_secs = max(0.6, min(1.8, 1.8 - active / 150.0)) + st.markdown( + _BANNER_TEMPLATE.replace("__SCHEME__", scheme).replace("__FLOWDUR__", f"{flow_secs:.2f}s"), + unsafe_allow_html=True, + ) + + +def _render_online_grid(display: pd.DataFrame, theme: str = "base"): + """AgGrid table - click anywhere on a row (no checkbox) to drill the detail + in. AgGrid returns the selected ROW DATA, so we read the user_id directly + and there is no row-index drift on refresh. Writes session_state["selected_user"].""" + from st_aggrid import AgGrid, GridOptionsBuilder, JsCode + + gb = GridOptionsBuilder.from_dataframe(display) + gb.configure_selection(selection_mode="single", use_checkbox=False) + gb.configure_grid_options(rowHeight=34, headerHeight=36, suppressCellFocus=True) + gb.configure_column("last active", type=["numericColumn"], + valueFormatter=JsCode("function(p){return (p.value).toFixed(1)+' s';}")) + # Stage cell tinted by meaning (mirrors STAGE_COLORS). + gb.configure_column("stage", cellStyle=JsCode(""" + function(p){ + var m={'converting':['rgba(52,168,83,0.18)','#1e8e3e'], + 'engaged':['rgba(26,115,232,0.18)','#1a73e8'], + 'browsing':['rgba(140,140,150,0.18)','#9aa0a6'], + 'needs help':['rgba(234,67,53,0.20)','#e06055']}; + var c=m[p.value]||['transparent','inherit']; + return {backgroundColor:c[0],color:c[1],fontWeight:600}; + }""")) + # Engagement rendered as an inline progress bar. + gb.configure_column("engagement", cellRenderer=JsCode(""" + class B{init(p){var v=Math.max(0,Math.min(100,p.value||0)); + var e=document.createElement('div'); + e.style.cssText='display:flex;align-items:center;gap:6px;height:100%'; + e.innerHTML='
' + +'
' + +''+v+''; + this.e=e;}getGui(){return this.e;}}""")) + + # Folding the detected theme into the key forces AgGrid to remount when the + # theme flips (it won't re-theme in place). Relies on theme detection working. + resp = AgGrid( + display, + gridOptions=gb.build(), + update_on=["selectionChanged"], + allow_unsafe_jscode=True, + fit_columns_on_grid_load=True, + height=470, + theme="streamlit", + key=f"online_aggrid_{theme}", + ) + + uid = _picked_user(resp) + if uid: + st.session_state["selected_user"] = uid + + +def _take_action(uid: str, kind: str, label: str): + """Act on the inspected session (route support / send incentive). Client-side + demo only - records the action, logs it, toasts. Nothing written to Lakebase.""" + _acted()[uid] = kind + log = st.session_state.setdefault("action_log", []) + log.insert(0, (datetime.now().strftime("%H:%M:%S"), f"{label} · {uid}")) + st.toast(f"{label} · {uid}") + st.rerun(scope="fragment") + + +def _render_action_log(): + log = st.session_state.get("action_log", []) + if not log: + return + rows = "".join( + f'
{t}   {msg}
' for t, msg in log[:8] + ) + st.markdown(f'
Action log
{rows}
', + unsafe_allow_html=True) + + +@st.fragment(run_every=REFRESH_SECONDS) +def render_console(): + try: + metrics = load_metrics() + except Exception as e: + st.error("Lakebase connection not ready yet.") + st.exception(e) + return + + if metrics.empty or int(metrics.iloc[0]["total_rows"] or 0) == 0: + st.info("Waiting for the stream... no session rows yet.") + return + + m = metrics.iloc[0] + freshest = m["freshest_secs"] + dot_color, live_text = _freshness_badge(freshest) + role = st.session_state.get("pg_role", os.getenv("PGUSER", "local")) + _theme = _detect_theme() # drives the banner palette + the AgGrid remount-on-toggle + + # The live online population is the single source for the banner flow rate, the + # KPIs, the table, and the detail - so every number on screen agrees. + grid = load_grid() + converting_count = int((grid["funnel_stage"] == "converting").sum()) if not grid.empty else 0 + needs_help_count = ( + int((grid["needs_help"] & ~grid["user_id"].isin(_support_acted_ids())).sum()) + if not grid.empty else 0 + ) + + # Architecture banner - the data-rail flow speed tracks live active sessions. + _render_banner(_theme, len(grid)) + + # Hero + st.markdown( + f""" +
+
+

Personalization Console

+
Live session state, served sub-second straight from Lakebase
+
+
+ {live_text} +
+
+ """, + unsafe_allow_html=True, + ) + + # Metric strip - each card carries a one-line read of what the number means. + c1, c2, c3, c4 = st.columns(4) + _metric_card(c1, "Active users", len(grid), + foot="live concurrency right now") + _metric_card(c2, "Converting", converting_count, + accent=STAGE_COLORS["converting"][1], foot="revenue-proximal, protect these") + _metric_card(c3, "Needs help", needs_help_count, + accent=STAGE_COLORS["needs help"][1] if needs_help_count else STAGE_COLORS["converting"][1], + foot="friction happening now") + fresh_str = f"{float(freshest):.1f}s" if freshest is not None else "n/a" + _metric_card(c4, "Freshest update", fresh_str, accent=dot_color, + foot="serving latency vs Postgres now()") + + st.write("") + + # Segment filter - focus the live sessions table on a cohort. + seg = st.segmented_control( + "Focus the sessions", options=SEGMENTS, default="All", key="seg_filter", + ) or "All" + + filtered = _apply_segment(grid, seg) + left, right = st.columns([3, 2], gap="large") + + with left: + st.markdown(f"##### Live online sessions · {seg} ({len(filtered)})") + if filtered.empty: + st.info("No sessions in this segment right now.") + else: + display = filtered.copy() + display["stage"] = display.apply(stage_label, axis=1) + display["last active"] = display["last_active_secs"].astype(float).round(1) + display = display[ + ["user_id", "current_surface", "stage", "engagement_score", "last active"] + ].rename(columns={ + "user_id": "user", "current_surface": "surface", + "engagement_score": "engagement", + }) + _render_online_grid(display, theme=_theme) + + # Resolve the drilled-in user: an explicit pick (grid row or queue button) + # wins, otherwise fall back to the most engaged session in view. + selected_user = st.session_state.get("selected_user") + if not selected_user: + if not filtered.empty: + selected_user = filtered.iloc[0]["user_id"] + elif not grid.empty: + selected_user = grid.iloc[0]["user_id"] + + with right: + st.markdown("##### Session detail") + if selected_user is None: + st.caption("Select a session to see its next best action.") + else: + detail = load_detail(selected_user) + if detail.empty: + st.caption("That session is no longer present.") + else: + r = detail.iloc[0] + act = _available_action(r) + done = bool(act and _acted().get(r["user_id"]) == act[0]) + resolved = " resolved" if done else "" + st.markdown(f"**{r['user_id']}**   {_pill(stage_label(r))}", + unsafe_allow_html=True) + st.markdown( + f'
Next best action
' + f'
{next_best_action(r)}
' + f'
Why: {nba_why(r)}
' + f'
Decision logic is illustrative - plug in your ' + f'own rules or a served model here.
', + unsafe_allow_html=True, + ) + # The action for the inspected session: an NBA that warrants a human move (route + # support / send incentive) gets a button. support = red, incentive = green. + if act and not done: + kind, label, _ = act + with st.container(key=f"actbtn_{kind}"): + if st.button(label, key=f"act_{r['user_id']}"): + _take_action(r["user_id"], kind, label) + score = int(r["engagement_score"]) + st.markdown( + f'
Current surface' + f'{_pill(r["current_surface"])}
' + f'
Funnel stage' + f'{_pill(r["funnel_stage"])}
' + f'
Status' + f'{r["status"]}
' + f'
Click count' + f'{int(r["click_count"])}
' + f'
Last updated' + f'{str(r["last_updated"])[:19]}
', + unsafe_allow_html=True, + ) + st.markdown( + f'
Engagement score ' + f'  {score}/100
' + f'
', + unsafe_allow_html=True, + ) + st.caption("Read live from Lakebase Postgres - no cache, no batch.") + + _render_action_log() + + st.caption( + f"Auto-refreshing every {REFRESH_SECONDS}s. Connected as Postgres role " + f"`{role}`{' (mock data)' if MOCK else ' via a minted short-lived credential'}." + ) + + +def _startup_selftest(): + """One-time connectivity self-test, logged to stdout so the Lakebase read path is visible in the app logs.""" + if MOCK or st.session_state.get("_selftest_done"): + return + st.session_state["_selftest_done"] = True + try: + df = run_query( + "SELECT COUNT(*) AS n, " + "EXTRACT(EPOCH FROM (now() - MAX(last_updated))) AS fresh " + f"FROM {SCHEMA}.{TABLE}" + ) + n = int(df.iloc[0]["n"]) + fresh = df.iloc[0]["fresh"] + print( + f"[SELFTEST] Lakebase OK as role={st.session_state.get('pg_role')} " + f"rows={n} freshest_secs={float(fresh):.2f}", + flush=True, + ) + except Exception as e: # noqa: BLE001 + print(f"[SELFTEST] Lakebase FAILED: {type(e).__name__}: {e}", flush=True) + + +_inject_css() +_startup_selftest() +render_console() diff --git a/product_demos/realtime-clickstream-serving/app/app.yaml b/product_demos/realtime-clickstream-serving/app/app.yaml new file mode 100644 index 00000000..0bec9013 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/app/app.yaml @@ -0,0 +1,4 @@ +command: + - sh + - -c + - python preflight.py; streamlit run app.py diff --git a/product_demos/realtime-clickstream-serving/app/banner.html b/product_demos/realtime-clickstream-serving/app/banner.html new file mode 100644 index 00000000..75ab27c1 --- /dev/null +++ b/product_demos/realtime-clickstream-serving/app/banner.html @@ -0,0 +1,31 @@ +
Real-time personalization on live session stateLIVE · SUB-SECOND
Clickstream
Live events: surface + action
Real-Time Mode
transformWithState sessionizes into live per-user state
Lakebase
Continuous Upsert to live.sessions, one row per user
Personalization Console
Point lookup by user, serves the next best action
Live session state, read straight from Lakebase in milliseconds — no batch, no cache.
\ No newline at end of file diff --git a/product_demos/realtime-clickstream-serving/app/preflight.py b/product_demos/realtime-clickstream-serving/app/preflight.py new file mode 100644 index 00000000..d385e0dd --- /dev/null +++ b/product_demos/realtime-clickstream-serving/app/preflight.py @@ -0,0 +1,52 @@ +"""Startup pre-flight: prove the app service principal can connect to Lakebase +and read live.sessions, BEFORE Streamlit launches. Runs in the real app runtime +as the app SP. Logs a single [PREFLIGHT] line to stdout (visible in app logs). +Never fails the boot - Streamlit starts regardless so the UI can show its own +graceful 'connection not ready' state if something is off. +""" +import sys +import uuid + + +def main() -> None: + try: + import psycopg + from psycopg.rows import dict_row + from databricks.sdk import WorkspaceClient + + w = WorkspaceClient() + inst = w.database.get_database_instance(name="clickstream-sessions") + cred = w.database.generate_database_credential( + request_id=str(uuid.uuid4()), instance_names=["clickstream-sessions"] + ) + import os + + user = os.getenv("PGUSER") or w.current_user.me().user_name + conn = psycopg.connect( + f"host={inst.read_write_dns} dbname=databricks_postgres " + f"user={user} password={cred.token} sslmode=require", + row_factory=dict_row, + connect_timeout=10, + ) + with conn.cursor() as cur: + cur.execute( + "SELECT COUNT(*) AS n, " + "COUNT(*) FILTER (WHERE status='online') AS online, " + "EXTRACT(EPOCH FROM (now() - MAX(last_updated))) AS fresh " + "FROM live.sessions" + ) + row = cur.fetchone() + conn.close() + print( + f"[PREFLIGHT] Lakebase OK as role={user} " + f"total_rows={row['n']} online={row['online']} " + f"freshest_secs={float(row['fresh']):.2f}", + flush=True, + ) + except Exception as e: # noqa: BLE001 + print(f"[PREFLIGHT] Lakebase FAILED: {type(e).__name__}: {e}", flush=True) + + +if __name__ == "__main__": + main() + sys.exit(0) diff --git a/product_demos/realtime-clickstream-serving/app/requirements.txt b/product_demos/realtime-clickstream-serving/app/requirements.txt new file mode 100644 index 00000000..b1c152ef --- /dev/null +++ b/product_demos/realtime-clickstream-serving/app/requirements.txt @@ -0,0 +1,6 @@ +streamlit>=1.40,<2 +streamlit-aggrid +psycopg[binary] +databricks-sdk>=0.40 +protobuf==5.29.5 +pandas