Fix Lakebase app schema ownership

Author: SatyKrish

.github/workflows/deploy.yml

@@ -137,14 +137,24 @@ jobs:
         run: |
           databricks apps get doc-intel-analyst-demo --output json > /tmp/app.json
           app_obo_required="$(python -c "import yaml; d=yaml.safe_load(open('databricks.yml')); default=d.get('variables',{}).get('app_obo_required',{}).get('default','true'); value=d.get('targets',{}).get('demo',{}).get('variables',{}).get('app_obo_required', default); print(str(value).lower())")"
+          lakebase_name="$(python -c "import yaml; d=yaml.safe_load(open('databricks.yml')); print(d.get('targets',{}).get('demo',{}).get('variables',{}).get('lakebase_instance','docintel-demo-state-v1'))")"
+          python -c "import json; app=json.load(open('/tmp/app.json')); vals=[str(app.get(k)) for k in ('service_principal_client_id','service_principal_name','service_principal_id') if app.get(k) is not None]; print('\n'.join(dict.fromkeys(v for v in vals if v)))" > /tmp/app-sp-candidates.txt
+          db_granted=0
+          while IFS= read -r principal; do
+            grant_json="$(python -c "import json, sys; print(json.dumps({'access_control_list':[{'service_principal_name':sys.argv[1],'permission_level':'CAN_USE'}]}))" "$principal")"
+            if databricks permissions update database-instances "$lakebase_name" --json "$grant_json"; then
+              db_granted=1
+              break
+            fi
+          done < /tmp/app-sp-candidates.txt
+          test "$db_granted" = "1"
           if [ "$app_obo_required" = "true" ]; then
             # `bundle run` may wipe user_api_scopes (documented destructive-update
             # behavior). Fail loudly if required user scopes are missing.
             python -c "import json; app=json.load(open('/tmp/app.json')); scopes=set(app.get('user_api_scopes') or []); required={'serving.serving-endpoints','sql'}; missing=required-scopes; assert not missing, f'OBO scopes missing: {sorted(missing)} (got {sorted(scopes)})'"
           else
             python -c "import json; app=json.load(open('/tmp/app.json')); scopes=app.get('user_api_scopes'); assert not scopes, f'demo App-SP mode expected no user_api_scopes, got {scopes}'"
             endpoint_id="$(databricks serving-endpoints get "$AGENT_ENDPOINT_NAME" --output json | python -c "import json, sys; e=json.load(sys.stdin); print(e.get('id') or e.get('name'))")"
-            python -c "import json; app=json.load(open('/tmp/app.json')); vals=[str(app.get(k)) for k in ('service_principal_client_id','service_principal_name','service_principal_id') if app.get(k) is not None]; print('\n'.join(dict.fromkeys(v for v in vals if v)))" > /tmp/app-sp-candidates.txt
             granted=0
             while IFS= read -r principal; do
               grant_json="$(python -c "import json, sys; print(json.dumps({'access_control_list':[{'service_principal_name':sys.argv[1],'permission_level':'CAN_QUERY'}]}))" "$principal")"

app/README.md

@@ -35,6 +35,7 @@ export DATABRICKS_CLIENT_SECRET=<app-sp-secret>
 # Locally, set the same connection fields and let lakebase_client.py mint the
 # OAuth database password through the Databricks SDK.
 export DOCINTEL_LAKEBASE_INSTANCE=docintel-demo-state-v1
+export DOCINTEL_LAKEBASE_SCHEMA=docintel_app
 export PGDATABASE=docintel-demo-state-v1
 export PGUSER=<app-sp-application-id>
 export PGPORT=5432
@@ -49,13 +50,11 @@ streamlit run app/app.py
 
 Local runs do not have the Databricks Apps `x-forwarded-access-token` header, so they cannot validate the Agent Bricks OBO path. Use a deployed OBO-enabled target for prod identity validation.
 
