feat(appkit): support Lakebase Autoscaling x Apps integration natively

pkosiec · pkosiec · commit d422879533ab · 2026-02-24T15:34:46.000+01:00
diff --git a/docs/docs/plugins/lakebase.md b/docs/docs/plugins/lakebase.md
@@ -11,6 +11,7 @@ Provides a PostgreSQL connection pool for Databricks Lakebase Autoscaling with a
 - Automatic OAuth token refresh (1-hour tokens, 2-minute refresh buffer)
 - Token caching to minimize API calls
 - Built-in OpenTelemetry instrumentation (query duration, pool connections, token refresh)
+- AppKit logger configured by default for query and connection events
 
 ## Setting up Lakebase
 
@@ -108,23 +109,26 @@ The required environment variables:
 
 | Variable | Description |
 |---|---|
-| `PGHOST` | Lakebase host |
-| `PGDATABASE` | Database name |
 | `LAKEBASE_ENDPOINT` | Endpoint resource path (e.g. `projects/.../branches/.../endpoints/...`) |
-| `PGSSLMODE` | TLS mode — set to `require` |
+| `PGHOST` | Lakebase host (auto-injected in production by the `postgres` Databricks Apps resource) |
+| `PGDATABASE` | Database name (auto-injected in production by the `postgres` Databricks Apps resource) |
+| `PGSSLMODE` | TLS mode — set to `require` (auto-injected in production by the `postgres` Databricks Apps resource) |
 
-Ensure that those environment variables are set both for local development (`.env` file) and for deployment (`app.yaml` file):
+When deployed to Databricks Apps with a `postgres` database resource configured, `PGHOST`, `PGDATABASE`, `PGSSLMODE`, `PGUSER`, `PGPORT`, and `PGAPPNAME` are automatically injected by the platform. Only `LAKEBASE_ENDPOINT` must be set explicitly:
 
 ```yaml
 env:
   - name: LAKEBASE_ENDPOINT
-    value: projects/{project-id}/branches/{branch-id}/endpoints/primary
-  - name: PGHOST
-    value: {your-lakebase-host}
-  - name: PGDATABASE
-    value: databricks_postgres
-  - name: PGSSLMODE
-    value: require
+    valueFrom: postgres
+```
+
+For local development, set all variables in your `.env` file:
+
+```env
+PGHOST=your-lakebase-host.databricks.com
+PGDATABASE=databricks_postgres
+LAKEBASE_ENDPOINT=projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>
+PGSSLMODE=require
 ```
 
 For the full configuration reference (SSL, pool size, timeouts, logging, ORM examples), see the [`@databricks/lakebase` README](https://github.com/databricks/appkit/blob/main/packages/lakebase/README.md).
diff --git a/docs/static/schemas/plugin-manifest.schema.json b/docs/static/schemas/plugin-manifest.schema.json
@@ -102,6 +102,7 @@
         "uc_function",
         "uc_connection",
         "database",
+        "postgres",
         "genie_space",
         "experiment",
         "app"
@@ -153,6 +154,11 @@
       "enum": ["CAN_CONNECT_AND_CREATE"],
       "description": "Permission for database resources"
     },
+    "postgresPermission": {
+      "type": "string",
+      "enum": ["CAN_CONNECT_AND_CREATE"],
+      "description": "Permission for Postgres resources"
+    },
     "genieSpacePermission": {
       "type": "string",
       "enum": ["CAN_VIEW", "CAN_RUN", "CAN_EDIT", "CAN_MANAGE"],
@@ -328,6 +334,17 @@
             }
           }
         },
