Update agent examples and quickstart to support Lakebase Autoscaling by jennsun · Pull Request #151 · databricks/app-templates

jennsun · 2026-03-04T19:54:55Z

updated quickstart experience to select instance:

examples work

ex app: https://dbc-fb904e0e-bca7.staging.cloud.databricks.com/apps/agent-longtermj?o=4146361371977516

- Add autoscaling (project/branch) as an alternative to provisioned Lakebase instances across all 3 memory templates - Detect app environment via DATABRICKS_APP_NAME and use PGENDPOINT (auto-injected by platform) for autoscaling in deployed apps - Update quickstart to support creating new Lakebase instances or connecting to existing provisioned/autoscaling instances - Use autoscaling_endpoint parameter (renamed in databricks-ai-bridge) - Add databricks-ai-bridge git dependency for autoscaling support - Update tests to remove update_databricks_yml_lakebase references Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.scripts/source/test_quickstart.py

jennsun · 2026-03-05T23:56:35Z

agent-langgraph-long-term-memory/pyproject.toml

+    # Note: Using git reference until official release
+    "databricks-langchain[memory] @ git+https://github.com/databricks/databricks-ai-bridge@lakebase-autoscaling-support#subdirectory=integrations/langchain",
+    "databricks-ai-bridge @ git+https://github.com/databricks/databricks-ai-bridge@lakebase-autoscaling-support",
    "databricks-agents>=1.9.3",


will remove after release

dhruv0811

Looks good overall!! Left some comments for local testing cleanup and some clarifications on lakebase options.

.gitignore

dhruv0811 · 2026-03-05T23:46:12Z

agent-openai-agents-sdk-short-term-memory/databricks.yml

-          database:
-            instance_name: '<your-lakebase-instance-name>'
-            database_name: 'databricks_postgres'
+        - name: 'postgres'


nit: couple comments around here saying this is for autoscaling, see below for provisioned?

dhruv0811 · 2026-03-05T23:47:47Z

agent-openai-agents-sdk-short-term-memory/.env.example

+# TODO: Update with your Lakebase instance for session storage
+# Option 1: Provisioned instance (set instance name)
+# LAKEBASE_INSTANCE_NAME=
+# Option 2: Autoscaling instance (set project and branch)


Did we not have a third option with just autoscaling_endpoint?