-If you accidentally run Lakebase schema initialization with user creds (`DATABRICKS_CLIENT_ID`/`SECRET` unset), `lakebase_client.init_schema()` logs a warning identifying the mismatch. The tables get created under your user account, not the App SP, and the deployed App will lose write access. Drop the user-owned tables and re-init under the App SP to recover:
+If you accidentally run Lakebase schema initialization with user creds (`DATABRICKS_CLIENT_ID`/`SECRET` unset), `lakebase_client.init_schema()` logs a warning identifying the mismatch. The schema gets created under your user account, not the App SP, and the deployed App will lose write access. Drop the user-owned schema and re-init under the App SP to recover:
 
 ```sql
 -- connected as the App SP via the local-dev env above
-DROP TABLE IF EXISTS feedback CASCADE;
-DROP TABLE IF EXISTS query_logs CASCADE;
-DROP TABLE IF EXISTS conversation_history CASCADE;
+DROP SCHEMA IF EXISTS docintel_app CASCADE;
 -- next streamlit run will re-init under the App SP
 ```
 

app/lakebase_client.py

@@ -25,37 +25,56 @@
 from typing import Iterator
 
 import psycopg
+from psycopg import sql
 from databricks.sdk import WorkspaceClient
 
 _log = logging.getLogger(__name__)
 
 
-_SCHEMA = """
-CREATE TABLE IF NOT EXISTS conversation_history (
+def _lakebase_schema() -> str:
+    return os.environ.get("DOCINTEL_LAKEBASE_SCHEMA", "docintel_app")
+
+
+def _table(name: str) -> sql.Identifier:
+    return sql.Identifier(_lakebase_schema(), name)
+
+
+def _schema_ddl() -> sql.Composed:
+    conversation_history = _table("conversation_history")
+    query_logs = _table("query_logs")
+    feedback = _table("feedback")
+    return sql.SQL(
+        """
+CREATE TABLE IF NOT EXISTS {conversation_history} (
   conversation_id UUID PRIMARY KEY,
   user_email TEXT NOT NULL,
   started_at TIMESTAMPTZ NOT NULL DEFAULT now(),
   last_turn_at TIMESTAMPTZ NOT NULL DEFAULT now()
 );