+        {
+          "if": {
+            "properties": { "type": { "const": "postgres" } },
+            "required": ["type"]
+          },
+          "then": {
+            "properties": {
+              "permission": { "$ref": "#/$defs/postgresPermission" }
+            }
+          }
+        },
         {
           "if": {
             "properties": { "type": { "const": "genie_space" } },
diff --git a/packages/appkit/src/plugins/lakebase/manifest.json b/packages/appkit/src/plugins/lakebase/manifest.json
@@ -4,7 +4,21 @@
   "displayName": "Lakebase",
   "description": "SQL query execution against Databricks Lakebase Autoscaling",
   "resources": {
-    "required": [],
+    "required": [
+      {
+        "type": "postgres",
+        "alias": "Postgres",
+        "resourceKey": "postgres",
+        "description": "Lakebase Postgres database for persistent storage",
+        "permission": "CAN_CONNECT_AND_CREATE",
+        "fields": {
+          "name": {
+            "env": "LAKEBASE_ENDPOINT",
+            "description": "Lakebase endpoint resource path"
+          }
+        }
+      }
+    ],
     "optional": []
   }
 }
diff --git a/packages/appkit/src/registry/types.generated.ts b/packages/appkit/src/registry/types.generated.ts
@@ -12,6 +12,7 @@ export enum ResourceType {
   UC_FUNCTION = "uc_function",
   UC_CONNECTION = "uc_connection",
   DATABASE = "database",
+  POSTGRES = "postgres",
   GENIE_SPACE = "genie_space",
   EXPERIMENT = "experiment",
   APP = "app",
@@ -47,6 +48,9 @@ export type UcConnectionPermission = "USE_CONNECTION";
 /** Permissions for DATABASE resources */
 export type DatabasePermission = "CAN_CONNECT_AND_CREATE";
 
+/** Permissions for POSTGRES resources */
+export type PostgresPermission = "CAN_CONNECT_AND_CREATE";
+
 /** Permissions for GENIE_SPACE resources */
 export type GenieSpacePermission =
   | "CAN_VIEW"
@@ -71,6 +75,7 @@ export type ResourcePermission =
   | UcFunctionPermission
   | UcConnectionPermission
   | DatabasePermission
+  | PostgresPermission
   | GenieSpacePermission
   | ExperimentPermission
   | AppPermission;
@@ -89,6 +94,7 @@ export const PERMISSION_HIERARCHY_BY_TYPE: Record<
   [ResourceType.UC_FUNCTION]: ["EXECUTE"],
   [ResourceType.UC_CONNECTION]: ["USE_CONNECTION"],
   [ResourceType.DATABASE]: ["CAN_CONNECT_AND_CREATE"],
+  [ResourceType.POSTGRES]: ["CAN_CONNECT_AND_CREATE"],
   [ResourceType.GENIE_SPACE]: ["CAN_VIEW", "CAN_RUN", "CAN_EDIT", "CAN_MANAGE"],
   [ResourceType.EXPERIMENT]: ["CAN_READ", "CAN_EDIT", "CAN_MANAGE"],
   [ResourceType.APP]: ["CAN_USE"],
diff --git a/packages/lakebase/README.md b/packages/lakebase/README.md
@@ -283,18 +283,11 @@ The driver calls `trace.getTracer('@databricks/lakebase')` and `metrics.getMeter
 
 ## AppKit Integration
 
-This driver is also available as part of [@databricks/appkit](https://www.npmjs.com/package/@databricks/appkit):
+This driver is also available as part of [@databricks/appkit](https://www.npmjs.com/package/@databricks/appkit).
 
-```typescript
-import { createLakebasePool } from "@databricks/appkit";
-
-const pool = createLakebasePool();
-```
-
-**Differences between standalone and AppKit:**
+While you can use the standalone driver, it is recommended to use the AppKit plugin for a more complete experience.
 
-- **Standalone** (`@databricks/lakebase`): Silent by default - no logger configured
-- **AppKit** (`@databricks/appkit`): Automatically injects AppKit's logger with scope `appkit:connectors:lakebase`.
+Read the [AppKit documentation](https://databricks.github.io/appkit/docs/plugins/lakebase) for more information.
 
 ## Learn more about Lakebase Autoscaling
 
diff --git a/packages/shared/src/schemas/plugin-manifest.schema.json b/packages/shared/src/schemas/plugin-manifest.schema.json
@@ -102,6 +102,7 @@
         "uc_function",
         "uc_connection",
         "database",
+        "postgres",
         "genie_space",
         "experiment",
         "app"
@@ -153,6 +154,11 @@
       "enum": ["CAN_CONNECT_AND_CREATE"],
       "description": "Permission for database resources"
     },
+    "postgresPermission": {
+      "type": "string",
+      "enum": ["CAN_CONNECT_AND_CREATE"],
+      "description": "Permission for Postgres resources"
+    },
     "genieSpacePermission": {
       "type": "string",
       "enum": ["CAN_VIEW", "CAN_RUN", "CAN_EDIT", "CAN_MANAGE"],
@@ -328,6 +334,17 @@
             }
           }
         },
+        {
+          "if": {
+            "properties": { "type": { "const": "postgres" } },
+            "required": ["type"]
+          },
+          "then": {
+            "properties": {
+              "permission": { "$ref": "#/$defs/postgresPermission" }
+            }
+          }
+        },
         {
           "if": {
             "properties": { "type": { "const": "genie_space" } },
diff --git a/template/.env.tmpl b/template/.env.tmpl
@@ -7,10 +7,11 @@ DATABRICKS_HOST={{.workspaceHost}}
 {{.dotEnv.content}}
 {{- end}}
 {{- if .plugins.lakebase}}
-PGHOST='' # Copy from the Lakebase Postgres UI
-PGDATABASE='databricks_postgres' # Copy from the Lakebase Postgres UI
-LAKEBASE_ENDPOINT='' # Run: databricks postgres list-endpoints projects/{project-id}/branches/{branch-id}
-# PGUSER=''  # optional, defaults to DATABRICKS_CLIENT_ID
+PGHOST=
+PGDATABASE=
+# Run: databricks postgres list-endpoints projects/{project-id}/branches/{branch-id}
+LAKEBASE_ENDPOINT=
+# PGUSER=  # optional, defaults to DATABRICKS_CLIENT_ID
 PGSSLMODE=require
 {{- end}}
 DATABRICKS_APP_PORT=8000
diff --git a/template/app.yaml.tmpl b/template/app.yaml.tmpl
@@ -1,18 +1,5 @@
 command: ['npm', 'run', 'start']
+{{- if .appEnv}}
 env:
-{{- if .plugins.analytics}}
-  - name: DATABRICKS_WAREHOUSE_ID
-    valueFrom: sql-warehouse
+{{.appEnv}}
 {{- end}}
-{{- if .plugins.lakebase}}
-  - name: PGHOST
-    value: "" # Copy from the Lakebase Postgres UI
-  - name: PGDATABASE
-    value: "databricks_postgres" # Copy from the Lakebase Postgres UI
-  - name: LAKEBASE_ENDPOINT
-    value: "" # Run: databricks postgres list-endpoints projects/{project-id}/branches/{branch-id}
-  - name: PGSSLMODE
-    value: "require"
-#  - name: PGUSER
-#    value: "" # optional, defaults to DATABRICKS_CLIENT_ID
-{{- end}}
diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json
@@ -32,7 +32,21 @@
       "description": "SQL query execution against Databricks Lakebase Autoscaling",
       "package": "@databricks/appkit",
       "resources": {
-        "required": [],
+        "required": [
+          {
+            "type": "postgres",
+            "alias": "Postgres",
+            "resourceKey": "postgres",
+            "description": "Lakebase Postgres database for persistent storage",
+            "permission": "CAN_CONNECT_AND_CREATE",
+            "fields": {
+              "name": {
+                "env": "LAKEBASE_ENDPOINT",
+                "description": "Lakebase endpoint resource path"
+              }
+            }
+          }
+        ],
         "optional": []
       }
     },