the thinking here is that when users develop locally, they do so via passing in project/branch name. the autoscaling_endpoint is only used when we are in the app environment and we can directly read from PGENDPOINT (since it's not as user friendly to pass in) - hence adding this check: https://github.com/databricks/app-templates/pull/151/changes#diff-1e8ffff9ea9c91fc9dce7d8af129175d1fcfe357eabbcec14705feff13eeacacR53

agent-langgraph-short-term-memory/databricks.yml

agent-langgraph-short-term-memory/.gitignore

agent-langgraph-short-term-memory/agent_server/start_server.py

dhruv0811 · 2026-03-05T23:53:36Z

agent-langgraph-short-term-memory/agent_server/agent.py

        # Check for Lakebase access/connection errors
        if any(keyword in error_msg for keyword in ["permission"]):
            logger.error(f"Lakebase access error: {e}")
+            lakebase_desc = LAKEBASE_INSTANCE_NAME or LAKEBASE_AUTOSCALING_ENDPOINT or f"{LAKEBASE_AUTOSCALING_PROJECT}/{LAKEBASE_AUTOSCALING_BRANCH}"


Is it worth doing some validation like we did in the SDK to ensure that we only have one valid combination of variables set here? Say a user has project and branch set, but accidentally also set instance name, they should be notified rather than silently picking up instance name and short circuiting here.

since the sdk covers all these cases/will throw errors, I think adding validation here would be a bit redundant

agent-langgraph-long-term-memory/databricks.yml

jennsun · 2026-03-06T02:41:55Z

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py

+LAKEBASE_AUTOSCALING_ENDPOINT = os.getenv("PGENDPOINT") if _is_app_env else None
+LAKEBASE_AUTOSCALING_PROJECT = os.getenv("LAKEBASE_AUTOSCALING_PROJECT") or None
+LAKEBASE_AUTOSCALING_BRANCH = os.getenv("LAKEBASE_AUTOSCALING_BRANCH") or None
+


bit of a chicken & egg prob here when DABs are not ready -

in order for us to be able to pass in PGENDPOINT as a var, we need the autoscaling instance as a resource on the app

DAB deployment will overwrite app resources, so it would remove manually-added postgres resource when we deploy

we will therefore use LAKEBASE_AUTOSCALING_PROJECT and LAKEBASE_AUTOSCALING_BRANCH as
static env vars supported by our agent SDK

if we want endpoint extension in future, code will look something like:

# Autoscaling params: in the app environment, PGENDPOINT is provided automatically; # for local dev, use project/branch names directly. _is_app_env = bool(os.getenv("DATABRICKS_APP_NAME")) LAKEBASE_AUTOSCALING_ENDPOINT = os.getenv("PGENDPOINT") if _is_app_env else None LAKEBASE_AUTOSCALING_PROJECT = os.getenv("LAKEBASE_AUTOSCALING_PROJECT") or None LAKEBASE_AUTOSCALING_BRANCH = os.getenv("LAKEBASE_AUTOSCALING_BRANCH") or None

dhruv0811

Added some comments based on a dry run with Claude.

dhruv0811 · 2026-03-09T20:11:51Z

agent-langgraph-long-term-memory/scripts/quickstart.py

    content = databricks_yml.read_text()
-    return "LAKEBASE_INSTANCE_NAME" in content
+    return (
+        "LAKEBASE_INSTANCE_NAME" in content


nit: Is this supposed to be LAKEBASE_INSTANCE_NAME OR (LAKEBASE_AUTOSCALING_PROJECT AND LAKEBASE_AUTOSCALING_BRANCH)

dhruv0811 · 2026-03-09T20:55:14Z

.scripts/source/quickstart.py

        help="Databricks workspace URL (for initial setup)",
        metavar="URL",
    )
    parser.add_argument(


Could we add arguments for lakebase project and branch here? Makes it easier for agents to run quickstart if we support non-interactive mode.

dhruv0811 · 2026-03-09T20:59:16Z

agent-langgraph-long-term-memory/scripts/grant_lakebase_permissions.py

Do we want to sync this script like we sync quickstart btw?

dhruv0811 · 2026-03-09T21:03:02Z

agent-langgraph-long-term-memory/scripts/grant_lakebase_permissions.py

+    )
+    parser.add_argument(
+        "--autoscaling-endpoint",
+        default=os.getenv("LAKEBASE_AUTOSCALING_ENDPOINT"),


This script is the only place that we use LAKEBASE_AUTOSCALING_ENDPOINT. Seems to be confusing for agents, can we either add documentation for this env var on what it refers to?

dhruv0811 · 2026-03-09T21:06:14Z

agent-langgraph-short-term-memory/databricks.yml

-          - name: LAKEBASE_INSTANCE_NAME
-            value: "<your-lakebase-instance-name>"
+          # Autoscaling Lakebase config
+          - name: LAKEBASE_AUTOSCALING_PROJECT


Does it make it easier if we add a step in the quickstart script to update these variables based on what the user provides?

dhruv0811 · 2026-03-09T21:06:38Z

agent-langgraph-short-term-memory/databricks.yml

            value: "300"
          - name: MLFLOW_EXPERIMENT_ID
            value_from: "experiment"
-          - name: LAKEBASE_INSTANCE_NAME


Do we want to keep this var? Maybe just have it commented out, in case people don't go down the autoscaling path?

jennsun mentioned this pull request Mar 4, 2026

Update agent examples to support Lakebase Autoscaling #147

Closed

jennsun changed the title ~~Update agent examples to support Lakebase Autoscaling~~ Update agent examples and quickstart to support Lakebase Autoscaling Mar 4, 2026

jennsun marked this pull request as ready for review March 4, 2026 21:44

jennsun force-pushed the lakebase-autoscaling-updates-v2 branch from 248b026 to e39267c Compare March 5, 2026 23:21

rm dab postgres resource since not ready in apps yet

0722c13

jennsun commented Mar 5, 2026

View reviewed changes

.scripts/source/test_quickstart.py Outdated Show resolved Hide resolved

restore to template scaffolding code

e02ef37

jennsun commented Mar 5, 2026

View reviewed changes

dhruv0811 reviewed Mar 5, 2026

View reviewed changes

jennsun added 9 commits March 5, 2026 15:58

revert gitignore fixes

1e3fd13

.

4793304

revert dab blocks

e8868ab

.

6530c00

revert quickstart test

743d7e2

autoscaling lakebase claude skills

6f5bbb3

specify sp client id

544d3e6

rm PGENDPOINT

9042b65

.

bcfdc4e

jennsun requested review from bbqiu and dhruv0811 March 6, 2026 02:36

jennsun commented Mar 6, 2026

View reviewed changes

jennsun added 6 commits March 5, 2026 18:43

rm passing in endpoint

2b131b5

bundle run

68814f7

append postgres resource and give explicit api for postgres

3a226f3

separate permission granting into separate script

16750c5

include permission granting script in existing lakebase-setup skill

f58f005

update autoscaling sdk to redeploy after adding postgres resource

4a9dee8

dhruv0811 reviewed Mar 9, 2026

View reviewed changes

Conversation

jennsun commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jennsun commented Mar 4, 2026 •

edited

Loading