-CREATE TABLE IF NOT EXISTS query_logs (
+CREATE TABLE IF NOT EXISTS {query_logs} (
   turn_id UUID PRIMARY KEY,
-  conversation_id UUID REFERENCES conversation_history(conversation_id),
+  conversation_id UUID REFERENCES {conversation_history}(conversation_id),
   question TEXT NOT NULL,
   answer TEXT NOT NULL,
   citations JSONB NOT NULL,
   latency_ms INT NOT NULL,
   agent_path TEXT NOT NULL,
   created_at TIMESTAMPTZ NOT NULL DEFAULT now()
 );
-CREATE TABLE IF NOT EXISTS feedback (
+CREATE TABLE IF NOT EXISTS {feedback} (
   feedback_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-  turn_id UUID REFERENCES query_logs(turn_id),
+  turn_id UUID REFERENCES {query_logs}(turn_id),
   user_email TEXT NOT NULL,
   rating TEXT NOT NULL CHECK (rating IN ('up','down')),
   comment TEXT,
   created_at TIMESTAMPTZ NOT NULL DEFAULT now()
 );
 """
+    ).format(
+        conversation_history=conversation_history,
+        query_logs=query_logs,
+        feedback=feedback,
+    )
 
 
 @contextmanager
@@ -98,7 +117,7 @@ def _generate_lakebase_password() -> str:
 
 
 def init_schema() -> None:
-    """Idempotent CREATE TABLE IF NOT EXISTS. Logs the connected role so
+    """Idempotent CREATE SCHEMA/TABLE IF NOT EXISTS. Logs the connected role so
     deployed-vs-local identity divergence is debuggable from app logs.
     """
     with _conn() as c, c.cursor() as cur:
@@ -115,23 +134,30 @@ def init_schema() -> None:
             )
         else:
             _log.info("Lakebase init connected as %r", connected_user)
-        cur.execute(_SCHEMA)
+        cur.execute(
+            sql.SQL("CREATE SCHEMA IF NOT EXISTS {}").format(sql.Identifier(_lakebase_schema()))
+        )
+        cur.execute(_schema_ddl())
 
 
 def ensure_conversation(conversation_id: uuid.UUID, user_email: str) -> None:
     with _conn() as c, c.cursor() as cur:
         cur.execute(
-            "INSERT INTO conversation_history (conversation_id, user_email) VALUES (%s, %s) "
-            "ON CONFLICT (conversation_id) DO UPDATE SET last_turn_at = now()",
+            sql.SQL(
+                "INSERT INTO {table} (conversation_id, user_email) VALUES (%s, %s) "
+                "ON CONFLICT (conversation_id) DO UPDATE SET last_turn_at = now()"
+            ).format(table=_table("conversation_history")),
             (conversation_id, user_email),
         )
 
 
 def log_turn(*, turn_id: str, conversation_id: uuid.UUID, response: dict, question: str) -> None:
     with _conn() as c, c.cursor() as cur:
         cur.execute(
-            "INSERT INTO query_logs (turn_id, conversation_id, question, answer, citations, latency_ms, agent_path) "
-            "VALUES (%s, %s, %s, %s, %s::jsonb, %s, %s)",
+            sql.SQL(
+                "INSERT INTO {table} (turn_id, conversation_id, question, answer, citations, latency_ms, agent_path) "
+                "VALUES (%s, %s, %s, %s, %s::jsonb, %s, %s)"
+            ).format(table=_table("query_logs")),
             (
                 turn_id,
                 conversation_id,
@@ -147,6 +173,8 @@ def log_turn(*, turn_id: str, conversation_id: uuid.UUID, response: dict, questi
 def write_feedback(*, turn_id: str, user_email: str, rating: str, comment: str | None) -> None:
     with _conn() as c, c.cursor() as cur:
         cur.execute(
-            "INSERT INTO feedback (turn_id, user_email, rating, comment) VALUES (%s, %s, %s, %s)",
+            sql.SQL("INSERT INTO {table} (turn_id, user_email, rating, comment) VALUES (%s, %s, %s, %s)").format(
+                table=_table("feedback")
+            ),
             (turn_id, user_email, rating, comment),
         )

docs/runbook.md

@@ -84,10 +84,10 @@ Metric key names can vary across MLflow/databricks-agents versions. The eval run
 | Agent answers ignore user UC permissions in prod | OBO scopes wiped by `bundle run` (documented destructive-update behavior — see [Databricks Apps deploy docs](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/deploy)) | Re-apply scopes to the target app: `databricks apps update <app-name> --user-api-scopes serving.serving-endpoints,sql,iam.access-control:read,iam.current-user:read` |
 | Agent deployment cannot grant endpoint query permission | Permissions API was called with endpoint name instead of internal endpoint ID, or the generated endpoint is not ready | Use current `agent/document_intelligence_agent.py`; it waits for readiness and grants by serving endpoint ID |
 | Streamlit user sees stale UC permissions | OBO token captured at WebSocket open; never refreshes ([Databricks Apps runtime docs](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/app-runtime)) | Reload the page after permission changes |
-| Lakebase tables not writable from deployed App | Local-dev `streamlit run` initialised schema under user identity, not App SP | Connect as App SP and `DROP TABLE feedback, query_logs, conversation_history`; next App run re-creates them under SP. See `app/README.md` |
+| Lakebase tables not writable from deployed App | Local-dev `streamlit run` initialised the `docintel_app` schema under user identity, not App SP | Connect as App SP and `DROP SCHEMA docintel_app CASCADE`; next App run re-creates it under SP. See `app/README.md` |
 | CLEARS Latency axis fails | Agent Bricks orchestration or Knowledge Assistant source is too broad | Narrow the Knowledge Assistant source, tune Supervisor instructions, or reduce structured-tool fan-out |
 | Citation chips render but filenames show `source` | Knowledge Assistant footnote format changed or omitted filename markers | Capture the raw Agent Bricks payload and compare it with `app/agent_bricks_response.py`'s markdown-footnote parser |
-| App errors connecting to Lakebase | Database resource binding missing connection fields, or OAuth credential minting failed | Check the `docintel-lakebase` resource binding plus `PGHOST`/`PGPORT`/`PGUSER`/`PGDATABASE`/`DOCINTEL_LAKEBASE_INSTANCE` in the App runtime. `PGPASSWORD` is minted at connection time by `app/lakebase_client.py` |
+| App errors connecting to Lakebase | Database resource binding missing connection fields, OAuth credential minting failed, or App SP lacks Lakebase instance `CAN_USE` | Check the `docintel-lakebase` resource binding plus `PGHOST`/`PGPORT`/`PGUSER`/`PGDATABASE`/`DOCINTEL_LAKEBASE_INSTANCE`/`DOCINTEL_LAKEBASE_SCHEMA` in the App runtime. `PGPASSWORD` is minted at connection time by `app/lakebase_client.py` |
 
 ## Verifying end-to-end OBO
 

resources/consumers/analyst.app.yml

@@ -13,6 +13,8 @@ resources:
             value: ${var.app_obo_required}
           - name: DOCINTEL_LAKEBASE_INSTANCE
             value: ${var.lakebase_instance}
+          - name: DOCINTEL_LAKEBASE_SCHEMA
+            value: docintel_app
 
       # Databricks Apps auto-grants Lakebase permissions to the App SP on
       # deploy — see https://docs.databricks.com/aws/en/dev-tools/databricks-apps/access-data.

scripts/bootstrap-demo.sh

@@ -204,19 +204,9 @@ for i in d.get('database_instances', []):
   done
 }
 
-grant_app_sp_endpoint_query() {
+app_sp_principals() {
   local app_json="$1"
-  local endpoint_json endpoint_id principals principal grant_json
-  endpoint_json=$(databricks serving-endpoints get "$AGENT_ENDPOINT_NAME" --output json)
-  endpoint_id=$(printf '%s' "$endpoint_json" | "$PYTHON" -c "
-import json, sys
-endpoint = json.load(sys.stdin)
-print(endpoint.get('id') or endpoint.get('name') or '$AGENT_ENDPOINT_NAME')
-")
-  principals=()
-  while IFS= read -r principal; do
-    [[ -n "$principal" ]] && principals+=("$principal")
-  done < <(printf '%s' "$app_json" | "$PYTHON" -c "
+  printf '%s' "$app_json" | "$PYTHON" -c "
 import json, sys
 app = json.load(sys.stdin)
 seen = set()
@@ -228,7 +218,50 @@ for key in ('service_principal_client_id', 'service_principal_name', 'service_pr
     if value and value not in seen:
         seen.add(value)
         print(value)
+"
+}
+
+grant_app_sp_lakebase_use() {
+  local app_json="$1"
+  local principals principal grant_json
+  principals=()
+  while IFS= read -r principal; do
+    [[ -n "$principal" ]] && principals+=("$principal")
+  done < <(app_sp_principals "$app_json")
+  if (( ${#principals[@]} == 0 )); then
+    die "app service principal was not returned by Databricks Apps API"
+  fi
+  for principal in "${principals[@]}"; do
+    grant_json=$("$PYTHON" -c "
+import json, sys
+print(json.dumps({
+    'access_control_list': [{
+        'service_principal_name': sys.argv[1],
+        'permission_level': 'CAN_USE',
+    }]
+}))
+" "$principal")
+    if databricks permissions update database-instances "$LAKEBASE_NAME" --json "$grant_json" >/dev/null 2>&1; then
+      log "  granted CAN_USE on Lakebase $LAKEBASE_NAME to App SP $principal"
+      return 0
+    fi
+  done
+  die "failed to grant CAN_USE on Lakebase $LAKEBASE_NAME to the App service principal"
+}
+
+grant_app_sp_endpoint_query() {
+  local app_json="$1"
+  local endpoint_json endpoint_id principals principal grant_json
+  endpoint_json=$(databricks serving-endpoints get "$AGENT_ENDPOINT_NAME" --output json)
+  endpoint_id=$(printf '%s' "$endpoint_json" | "$PYTHON" -c "
+import json, sys
+endpoint = json.load(sys.stdin)
+print(endpoint.get('id') or endpoint.get('name') or '$AGENT_ENDPOINT_NAME')
 ")
+  principals=()
+  while IFS= read -r principal; do
+    [[ -n "$principal" ]] && principals+=("$principal")
+  done < <(app_sp_principals "$app_json")
   if (( ${#principals[@]} == 0 )); then
     die "app service principal was not returned by Databricks Apps API"
   fi
@@ -365,6 +398,7 @@ databricks api patch \
 
 log "  verifying app auth mode on $APP_NAME"
 if app_state=$(databricks apps get "$APP_NAME" --output json 2>/dev/null); then
+  grant_app_sp_lakebase_use "$app_state"
   if [[ "$APP_OBO_REQUIRED" == "true" ]]; then
     "$PYTHON" -c "
 import json

specs/001-doc-intel-10k/data-model.md

@@ -1,6 +1,6 @@
 # Phase 1 Data Model
 
-All Delta tables live under the bundle-parameterized `${var.catalog}.${var.schema}`. Lakebase tables live in the bundle-managed Lakebase database instance `${var.lakebase_instance}`, exposed to SQL dashboards through the UC database catalog `${var.schema}_state`.
+All Delta tables live under the bundle-parameterized `${var.catalog}.${var.schema}`. Lakebase tables live in schema `docintel_app` inside the bundle-managed Lakebase database instance `${var.lakebase_instance}`, exposed to SQL dashboards through the UC database catalog `${var.schema}_state`.
 
 ## Bronze
 

specs/001-doc-intel-10k/tasks.md

@@ -38,7 +38,7 @@ This is a DAB plus Agent Bricks deployment project. SQL pipeline code is at `pip
 **⚠️ CRITICAL**: All user stories depend on these.
 
 - [x] T006 Define UC catalog/schema/volume in `resources/foundation/catalog.yml`: `${var.catalog}.${var.schema}` schema + `raw_filings` volume; grant `USE_CATALOG`, `USE_SCHEMA`, `READ_VOLUME` to a configurable analyst group
-- [x] T007 [P] Define the Lakebase instance/catalog in `resources/foundation/lakebase_instance.yml` and `resources/consumers/lakebase_catalog.yml`; the UC catalog is `${var.schema}_state`, while `app/lakebase_client.py` creates `conversation_history`, `query_logs`, and `feedback` tables at runtime using App resource binding fields plus Databricks-minted Lakebase OAuth credentials
+- [x] T007 [P] Define the Lakebase instance/catalog in `resources/foundation/lakebase_instance.yml` and `resources/consumers/lakebase_catalog.yml`; the UC catalog is `${var.schema}_state`, while `app/lakebase_client.py` creates `docintel_app.conversation_history`, `docintel_app.query_logs`, and `docintel_app.feedback` at runtime using App resource binding fields plus Databricks-minted Lakebase OAuth credentials
 - [x] T008 [P] Add JSON contracts under `specs/001-doc-intel-10k/contracts/`: `agent-request.json`, `agent-response.json`, `feedback-event.json`, `kpi-schema.json`
 
 **Checkpoint**: catalog, schema, volume, Lakebase database exist; bundle validates.

src/dashboards/usage.lvdash.json

@@ -26,7 +26,7 @@
         "SELECT created_at::date AS day, agent_path, count(*) AS turns,",
         "  percentile_approx(latency_ms, 0.95) AS latency_p95_ms,",
         "  sum(CASE WHEN array_size(citations) = 0 THEN 1 ELSE 0 END) AS ungrounded",
-        "FROM `__dataset_schema___state`.public.query_logs",
+        "FROM `__dataset_schema___state`.docintel_app.query_logs",
         "GROUP BY 1, 2 ORDER BY 1 DESC"
       ]
     }