wehub/omnigent-ai--omnigent
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:12:00 +08:00
commit 3de48288cb
2986 changed files with 1131193 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deploy/databricks/README.md
+294
View File
@@ -0,0 +1,294 @@
# Deploying Omnigent on Databricks Apps
This directory deploys the Omnigent server to
[Databricks Apps](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/)
via [Databricks Asset Bundles](https://docs.databricks.com/aws/en/dev-tools/bundles/):
- **Lakebase** (managed PostgreSQL) — the database for every store.
- **UC Volumes** — the artifact store for agent bundles and executor
storage snapshots.
> **Most Databricks users want the managed offering instead.**
> [Omnigent on Databricks](https://docs.databricks.com/aws/en/omnigent/)
> (Beta) runs the server for you, wired to workspace identity,
> Foundation Models, AI Gateway, and MLflow Tracing out of the box.
> Enable the **Omnigent** preview in your workspace settings and follow
> the quickstart there. Use this directory only when you need to
> self-manage the deployment: the managed service is not in your region
> yet, or you need control it does not expose today (custom YAML
> policies, bring-your-own provider API keys, custom egress controls).
The orchestrator at `deploy.py` builds the wheels, generates an app
`pyproject.toml` + `uv.lock`, and then runs
`databricks bundle deploy` + `bundle run` against the bundle config
in `databricks.yml`. App config (Lakebase, UC volume) lives
declaratively in `databricks.yml` — adding or removing a resource is
a YAML edit, not a Python SDK call.
Runs unchanged from a laptop. Re-runnable; every step is idempotent.
## Prerequisites
1. A Databricks workspace with Databricks Apps, Lakebase, and UC
Volumes enabled.
2. The [Databricks CLI](https://docs.databricks.com/aws/en/dev-tools/cli/install.md)
installed and authenticated. Either a CLI profile
(`DATABRICKS_CONFIG_PROFILE=<profile>`) or env-based auth
(`DATABRICKS_HOST` + `DATABRICKS_CLIENT_ID` + `DATABRICKS_CLIENT_SECRET`).
3. The repo's local venv with the `databricks` extra:
`uv sync --extra databricks` (use `uv`, not global pip).
4. Permissions to create or use:
- a Lakebase project (one per app — do not share with other apps);
- a UC volume whose parent catalog/schema can grant access to the
app service principal;
- (optional) Databricks secrets for LLM API keys.
Set your workspace URL in `databricks.yml` under
`targets.prod.workspace.host` (it ships as a `https://example.databricks.com`
placeholder; DAB reads `workspace.host` before resolving variables, so it
must be a literal).
## One-time bootstrap
### 1. Lakebase project (one per app — never share)
Reusing a shared autoscaling project causes the migrate-on-boot hook
to fail with "permission denied for table agents" because the tables
are owned by whoever ran migrations first. Always start fresh:
```python
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.postgres import Project
wc = WorkspaceClient(profile="<your-profile>")
wc.postgres.create_project(project=Project(), project_id="omnigent")
branch = "projects/omnigent/branches/production"
endpoint = f"{branch}/endpoints/primary"
import time
for _ in range(120):
ep = wc.postgres.get_endpoint(name=endpoint)
if ep.status and ep.status.current_state == "ACTIVE":
break
time.sleep(5)
else:
raise TimeoutError(endpoint)
database = next(iter(wc.postgres.list_databases(parent=branch)))
print("database resource path:", database.name)
```
### 2. UC Volumes
```sql
CREATE SCHEMA IF NOT EXISTS main.omnigent;
CREATE VOLUME IF NOT EXISTS main.omnigent.artifacts;
```
The `artifacts` volume is referenced declaratively in `databricks.yml`
(the app resource) via `--var volume_name=…`.
### 3. First deploy — creates the app and its service principal
Run the [Deploy](#deploy) command once. The first
`databricks bundle deploy` creates the app and provisions its service
principal (SP). This first pass will **not** pass its `/health` check
yet: the SP has no Lakebase schema grants, so the migrate-on-boot hook
fails with `permission denied for schema public`. That's expected —
the next step grants those privileges.
### 4. Grant the app SP Lakebase privileges
Now that the app (and its SP) exist, grant the SP the public schema
privileges Alembic needs, then re-run the deploy:
```bash
python deploy/databricks/grant_sp_perms.py \
--app-name omnigent \
--lakebase-endpoint projects/omnigent/branches/production/endpoints/primary \
--database databricks_postgres \
--profile <your-profile>
```
> [!NOTE]
> Lakebase uses two spellings for the same database. The **resource
> path** uses a hyphenated slug — `…/databases/databricks-postgres`
> (what `deploy.py --lakebase-database` and `databricks.yml` want) —
> while the underlying **PostgreSQL database name** uses an underscore,
> `databricks_postgres` (what `grant_sp_perms.py --database` and the
> app's `PGDATABASE` use). Pass each form where shown.
After this one-time grant, re-running the deploy boots the app cleanly
and `/health` returns 200. Subsequent redeploys are a single
`deploy.py` invocation.
## Deploy
```bash
uv run python deploy/databricks/deploy.py \
--app-name omnigent \
--profile <your-profile> \
--lakebase-branch projects/omnigent/branches/production \
--lakebase-database projects/omnigent/branches/production/databases/databricks-postgres \
--volume-name main.omnigent.artifacts
```
The script builds wheels, classifies them by size, copies wheels into
`src/`, regenerates `src/pyproject.toml` and `src/uv.lock`, runs
`databricks bundle deploy --target prod`, runs
`databricks bundle run omnigent --target prod`, and polls `/health`
with backoff until 200.
All Omnigent wheels must fit under the Databricks Apps source
snapshot limit (10 MB). If a wheel exceeds it, rebuild with
`--skip-web-ui` or reduce the wheel size; uv lockfiles cannot point at
UC Volume wheel paths because `uv lock` validates path sources locally.
Re-running is safe — every step is idempotent.
> [!TIP]
> To lock against a private PyPI mirror or proxy instead of public
> PyPI, set `UV_INDEX_URL` before running `deploy.py`.
## Smoke check
`deploy.py` polls `/health` automatically. To check other endpoints:
```bash
TOKEN="$(databricks auth token <your-profile> --output json \
| python -c 'import json, sys; print(json.load(sys.stdin)["access_token"])')"
curl --http1.1 -fsS \
-H "Authorization: Bearer ${TOKEN}" \
https://<app>.databricksapps.com/health
```
## How it works
### Authentication
The app runs as a Databricks service principal. Credentials are
managed automatically:
- **Lakebase** — OAuth tokens generated via
`WorkspaceClient.postgres.generate_database_credential()`, injected
into every new SQLAlchemy connection via a class-level
`do_connect` event hook in `src/app.py`.
- **UC Volumes** — Workspace credentials used by the Databricks SDK
(ambient in Apps).
- **TUI / API access** — Browser-based OAuth using the
`databricks-cli` OIDC client with PKCE.
The Databricks Apps proxy injects `X-Forwarded-Email` on every
request, so the app pins `OMNIGENT_AUTH_PROVIDER=header` (see
`src/app.py`).
> [!IMPORTANT]
> Header auth trusts the `X-Forwarded-Email` header verbatim. This is
> safe **only** because the Databricks Apps platform terminates auth at
> its proxy, strips any client-supplied copy of the header, and the app
> port is never reachable except through that proxy. Don't expose the
> app process directly (e.g. a port forward or alternate ingress that
> bypasses the proxy): a caller who can set the header themselves could
> then impersonate any user. If you front the app with anything other
> than the standard Apps proxy, ensure it sanitizes the header too.
### Token lifecycle
Lakebase OAuth tokens expire after 60 minutes. The SQLAlchemy
connection pool recycles connections every 5 minutes by default
(configurable via `AP_POOL_RECYCLE_SECONDS`), ensuring fresh tokens
on new connections.
### Storage
| Component | Backend | Purpose |
|---|---|---|
| Agent specs, tasks, conversations | Lakebase (PostgreSQL) | Durable metadata |
| Agent bundles, executor snapshots | UC Volumes | Binary blob storage |
| DBOS workflow state | Lakebase (same DB) | Workflow recovery |
| Executor working dirs | Local ephemeral disk | Cache (restored from UC Volumes) |
## Configuration reference
Environment variables read by `src/app.py`:
| Variable | Source | Description |
|---|---|---|
| `PGHOST` | Databricks runtime | Lakebase hostname |
| `PGPORT` | Databricks runtime | Lakebase port (default 5432) |
| `PGDATABASE` | Databricks runtime | Lakebase database name |
| `PGUSER` | Databricks runtime | Lakebase user (service principal) |
| `PGSSLMODE` | Databricks runtime | SSL mode (default `require`) |
| `AP_LAKEBASE_ENDPOINT` | app resource `valueFrom: postgres` | Lakebase endpoint for token generation |
| `AP_ARTIFACT_VOLUME_PATH` | app resource `valueFrom: artifact_volume` | UC Volume path for artifacts |
| `DATABRICKS_APP_PORT` | Databricks runtime | App port (default 8000) |
| `AP_POOL_RECYCLE_SECONDS` | Optional | Connection pool recycle interval (default 300) |
## Multi-app safety — one bundle, many apps
The same bundle directory can deploy many apps (one per `--app-name`).
Terraform can only delete or replace what is tracked in the state it
loads, so the blast radius of a deploy is exactly that state file.
- **Remote state is per app.** `targets.<t>.workspace.root_path` ends
in `${var.app_name}`, so `--app-name X` reads and writes state only
under `.bundle/omnigent/X`. A deploy of X cannot mutate app Y.
- **The app resource's `name` is `${var.app_name}`.** If the loaded
state tracks app X but you pass `app_name=Y`, terraform sees a name
change and plans a **destroy of X + create of Y**. Never bind the
bundle resource to one app and then deploy with a different
`--app-name`.
- The **local** cache at `deploy/databricks/.databricks/bundle/<target>/`
is per-*target*. Before deploying a *different* app on a target
you've used before, drop it: `rm -rf deploy/databricks/.databricks/bundle/<target>`
(it's only a cache; the per-app remote state is the source of truth).
If a `bundle deploy` plan ever shows a delete or replace of a
`databricks_app`, abort and re-check the bind and `--app-name` —
routine redeploys only ever update in place.
## Common deploy modes
```bash
# Iterate without rebuilding wheels (reuses dist/; useful when you only
# changed app.py / app.yaml). Skips the clean-tree check.
uv run python deploy/databricks/deploy.py --skip-build --allow-dirty ...
# API-only deploy (drops the SPA from the main wheel).
uv run python deploy/databricks/deploy.py --skip-web-ui ...
```
## Troubleshooting
| Symptom | Cause | Fix |
|---|---|---|
| Deploy refuses: "working tree has uncommitted changes" / "HEAD is not at origin/main" | Clean-tree assertion | Commit/stash, `git checkout main && git pull`, or pass `--allow-dirty` |
| `bundle deploy` fails: "Resource already managed by Terraform" | App already bound to another bundle directory | Run from that directory, or unbind: `databricks bundle deployment unbind omnigent` |
| `bundle deploy` fails: "An app with the same name already exists" | App exists but isn't bound to this bundle (or a stale per-target local cache from a *different* app made `deploy.py` skip the bind) | `rm -rf deploy/databricks/.databricks/bundle/<target>`, then bind: `databricks bundle deployment bind omnigent <app-name> --target <target> --auto-approve --var ...` |
| App fails "Error installing packages"; `/logz` shows "Ignoring existing lockfile due to … exclude newer …" then a PyPI fetch timeout | The Apps runtime pins a global uv `exclude-newer` cutoff; a lock generated without the matching option is re-resolved in-container, where PyPI is unreachable | Read the cutoff from `/logz` ("change of exclude newer timestamp from X to Y") and redeploy with `UV_EXCLUDE_NEWER=<cutoff>` in the environment |
| `permission denied for table agents` | Lakebase tables owned by wrong user | Connect as the owner and `DROP TABLE … CASCADE`; redeploy |
| `schema "dbos" already exists` | Same for the DBOS schema | `DROP SCHEMA dbos CASCADE` and redeploy |
| `permission denied for schema public` | App SP missing schema grants | Run `grant_sp_perms.py` (one-time) |
| `Field 'spec.role' cannot be empty` | Lakebase requires explicit role for extra databases | Use the project's default database; don't create extras |
| Deploy refuses because a wheel is over 10 MB | uv app payload requires local wheel path sources | Rebuild with `--skip-web-ui` or reduce wheel size |
| App starts cleanly but the first agent request 403s on the artifact volume | App SP has `WRITE_VOLUME` on the leaf but no `USE_CATALOG` / `USE_SCHEMA` on the parents | `deploy.py` grants both automatically — for a fresh catalog, redeploy or grant manually via `databricks grants update` |
## Files in this directory
| File | Purpose |
|---|---|
| `deploy.py` | Orchestrator. Single entry point. |
| `databricks.yml` | DAB bundle config. Declares the app + its resources. |
| `build.sh` | Cleans static, builds the web UI, builds three wheels. |
| `grant_sp_perms.py` | One-time Lakebase `public` schema grant for the app SP. |
| `src/app.py` | The app process. SQLAlchemy `do_connect` token hook + Alembic-on-boot + uvicorn. |
| `src/app.yaml` | App startup config — command + env-var wiring. |
| `src/pyproject.toml` / `src/uv.lock` | Regenerated per deploy; not committed (they pin the per-deploy wheel version). |
## See also
- [`databricks.yml`](./databricks.yml) — DAB bundle config.
- [Databricks Apps docs](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/).
- [Databricks Asset Bundles docs](https://docs.databricks.com/aws/en/dev-tools/bundles/).
deploy/databricks/build.sh
Executable
+50
View File
@@ -0,0 +1,50 @@
#!/usr/bin/env bash
# Build the wheels and optional web UI assets needed for a Databricks Apps
# deployment of Omnigent.
#
# Inputs:
# SKIP_WEB_UI=1 Skip the web SPA build for API-only deployments.
#
# Outputs:
# dist/omnigent-<version>-py3-none-any.whl
# dist/omnigent_client-<version>-py3-none-any.whl
# dist/omnigent_ui_sdk-<version>-py3-none-any.whl
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# This file lives at deploy/databricks/ — two levels deep — so the repo root is
# two parents up.
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
cd "${REPO_ROOT}"
# Each Vite build emits uniquely-hashed JS chunk filenames. Without a
# sweep, orphaned chunks from prior builds accumulate in the static
# dir, end up in the main wheel, and push it over the 10 MB Workspace
# upload cap. Always start from a clean slate.
echo "==> Cleaning stale static assets and build outputs"
rm -rf omnigent/server/static/web-ui dist build omnigent.egg-info
if [[ "${SKIP_WEB_UI:-}" != "1" ]]; then
echo "==> Building web SPA into omnigent/server/static/web-ui/"
cd web
npm install
npm run build
cd "${REPO_ROOT}"
else
echo "==> SKIP_WEB_UI=1: skipping web build"
fi
echo "==> Building omnigent-client wheel"
uv build --wheel --out-dir dist/ sdks/python-client/
echo "==> Building omnigent-ui-sdk wheel"
uv build --wheel --out-dir dist/ sdks/ui/
echo "==> Building omnigent wheel"
uv build --wheel --out-dir dist/ .
echo ""
echo "Built wheels:"
ls -1 dist/
deploy/databricks/databricks.yml
+84
View File
@@ -0,0 +1,84 @@
# Databricks Asset Bundle config for the Omnigent Databricks App.
#
# Every knob the app exposes lives here. deploy.py wraps `databricks
# bundle deploy` + `bundle run`; per-app values are passed as `--var`
# pairs (see deploy.py / README.md).
bundle:
name: omnigent
variables:
app_name:
description: "Databricks App name (e.g. omnigent)."
lakebase_branch:
description: "projects/<app>/branches/production"
lakebase_database:
description: "projects/<app>/branches/production/databases/databricks-postgres"
volume_name:
description: "<catalog>.<schema>.<volume> for artifact storage"
otel_table_schema:
description: >
UC schema (catalog.schema) holding the OTel destination tables.
The platform writes to <schema>.otel_logs, otel_metrics, otel_spans.
default: main.omnigent_logs
resources:
apps:
# Resource key must match _BUNDLE_RESOURCE_KEY in deploy.py. The app's
# display name comes from ${var.app_name}.
omnigent:
name: ${var.app_name}
description: "Omnigent server backed by Lakebase + UC Volumes."
source_code_path: ./src
# compute_size is intentionally NOT declared here. Terraform's
# databricks_app provider uses an apps update API that rejects
# compute_size changes ("not supported in this update API"). We
# set it out-of-band via apps.create_update in deploy.py before
# the bundle deploy runs.
config:
command: ["opentelemetry-instrument", "python", "app.py"]
env:
- name: AP_LAKEBASE_ENDPOINT
value_from: postgres
- name: AP_ARTIFACT_VOLUME_PATH
value_from: artifact_volume
- name: OTEL_TRACES_SAMPLER
value: 'always_on'
resources:
- name: postgres
postgres:
branch: ${var.lakebase_branch}
database: ${var.lakebase_database}
permission: CAN_CONNECT_AND_CREATE
- name: artifact_volume
uc_securable:
securable_full_name: ${var.volume_name}
securable_type: VOLUME
permission: WRITE_VOLUME
targets:
# One target per workspace. workspace.host must be literal (DAB reads it
# before resolving vars). deploy.py selects via --target. To deploy to a
# different workspace, add another target block here and pass
# `deploy.py --target <name>`.
prod:
mode: production
workspace:
host: https://example.databricks.com # ← your workspace URL
# Isolate state per app so one bundle can deploy multiple apps.
root_path: /Workspace/Users/${workspace.current_user.userName}/.bundle/${bundle.name}/${var.app_name}
resources:
apps:
omnigent:
# Optional: Databricks Apps platform OTel export. The platform
# auto-injects OTEL_EXPORTER_OTLP_ENDPOINT into the app container
# so any OTLP emitter (including telemetry.init() in src/app.py)
# routes to the configured UC tables. The export API rejects tables
# in default storage ("Cannot export App telemetry to tables in
# default storage"), so this requires a workspace with a custom
# storage location — drop this block if yours has none.
telemetry_export_destinations:
- unity_catalog:
logs_table: ${var.otel_table_schema}.otel_logs
metrics_table: ${var.otel_table_schema}.otel_metrics
traces_table: ${var.otel_table_schema}.otel_spans
deploy/databricks/deploy.py
+933
View File
@@ -0,0 +1,933 @@
#!/usr/bin/env python3
"""Deploy Omnigent to a Databricks App via Databricks Asset Bundles.
End-to-end orchestrator that wraps `databricks bundle deploy` +
`databricks bundle run`. The build pieces (version stamp, wheel
build, uv lock generation, stale-wheel sweep) stay in Python; the
deploy itself is a DAB so the app's resource definition (Lakebase,
UC volume) lives declaratively in ``databricks.yml``.
Runs unchanged from a laptop or from CI. Re-runnable;
every step is idempotent.
Usage example:
python deploy/databricks/deploy.py \\
--app-name omnigent --profile <your-profile> \\
--lakebase-branch projects/omnigent/branches/production \\
--lakebase-database \\
projects/omnigent/branches/production/databases/databricks-postgres \\
--volume-name main.omnigent.artifacts
See ``README.md`` in the same directory for the full guide,
including first-time infrastructure setup.
"""
from __future__ import annotations
import argparse
import os
import re
import shutil
import subprocess
import sys
import time
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from databricks.sdk import WorkspaceClient
_WORKSPACE_WHEEL_LIMIT_BYTES = 10 * 1024 * 1024
_APP_REQUIRES_PYTHON = ">=3.12,<3.13"
# Public PyPI by default. Set UV_INDEX_URL to lock against a private mirror or
# proxy instead (see run_uv_lock).
_UV_DEFAULT_INDEX_URL = "https://pypi.org/simple"
# Leaving these in the env when we hand off to the CLI/SDK can
# silently route us to the wrong workspace, or upload code under the
# wrong account. The deploy must use --profile / DATABRICKS_HOST +
# DATABRICKS_CLIENT_ID explicitly.
_ENV_VARS_TO_CLEAR = (
"DATABRICKS_TOKEN",
"ANTHROPIC_API_KEY",
"CODEX",
"CLAUDE_CODE",
)
# Must match the `resources.apps.<key>` and `bundle.name` in databricks.yml.
_BUNDLE_RESOURCE_KEY = "omnigent"
_WHEEL_PREFIXES = ("omnigent-", "omnigent_client-", "omnigent_ui_sdk-")
def _log(msg: str) -> None:
print(f"[deploy] {msg}", flush=True)
def _repo_root() -> Path:
# deploy/databricks/deploy.py → repo root is two parents up.
return Path(__file__).resolve().parents[2]
def _deploy_dir() -> Path:
return Path(__file__).resolve().parent
def _src_dir() -> Path:
return _deploy_dir() / "src"
def _pyproject_paths() -> list[Path]:
root = _repo_root()
return [
root / "pyproject.toml",
root / "sdks" / "python-client" / "pyproject.toml",
root / "sdks" / "ui" / "pyproject.toml",
]
def _read_base_version() -> str:
"""Read the base version from the top-level pyproject.toml.
The three pyprojects share the same version; we only need to read
one. The base value is the on-disk version; deploys append a
`.postN` suffix so pip's wheel cache treats every deploy as a
distinct release.
"""
text = (_repo_root() / "pyproject.toml").read_text()
match = re.search(r'^version\s*=\s*"([^"]+)"', text, re.MULTILINE)
if not match:
raise RuntimeError("could not find version in pyproject.toml")
return match.group(1)
def _compute_deploy_version(base: str, explicit: str | None) -> str:
if explicit:
# Caller knows what they want — let it through after a sanity check.
if not re.match(r"^\d+(\.\d+)*(\.dev\d+|\.post\d+|[+\-][\w.]+)?$", explicit):
raise SystemExit(f"--version {explicit!r} is not a recognizable PEP 440 version")
return explicit
# Strip any existing `.postN` / `.devN` suffix so we don't stack
# them if a prior deploy left pyproject.toml dirty (or someone
# committed the bumped value). Without this, `0.1.0.post<old>`
# would become `0.1.0.post<old>.post<new>` which isn't valid
# PEP 440 and fails the wheel build.
base = re.sub(r"(\.post\d+|\.dev\d+)+$", "", base)
# Post-release, not dev: pip treats `.dev` as a pre-release and
# ignores it when resolving `>=` constraints, so a deploy that
# bumps via `.dev` clashes with `omnigent-ui-sdk` declaring
# `omnigent-client>=0.1.0`. `.post` is a final release and
# sorts strictly above the base.
return f"{base}.post{int(time.time())}"
def set_version_in_pyproject(path: Path, new_version: str) -> str:
"""Rewrite the version line and lockstep sibling pins in a pyproject.
The release process pins the sibling SDK packages in lockstep
(e.g. ``"omnigent-client==0.1.0rc2"``). Stamping only the
``version = "..."`` line would leave those exact pins pointing at
the unstamped base version, which no built wheel carries — so the
app-level ``uv lock`` becomes unsatisfiable. Rewrite both.
:param path: Pyproject file to rewrite, e.g.
``<repo>/sdks/ui/pyproject.toml``.
:param new_version: Stamped deploy version, e.g. ``"0.1.0rc2.post123"``.
:returns: The original file text, for restore after the build.
"""
original = path.read_text()
updated, count = re.subn(
r'(?m)^version\s*=\s*"[^"]+"',
f'version = "{new_version}"',
original,
count=1,
)
if count != 1:
raise RuntimeError(f"could not rewrite version in {path}")
# The lockstep graph is circular: omnigent pins both SDKs, the
# client pins omnigent back, and the ui-sdk pins the client. All
# three names must be stamped or the resolver still dead-ends.
updated = re.sub(
r'"(omnigent(?:-client|-ui-sdk)?)==[^"]+"',
rf'"\1=={new_version}"',
updated,
)
path.write_text(updated)
return original
def _stamp_versions(new_version: str) -> dict[Path, str]:
"""Stamp `new_version` into all three pyprojects. Returns originals for restore."""
backups: dict[Path, str] = {}
for path in _pyproject_paths():
backups[path] = set_version_in_pyproject(path, new_version)
return backups
def _restore_versions(backups: dict[Path, str]) -> None:
for path, original in backups.items():
path.write_text(original)
def _clean_build_artifacts() -> None:
"""Remove stale build outputs.
Old Vite bundles in ``omnigent/server/static/web-ui/``
accumulate uniquely-hashed JS chunk filenames between builds.
Without this sweep, orphan files get bundled into the main wheel
and push it over the 10 MB Workspace upload limit.
"""
root = _repo_root()
targets = [
root / "omnigent" / "server" / "static" / "web-ui",
root / "dist",
root / "build",
root / "omnigent.egg-info",
]
for target in targets:
if target.exists():
_log(f"cleaning {target.relative_to(root)}")
shutil.rmtree(target)
def _build_wheels(skip_web_ui: bool) -> list[Path]:
"""Invoke build.sh and return the resulting wheel paths."""
root = _repo_root()
build_sh = _deploy_dir() / "build.sh"
env = os.environ.copy()
if skip_web_ui:
env["SKIP_WEB_UI"] = "1"
_log(f"$ {build_sh}" + (" (SKIP_WEB_UI=1)" if skip_web_ui else ""))
subprocess.run([str(build_sh)], cwd=root, env=env, check=True)
wheels = sorted((root / "dist").glob("*.whl"))
if not wheels:
raise RuntimeError("build.sh produced no wheels")
return wheels
@dataclass(frozen=True)
class _ClassifiedWheels:
"""Result of sorting built wheels by size for upload routing.
:param main: The top-level ``omnigent`` wheel — always uploaded
with the ``[databricks]`` extra.
:param small: Wheels ≤ 10 MB. Uploaded into the bundle's
``source_code_path`` and referenced by relative path.
:param oversize: Wheels > 10 MB. These cannot be used by the
uv-based app payload because ``uv lock`` validates local path
sources before the bundle is synced.
"""
main: Path
small: list[Path]
oversize: list[Path]
def _classify_wheels(wheels: Iterable[Path]) -> _ClassifiedWheels:
main_wheel = next(w for w in wheels if w.name.startswith("omnigent-"))
small: list[Path] = []
oversize: list[Path] = []
for wheel in wheels:
if wheel.stat().st_size <= _WORKSPACE_WHEEL_LIMIT_BYTES:
small.append(wheel)
else:
oversize.append(wheel)
return _ClassifiedWheels(main=main_wheel, small=small, oversize=oversize)
def _wheel_version(wheel: Path, prefix: str) -> str:
"""Extract a deploy version from an Omnigent wheel filename.
:param wheel: Built wheel path, e.g.
``dist/omnigent-0.1.0.post123-py3-none-any.whl``.
:param prefix: Expected wheel filename prefix, e.g. ``"omnigent-"``.
:returns: Version embedded in the wheel filename, e.g.
``"0.1.0.post123"``.
:raises RuntimeError: If the wheel filename does not match the
deploy script's expected wheel naming convention.
"""
pattern = rf"^{re.escape(prefix)}(?P<version>.+)-py3-none-any\.whl$"
match = re.match(pattern, wheel.name)
if not match:
raise RuntimeError(f"unexpected wheel name for {prefix}: {wheel.name}")
return match.group("version")
def _derive_deploy_version_from_wheels(wheels: list[Path]) -> str:
"""Derive the deploy version from reused ``dist/`` wheels.
:param wheels: Wheel files from ``dist/``.
:returns: Shared wheel version, e.g. ``"0.1.0.post123"``.
:raises RuntimeError: If any required wheel is missing or the
wheel versions do not match.
"""
expected_prefixes = ("omnigent-", "omnigent_client-", "omnigent_ui_sdk-")
versions = []
for prefix in expected_prefixes:
matching = [wheel for wheel in wheels if wheel.name.startswith(prefix)]
if len(matching) != 1:
raise RuntimeError(f"expected exactly one {prefix} wheel, found {len(matching)}")
versions.append(_wheel_version(matching[0], prefix))
if len(set(versions)) != 1:
raise RuntimeError(f"wheel versions do not match: {versions}")
return versions[0]
def _sweep_local_src_wheels(keep: set[str]) -> None:
"""Delete Omnigent wheels from src/ whose filename is not in `keep`.
Old deploys accumulate wheels here. Databricks Apps installs the
source directory as a project, so stale wheels can keep local path
sources or lockfile entries alive if we leave them around. Trim to
exactly the wheels we just built.
"""
src = _src_dir()
if not src.exists():
return
for entry in src.iterdir():
if not entry.is_file() or entry.suffix != ".whl":
continue
if not any(entry.name.startswith(p) for p in _WHEEL_PREFIXES):
continue
if entry.name in keep:
continue
_log(f"removing stale local wheel {entry.relative_to(_repo_root())}")
entry.unlink()
def _toml_string(value: str) -> str:
"""Return ``value`` encoded as a TOML basic string.
:param value: String to encode for generated TOML, e.g.
``"./omnigent-0.1.0-py3-none-any.whl"``.
:returns: TOML string literal with quotes and backslashes escaped.
"""
return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"'
def _wheel_source_path(wheel: Path) -> str:
"""Return the path Databricks Apps should use for a wheel source.
:param wheel: Built wheel path, e.g.
``dist/omnigent-0.1.0-py3-none-any.whl``.
:returns: Relative source path for bundled wheels.
:raises SystemExit: If ``wheel`` is oversize.
"""
if wheel.stat().st_size <= _WORKSPACE_WHEEL_LIMIT_BYTES:
return f"./{wheel.name}"
raise SystemExit(
f"wheel {wheel.name} is over 10 MB; uv-based Databricks Apps "
"deploys require all wheels in the app source snapshot"
)
def _uv_source_lines(
main_wheel: Path,
small_wheels: list[Path],
oversize_wheels: list[Path],
) -> list[str]:
"""Build ``[tool.uv.sources]`` lines for the three deploy wheels.
:param main_wheel: Top-level ``omnigent`` wheel path, e.g.
``dist/omnigent-0.1.0-py3-none-any.whl``.
:param small_wheels: Wheels copied into the app source directory.
:param oversize_wheels: Wheels too large for the app source directory.
:returns: TOML lines mapping package names to wheel paths.
"""
wheels = {wheel.name: wheel for wheel in [*small_wheels, *oversize_wheels]}
if main_wheel.name not in wheels:
raise RuntimeError(f"main wheel {main_wheel.name} was not classified for deployment")
source_lines = []
for package_name, wheel_prefix in (
("omnigent", "omnigent-"),
("omnigent-client", "omnigent_client-"),
("omnigent-ui-sdk", "omnigent_ui_sdk-"),
):
wheel = next(wheel for name, wheel in wheels.items() if name.startswith(wheel_prefix))
source = _wheel_source_path(wheel)
source_lines.append(f"{package_name} = {{ path = {_toml_string(source)} }}")
return source_lines
def build_uv_pyproject(
main_wheel: Path,
small_wheels: list[Path],
oversize_wheels: list[Path],
deploy_version: str,
) -> str:
"""Compose the app-level ``pyproject.toml`` for Databricks Apps.
:param main_wheel: Top-level ``omnigent`` wheel path, e.g.
``dist/omnigent-0.1.0-py3-none-any.whl``.
:param small_wheels: Wheels copied into the app source directory.
:param oversize_wheels: Wheels too large for the app source directory.
:param deploy_version: Version stamped into the wheels, e.g.
``"0.1.0.post123"``.
:returns: Complete TOML text for ``src/pyproject.toml``.
"""
source_lines = _uv_source_lines(main_wheel, small_wheels, oversize_wheels)
dependencies = [
f'"omnigent[databricks]=={deploy_version}"',
f'"omnigent-client=={deploy_version}"',
f'"omnigent-ui-sdk=={deploy_version}"',
]
return (
"[project]\n"
'name = "omnigent-databricks-app"\n'
'version = "0.0.0"\n'
f"requires-python = {_toml_string(_APP_REQUIRES_PYTHON)}\n"
"dependencies = [\n"
+ "".join(f" {dependency},\n" for dependency in dependencies)
+ "]\n\n"
"[tool.uv.sources]\n" + "\n".join(source_lines) + "\n"
)
def run_uv_lock(src: Path) -> None:
"""Generate ``uv.lock`` for the Databricks Apps source directory.
:param src: App source directory containing ``pyproject.toml``,
e.g. ``deploy/databricks/src``.
"""
# Honor a caller-supplied UV_INDEX_URL (e.g. a private mirror or proxy);
# otherwise default to public PyPI. UV_INDEX / UV_DEFAULT_INDEX are dropped
# so a stray value in the shell can't shadow the index we lock against.
index_url = os.environ.get("UV_INDEX_URL") or _UV_DEFAULT_INDEX_URL
env = os.environ.copy()
env.pop("UV_INDEX", None)
env.pop("UV_DEFAULT_INDEX", None)
env["UV_INDEX_URL"] = index_url
_log(f"uv lock --python 3.12 --index-url {index_url}")
subprocess.run(
["uv", "lock", "--python", "3.12", "--index-url", index_url],
cwd=src,
env=env,
check=True,
)
def write_uv_dependency_files(
src: Path,
main_wheel: Path,
small_wheels: list[Path],
oversize_wheels: list[Path],
deploy_version: str,
) -> None:
"""Write the uv dependency files Databricks Apps should install.
:param src: App source directory, e.g. ``deploy/databricks/src``.
:param main_wheel: Top-level ``omnigent`` wheel path, e.g.
``dist/omnigent-0.1.0-py3-none-any.whl``.
:param small_wheels: Wheels copied into ``src``.
:param oversize_wheels: Wheels too large for ``src``.
:param deploy_version: Version stamped into the wheels, e.g.
``"0.1.0.post123"``.
"""
requirements = src / "requirements.txt"
if requirements.exists():
_log(f"removing {requirements}; Databricks Apps must use uv")
requirements.unlink()
pyproject = build_uv_pyproject(
main_wheel,
small_wheels,
oversize_wheels,
deploy_version,
)
(src / "pyproject.toml").write_text(pyproject)
_log("src/pyproject.toml:\n" + pyproject)
run_uv_lock(src)
def _smoke_check(wc: WorkspaceClient, app_url: str) -> None:
"""Poll /health on the running app and fail if it never returns 200.
``databricks bundle run`` returns as soon as the app start is
signalled, but uvicorn takes a few extra seconds to bind. Retry
up to a minute to ride out the warm-up; surface the most recent
error if it never goes green.
"""
import urllib.error
import urllib.request
token = wc.config.authenticate()["Authorization"].removeprefix("Bearer ").strip()
url = f"{app_url.rstrip('/')}/health"
last_err: str = ""
for attempt in range(12):
req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"})
try:
with urllib.request.urlopen(req, timeout=30) as resp:
body = resp.read().decode()
if resp.status == 200:
_log(f"/health ok: {body!r}")
return
last_err = f"HTTP {resp.status}: {body!r}"
except urllib.error.HTTPError as exc:
last_err = f"HTTP {exc.code}: {exc.reason}"
except urllib.error.URLError as exc:
last_err = f"URLError: {exc.reason}"
if attempt < 11:
time.sleep(5)
raise RuntimeError(f"/health did not return 200 within 60s; last: {last_err}")
def _assert_clean_tree(skip: bool) -> None:
"""Refuse to deploy if the working tree has uncommitted changes or
HEAD is not at origin/main.
Pass ``--allow-dirty`` to override. Reason: deploys are intended
to come from a known commit on main, so the deployed app code is
reproducible from git history. A dirty tree means whatever we
deploy is not reachable from any commit, which is the cause of
nearly every "wait, what's actually live?" debugging session.
"""
root = _repo_root()
if skip:
_log("--allow-dirty: skipping clean-tree assertion")
return
status = subprocess.run(
["git", "status", "--porcelain"],
cwd=root,
check=True,
capture_output=True,
text=True,
).stdout
if status.strip():
raise SystemExit(
"working tree has uncommitted changes:\n"
+ status
+ "\ncommit or stash, or pass --allow-dirty to override."
)
subprocess.run(["git", "fetch", "origin", "main", "--quiet"], cwd=root, check=True)
head = subprocess.run(
["git", "rev-parse", "HEAD"],
cwd=root,
check=True,
capture_output=True,
text=True,
).stdout.strip()
main = subprocess.run(
["git", "rev-parse", "origin/main"],
cwd=root,
check=True,
capture_output=True,
text=True,
).stdout.strip()
if head != main:
raise SystemExit(
f"HEAD ({head}) is not at origin/main ({main}); "
f"rebase or pass --allow-dirty to override."
)
_log(f"clean tree at origin/main {head[:12]}")
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--app-name",
required=True,
help="Databricks App name, e.g. 'omnigent'.",
)
parser.add_argument(
"--lakebase-branch",
required=True,
help=("Full Lakebase branch resource path, e.g. 'projects/omnigent/branches/production'."),
)
parser.add_argument(
"--lakebase-database",
required=True,
help=(
"Full Lakebase database resource path, e.g. "
"'projects/omnigent/branches/production/databases/databricks-postgres'."
),
)
parser.add_argument(
"--volume-name",
required=True,
help=(
"UC Volume full name (catalog.schema.volume) for artifact storage, "
"e.g. 'main.omnigent.artifacts'."
),
)
parser.add_argument(
"--compute-size",
default="LARGE",
choices=["SMALL", "MEDIUM", "LARGE"],
help="App compute size. Pinned via databricks.yml so it doesn't drift.",
)
parser.add_argument(
"--otel-table-schema",
default="main.omnigent_logs",
help=(
"UC schema (catalog.schema) holding the OTel destination tables. "
"The Databricks Apps platform writes logs/metrics/spans to "
"<schema>.otel_{logs,metrics,spans}."
),
)
parser.add_argument(
"--target",
default="prod",
help=(
"DAB target from databricks.yml selecting the destination "
"workspace, e.g. 'prod'. The authenticating identity "
"(--profile or env) must belong to the target's workspace."
),
)
parser.add_argument(
"--profile",
default=None,
help=(
"Databricks CLI profile to authenticate with. Omit when running "
"with env-based auth (DATABRICKS_HOST + DATABRICKS_CLIENT_ID + "
"DATABRICKS_CLIENT_SECRET)."
),
)
parser.add_argument(
"--version",
default=None,
help=(
"Explicit PEP 440 version to stamp into pyprojects for this "
"deploy. Default: <base-version>.post<unix-ts>."
),
)
parser.add_argument(
"--skip-build",
action="store_true",
help="Reuse existing dist/ wheels — skip the npm + uv build.",
)
parser.add_argument(
"--skip-web-ui",
action="store_true",
help="Build without the SPA (API-only deploy).",
)
parser.add_argument(
"--app-url",
default=None,
help=(
"Override the smoke-check URL. By default the script reads "
"the URL from the App resource returned by the SDK."
),
)
parser.add_argument(
"--no-smoke-check",
action="store_true",
help="Skip the post-deploy /health check.",
)
parser.add_argument(
"--keep-version-bump",
action="store_true",
help=(
"Don't restore pyproject.toml files after build. Useful "
"when you want the bumped version to land in a commit "
"(e.g. CI auto-tagged release builds)."
),
)
parser.add_argument(
"--allow-dirty",
action="store_true",
help=(
"Skip the assertion that the working tree is clean and at "
"origin/main. Deploys are normally required to be from a "
"known commit on main."
),
)
return parser.parse_args()
def _clear_env_vars() -> None:
for name in _ENV_VARS_TO_CLEAR:
if name in os.environ:
_log(f"unsetting {name} to avoid leaking into the SDK")
del os.environ[name]
def _ensure_bound(args: argparse.Namespace) -> None:
"""Bind the bundle to the named app if it exists but is unbound.
If the app already exists in the workspace and the bundle has no
Terraform state for it (always true on a fresh checkout),
``bundle deploy`` fails with "An app with the same name already
exists". The fix is ``bundle deployment bind``, which adopts the
existing app into the bundle's state.
There's no clean way to ask "is the bundle already bound to this
app" — ``bundle summary`` reflects what's *declared* in the YAML,
not what's tracked by Terraform. So we just always try to bind
when the app exists and treat "already managed by Terraform" as
success.
"""
# Late-import so --help works without the SDK.
from databricks.sdk import WorkspaceClient
from databricks.sdk.errors.platform import NotFound
wc = WorkspaceClient(profile=args.profile) if args.profile else WorkspaceClient()
try:
wc.apps.get(name=args.app_name)
except NotFound:
_log(f"app {args.app_name!r} not found; bundle deploy will create it")
return
_log(f"binding bundle resource {_BUNDLE_RESOURCE_KEY!r} → app {args.app_name!r}")
result = subprocess.run(
[
"databricks",
"bundle",
"deployment",
"bind",
_BUNDLE_RESOURCE_KEY,
args.app_name,
"--target",
args.target,
"--auto-approve",
*_profile_arg(args),
*_bundle_vars(args),
],
cwd=_deploy_dir(),
capture_output=True,
text=True,
)
if result.returncode == 0:
return
combined = "".join(stream for stream in (result.stdout, result.stderr) if stream)
if "already managed" in combined.lower() or "already bound" in combined.lower():
_log("bundle already bound to this app; continuing")
return
sys.stderr.write(combined)
raise SystemExit(f"bundle deployment bind failed (exit {result.returncode})")
def _ensure_compute_size(
wc: WorkspaceClient,
app_name: str,
desired: str,
) -> None:
"""Resize the app to `desired` if it's not already there.
The bundle's Terraform databricks_app resource can't update
compute_size on an existing app (the old apps update API rejects
it). The newer ``apps.create_update`` endpoint can. Run that here
so the subsequent bundle deploy sees no diff and doesn't error.
"""
from databricks.sdk.errors.platform import NotFound
from databricks.sdk.service.apps import App, ComputeSize
try:
current = wc.apps.get(name=app_name)
except NotFound:
_log(f"app {app_name!r} not found; bundle deploy will create at {desired}")
return
current_value = current.compute_size.value if current.compute_size else None
if current_value == desired:
_log(f"compute_size already {desired}; skipping resize")
return
_log(f"resizing app {app_name!r}: {current_value} → {desired}")
wc.apps.create_update_and_wait(
app_name=app_name,
update_mask="compute_size",
app=App(name=app_name, compute_size=ComputeSize(desired)),
)
def _bundle_vars(args: argparse.Namespace) -> list[str]:
"""CLI args to pass to `databricks bundle` as --var pairs."""
return [
"--var",
f"app_name={args.app_name}",
"--var",
f"lakebase_branch={args.lakebase_branch}",
"--var",
f"lakebase_database={args.lakebase_database}",
"--var",
f"volume_name={args.volume_name}",
"--var",
f"otel_table_schema={args.otel_table_schema}",
]
def _profile_arg(args: argparse.Namespace) -> list[str]:
return ["--profile", args.profile] if args.profile else []
def _ensure_app_sp_uc_traversal(
args: argparse.Namespace,
app_sp: str | None,
) -> None:
"""Grant USE_CATALOG + USE_SCHEMA to the app SP on the volume's parents.
Apps' ``uc_securable`` only grants the leaf (WRITE_VOLUME); the
SP can boot but 403s on first volume read if the parent catalog
doesn't grant USE to ``account users``. Idempotent.
"""
if not app_sp:
_log("app SP not resolved yet; skipping UC traversal grants")
return
parts = args.volume_name.split(".")
if len(parts) != 3:
raise SystemExit(f"--volume-name {args.volume_name!r} must be catalog.schema.volume")
catalog, schema_only, _ = parts
schema_fqn = f"{catalog}.{schema_only}"
import json as _json
for kind, fqn, priv in (
("catalog", catalog, "USE_CATALOG"),
("schema", schema_fqn, "USE_SCHEMA"),
):
_log(f"granting {priv} on {kind} {fqn} → app SP {app_sp}")
payload = _json.dumps({"changes": [{"principal": app_sp, "add": [priv]}]})
subprocess.run(
[
"databricks",
"grants",
"update",
kind,
fqn,
*_profile_arg(args),
"--json",
payload,
],
check=True,
capture_output=True,
text=True,
)
def main() -> int:
args = _parse_args()
_clear_env_vars()
_assert_clean_tree(skip=args.allow_dirty)
base_version = _read_base_version()
deploy_version = _compute_deploy_version(base_version, args.version)
_log(f"deploy version: {deploy_version} (base: {base_version})")
backups: dict[Path, str] = {}
try:
if not args.skip_build:
_clean_build_artifacts()
backups = _stamp_versions(deploy_version)
wheels = _build_wheels(skip_web_ui=args.skip_web_ui)
else:
dist = _repo_root() / "dist"
wheels = sorted(dist.glob("*.whl"))
if not wheels:
raise SystemExit("--skip-build was set but dist/ has no wheels to redeploy")
wheel_version = _derive_deploy_version_from_wheels(wheels)
if args.version and args.version != wheel_version:
raise SystemExit(
f"--version {args.version!r} does not match reused wheel "
f"version {wheel_version!r}"
)
deploy_version = wheel_version
_log(f"deploy version from reused wheels: {deploy_version}")
_log(f"reusing wheels: {[w.name for w in wheels]}")
finally:
if backups and not args.keep_version_bump:
_restore_versions(backups)
classified = _classify_wheels(wheels)
for wheel in wheels:
size_mb = wheel.stat().st_size / 1024 / 1024
_log(f" {wheel.name} {size_mb:.2f} MB")
if classified.oversize:
raise SystemExit(
"uv-based Databricks Apps deploys require all Omnigent wheels "
"to fit in the app source snapshot. Rebuild with --skip-web-ui "
"or reduce wheel size; UC Volume wheel paths are not used "
"because uv lock validates path sources locally."
)
# Late-import the SDK so `--help` works without it installed.
from databricks.sdk import WorkspaceClient
wc = WorkspaceClient(profile=args.profile) if args.profile else WorkspaceClient()
# 1) Prep the bundle's source_code_path (src/) — sweep stale
# wheels locally, then copy the new small wheels in.
src = _src_dir()
src.mkdir(parents=True, exist_ok=True)
_sweep_local_src_wheels(keep={w.name for w in classified.small})
for wheel in classified.small:
dest = src / wheel.name
_log(f"copy {wheel.name} → {dest.relative_to(_repo_root())}")
shutil.copy2(wheel, dest)
# 2) Generate pyproject.toml + uv.lock. Remove requirements.txt
# first because Databricks Apps gives it precedence over uv.
write_uv_dependency_files(
src,
classified.main,
classified.small,
classified.oversize,
deploy_version,
)
# 4) Bind the bundle to the existing app (if any).
_ensure_bound(args)
# 4a) Reconcile compute_size out-of-band. Terraform's databricks_app
# update path doesn't support compute_size changes ("not supported
# in this update API"). The new SDK apps.create_update endpoint
# does — call it ourselves so the subsequent bundle deploy sees no
# diff. Skipped when already matching.
_ensure_compute_size(wc, args.app_name, args.compute_size)
# 5) databricks bundle deploy --target <target> (syncs src/ to the
# bundle workspace folder and creates/updates the app resource).
_log(f"databricks bundle deploy --target {args.target}")
subprocess.run(
[
"databricks",
"bundle",
"deploy",
"--target",
args.target,
*_profile_arg(args),
*_bundle_vars(args),
],
cwd=_deploy_dir(),
check=True,
)
# 5) databricks bundle run <key> --target <target> (starts/restarts
# the app with the just-deployed source).
_log(f"databricks bundle run {_BUNDLE_RESOURCE_KEY} --target {args.target}")
subprocess.run(
[
"databricks",
"bundle",
"run",
_BUNDLE_RESOURCE_KEY,
"--target",
args.target,
*_profile_arg(args),
*_bundle_vars(args),
],
cwd=_deploy_dir(),
check=True,
)
# 6) Resolve URL + smoke-check.
app = wc.apps.get(name=args.app_name)
app_url = args.app_url or app.url
# 7) Grant the app SP USE_CATALOG + USE_SCHEMA on the volume's
# parents (Apps' uc_securable resource only grants the leaf).
_ensure_app_sp_uc_traversal(args, app.service_principal_client_id)
if not args.no_smoke_check:
if not app_url:
raise SystemExit("no app URL available for smoke check")
_smoke_check(wc, app_url)
_log(f"done. app: {app_url}")
return 0
if __name__ == "__main__":
sys.exit(main())
deploy/databricks/grant_sp_perms.py
+179
View File
@@ -0,0 +1,179 @@
"""
Grant a Databricks App service principal Lakebase schema privileges.
Run this after ``wc.apps.create`` creates the app service principal and
before ``wc.apps.deploy`` starts the app. The app needs these grants so
Alembic can create and migrate Omnigent tables on first boot.
"""
from __future__ import annotations
import argparse
import sys
from typing import Protocol, cast
import psycopg
from databricks.sdk import WorkspaceClient
class _GrantArgs(Protocol):
"""
Parsed CLI arguments for the grant helper.
:param app_name: Databricks App name, e.g. ``"omnigent"``.
:param lakebase_endpoint: Full Lakebase endpoint resource path, e.g.
``"projects/omnigent/branches/production/endpoints/primary"``.
:param database: PostgreSQL database name, e.g.
``"databricks_postgres"``.
:param profile: Optional Databricks CLI profile name, e.g.
``"<your-profile>"``.
"""
app_name: str
lakebase_endpoint: str
database: str
profile: str | None
def _parse_args() -> _GrantArgs:
"""
Parse command-line arguments for the grant helper.
:returns: Parsed CLI arguments.
"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--app-name",
required=True,
help="Databricks App name, e.g. 'omnigent'.",
)
parser.add_argument(
"--lakebase-endpoint",
required=True,
help=(
"Full Lakebase endpoint resource path, e.g. "
"'projects/omnigent/branches/production/endpoints/primary'."
),
)
parser.add_argument(
"--database",
required=True,
help="PostgreSQL database name, e.g. 'databricks_postgres'.",
)
parser.add_argument(
"--profile",
default=None,
help="Optional Databricks CLI profile name, e.g. '<your-profile>'.",
)
return cast(_GrantArgs, parser.parse_args())
def _grant_sql(sp_uuid: str) -> str:
"""
Build the GRANT statement block for an app service principal.
:param sp_uuid: App service principal client ID, e.g.
``"00000000-0000-0000-0000-000000000000"``. Lakebase creates a
PostgreSQL role with this identifier when the app is created.
:returns: SQL statements granting schema, table, and sequence
privileges in the PostgreSQL ``public`` schema.
"""
escaped = sp_uuid.replace('"', '""')
quoted = f'"{escaped}"'
return f"""
GRANT ALL ON SCHEMA public TO {quoted};
GRANT ALL ON ALL TABLES IN SCHEMA public TO {quoted};
GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO {quoted};
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO {quoted};
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO {quoted};
"""
def _resolve_endpoint_host(wc: WorkspaceClient, endpoint_name: str) -> str | None:
"""
Resolve the Lakebase endpoint hostname.
:param wc: Databricks workspace client.
:param endpoint_name: Full Lakebase endpoint resource path, e.g.
``"projects/omnigent/branches/production/endpoints/primary"``.
:returns: Endpoint hostname, or ``None`` when the endpoint is not ready.
"""
endpoint = wc.postgres.get_endpoint(name=endpoint_name)
if endpoint.status is None or endpoint.status.hosts is None:
return None
return endpoint.status.hosts.host
def _build_conn_params(
wc: WorkspaceClient, host: str, database: str, endpoint_name: str
) -> dict[str, str]:
"""
Build psycopg connection params using a short-lived Lakebase OAuth token.
Returned as keyword params (not a hand-built conninfo string) so the
token — which we don't control the contents of — is never string-
interpolated into a DSN where whitespace or ``key=value`` metacharacters
could be mis-parsed.
:param wc: Databricks workspace client.
:param host: Lakebase endpoint hostname, e.g.
``"example.database.cloud.databricks.com"``.
:param database: PostgreSQL database name, e.g.
``"databricks_postgres"``.
:param endpoint_name: Full Lakebase endpoint resource path, e.g.
``"projects/omnigent/branches/production/endpoints/primary"``.
:returns: psycopg connection keyword params for the current user.
"""
cred = wc.postgres.generate_database_credential(endpoint=endpoint_name)
pg_user = wc.current_user.me().user_name
return {
"host": host,
"port": "5432",
"dbname": database,
"user": pg_user,
"password": cred.token,
"sslmode": "require",
}
def main() -> int:
"""
Resolve the app service principal and apply Lakebase grants.
:returns: Process exit code, ``0`` on success and ``1`` when the app
or Lakebase endpoint is not ready.
"""
args = _parse_args()
wc = WorkspaceClient(profile=args.profile) if args.profile else WorkspaceClient()
app = wc.apps.get(name=args.app_name)
sp_uuid = app.service_principal_client_id
if not sp_uuid:
print(
f"ERROR: app '{args.app_name}' has no service_principal_client_id "
"yet; wait for wc.apps.create() to finish.",
file=sys.stderr,
)
return 1
host = _resolve_endpoint_host(wc, args.lakebase_endpoint)
if not host:
print(
f"ERROR: endpoint '{args.lakebase_endpoint}' has no hostname "
"in status; wait for the endpoint to become ACTIVE.",
file=sys.stderr,
)
return 1
print(f"==> Granting public schema privileges to app SP {sp_uuid}")
params = _build_conn_params(wc, host, args.database, args.lakebase_endpoint)
with psycopg.connect(autocommit=True, **params) as conn, conn.cursor() as cur:
cur.execute(_grant_sql(sp_uuid))
print("Done. The app can create and migrate Omnigent tables on first boot.")
return 0
if __name__ == "__main__":
sys.exit(main())
deploy/databricks/src/app.py
+228
View File
@@ -0,0 +1,228 @@
"""Databricks Apps entry point for omnigent.
Starts omnigent with Lakebase (managed PostgreSQL) as the
database and UC Volumes as the artifact store.
"""
from __future__ import annotations
import logging
import os
import sys
import threading
import time
import traceback
logging.basicConfig(level=logging.INFO, stream=sys.stderr, force=True)
logger = logging.getLogger("omnigent-app")
# ── Lakebase token cache ──────────────────────────────────
#
# Lakebase tokens are valid for ~60 minutes. The previous design
# minted a fresh token on every new physical Postgres connection
# inside SQLAlchemy's ``do_connect`` event hook — a synchronous
# Databricks SDK HTTPS round-trip costing 100–300 ms per call.
# Under the 200-runner load test that meant ~20 mints/minute as
# the pool churned overflow connections, with each mint blocking
# whatever thread (sometimes the asyncio event-loop thread) was
# establishing the connection.
#
# This cache mints once per endpoint and reuses the token across
# all subsequent ``do_connect`` calls until the TTL expires. 50
# minutes leaves a 10-minute safety margin before Lakebase rejects
# the token. Concurrent first-time mints are NOT serialized — we
# release the lock around the SDK call so a thundering herd of
# initial connections all mints once each (worst case) rather than
# waiting on a single in-flight mint. The cache is then populated
# atomically with whichever mint finishes first; the late losers
# just overwrite with their own (equally-valid) token.
_TOKEN_TTL_SECONDS = 50 * 60
_token_cache: dict[str, tuple[str, float]] = {}
_token_cache_lock = threading.Lock()
try:
import sqlalchemy
from databricks.sdk import WorkspaceClient
# ── Configuration ──────────────────────────────────────────
# Required env vars — injected by Databricks Apps runtime from
# the resources declared in databricks.yml / app.yaml.
LAKEBASE_ENDPOINT = os.environ["AP_LAKEBASE_ENDPOINT"]
VOLUME_PATH = os.environ["AP_ARTIFACT_VOLUME_PATH"]
PGHOST = os.environ["PGHOST"]
PGDATABASE = os.environ["PGDATABASE"]
PGUSER = os.environ["PGUSER"]
# Optional with documented defaults.
# Databricks Apps expects the app to listen on DATABRICKS_APP_PORT
# (8000 by convention) — deliberately decoupled from the CLI's
# local-server default (6767 in host/local_server.py).
PORT = int(os.environ.get("DATABRICKS_APP_PORT", "8000"))
PGPORT = os.environ.get("PGPORT", "5432")
PGSSLMODE = os.environ.get("PGSSLMODE", "require")
# Recycle DB connections before Lakebase 60-min token expiry.
# 300s (5 min) is conservative — well under the 60-min token TTL.
POOL_RECYCLE_SECONDS = int(os.environ.get("AP_POOL_RECYCLE_SECONDS", "300"))
logger.info(
"Config: PGHOST=%s PGDATABASE=%s PGUSER=%s VOLUME=%s PORT=%d",
PGHOST,
PGDATABASE,
PGUSER,
VOLUME_PATH,
PORT,
)
# ── Lakebase token injection ──────────────────────────────
_workspace_client = WorkspaceClient()
def _get_cached_token(endpoint: str) -> str:
"""Return a cached Lakebase token for ``endpoint``, minting if needed.
Fast path: cached token whose expiry is in the future is returned
without contacting the workspace. Slow path: mint a new token via
the Databricks SDK (synchronous HTTPS). The mint runs OUTSIDE the
cache lock so multiple concurrent first-time mints don't serialize
behind one another — the last winner writes the cache, which is
safe since every minted token is independently valid for ~60 min.
:param endpoint: Lakebase endpoint resource name, e.g.
``"projects/foo/branches/production/endpoints/primary"``.
:returns: A Lakebase database credential token.
:raises RuntimeError: If the SDK returns a credential with no token.
"""
now = time.monotonic()
with _token_cache_lock:
cached = _token_cache.get(endpoint)
if cached is not None and cached[1] > now:
return cached[0]
credential = _workspace_client.postgres.generate_database_credential(
endpoint=endpoint,
)
if credential.token is None:
raise RuntimeError("Lakebase credential response did not include a token")
with _token_cache_lock:
_token_cache[endpoint] = (credential.token, now + _TOKEN_TTL_SECONDS)
return credential.token
# SQLAlchemy fixes the signature of do_connect; the dialect /
# conn_rec / cargs args aren't used here, but they have to be
# named so the hook accepts them. Underscore prefix tells the
# linter we know they're unused.
@sqlalchemy.event.listens_for(sqlalchemy.engine.Engine, "do_connect")
def _inject_lakebase_credentials(_dialect, _conn_rec, _cargs, cparams):
if cparams.get("host") != PGHOST:
return
cparams["password"] = _get_cached_token(LAKEBASE_ENDPOINT)
cparams["sslmode"] = PGSSLMODE
# ── Start omnigent ─────────────────────────────────────
import tempfile
from pathlib import Path
import uvicorn
from omnigent.runtime import init as init_runtime
from omnigent.runtime import telemetry
from omnigent.runtime.agent_cache import AgentCache
from omnigent.runtime.caps import RuntimeCaps
from omnigent.server.app import create_app
from omnigent.server.auth import create_auth_provider
# OTel: the Databricks Apps platform auto-injects
# OTEL_EXPORTER_OTLP_ENDPOINT when `telemetry_export_destinations`
# is set on the app — telemetry.init() picks that up and routes
# OTLP to the platform collector, which writes to the configured
# UC tables. No-op if neither OTEL nor MLflow env vars are set.
telemetry.init()
from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore
from omnigent.stores.artifact_store.databricks_volumes import (
DatabricksVolumesArtifactStore,
)
from omnigent.stores.comment_store.sqlalchemy_store import (
SqlAlchemyCommentStore,
)
from omnigent.stores.conversation_store.sqlalchemy_store import (
SqlAlchemyConversationStore,
)
from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore
from omnigent.stores.host_store import HostStore
from omnigent.stores.permission_store.sqlalchemy_store import (
SqlAlchemyPermissionStore,
)
from omnigent.stores.policy_store.sqlalchemy_store import SqlAlchemyPolicyStore
DB_URI = f"postgresql+psycopg://{PGUSER}@{PGHOST}:{PGPORT}/{PGDATABASE}"
ARTIFACT_URI = f"dbfs:{VOLUME_PATH}"
CACHE_DIR = Path(tempfile.mkdtemp(prefix="ap_cache_"))
logger.info("DB_URI: %s", DB_URI[:80])
logger.info("ARTIFACT_URI: %s", ARTIFACT_URI)
# The app SP owns the tables — run any pending Alembic upgrades
# before the stores boot, since the verify-schema check refuses
# to start a stale DB. Idempotent: a no-op when the DB is at head.
from omnigent.db.utils import _run_migrations as _run_alembic_upgrade
_migration_engine = sqlalchemy.create_engine(DB_URI)
try:
_run_alembic_upgrade(_migration_engine, DB_URI)
finally:
_migration_engine.dispose()
agent_store = SqlAlchemyAgentStore(DB_URI)
file_store = SqlAlchemyFileStore(DB_URI)
conversation_store = SqlAlchemyConversationStore(DB_URI)
artifact_store = DatabricksVolumesArtifactStore(ARTIFACT_URI)
file_comment_store = SqlAlchemyCommentStore(DB_URI)
permission_store = SqlAlchemyPermissionStore(DB_URI)
policy_store = SqlAlchemyPolicyStore(DB_URI)
host_store = HostStore(DB_URI)
agent_cache = AgentCache(artifact_store=artifact_store, cache_dir=CACHE_DIR)
init_runtime(
agent_cache=agent_cache,
caps=RuntimeCaps(),
agent_store=agent_store,
file_store=file_store,
conversation_store=conversation_store,
artifact_store=artifact_store,
comment_store=file_comment_store,
policy_store=policy_store,
)
# The Databricks Apps proxy injects ``X-Forwarded-Email`` on
# every request, so we run in header mode. Header is the
# framework default, but pin it explicitly so the hosted product
# keeps its existing behavior regardless of any ambient
# OMNIGENT_AUTH_ENABLED in the deploy env (an explicit
# provider always wins over the enable switch).
os.environ.setdefault("OMNIGENT_AUTH_PROVIDER", "header")
auth_provider = create_auth_provider()
app = create_app(
agent_store=agent_store,
file_store=file_store,
conversation_store=conversation_store,
artifact_store=artifact_store,
agent_cache=agent_cache,
comment_store=file_comment_store,
permission_store=permission_store,
policy_store=policy_store,
host_store=host_store,
auth_provider=auth_provider,
)
if __name__ == "__main__":
logger.info("Starting omnigent on 0.0.0.0:%d", PORT)
uvicorn.run(app, host="0.0.0.0", port=PORT)
except Exception: # noqa: BLE001 — startup catch-all; we want every failure logged
logger.error("FATAL: omnigent failed to start:\n%s", traceback.format_exc())
# Keep the process alive briefly so logs can be captured
import time
time.sleep(30)
sys.exit(1)
deploy/databricks/src/app.yaml
+8
View File
@@ -0,0 +1,8 @@
command: ["opentelemetry-instrument", "python", "app.py"]
env:
- name: AP_LAKEBASE_ENDPOINT
valueFrom: postgres
- name: AP_ARTIFACT_VOLUME_PATH
valueFrom: artifact_volume
- name: OTEL_TRACES_SAMPLER
value: 'always_on'