Add Lakebase Autoscaling support for Databricks Asset Bundles

NEXT_CHANGELOG.md

@@ -11,6 +11,7 @@
 
 ### Bundles
 
+* Add support for Lakebase resources (postgres_projects, postgres_branches, postgres_endpoints) ([#4423](https://github.com/databricks/cli/pull/4423))
 * Add missing values to SchemaGrantPrivilege enum ([#4380](https://github.com/databricks/cli/pull/4380))
 * Added support for UC catalogs (only in direct mode) ([#4342](https://github.com/databricks/cli/pull/4342))
 * engine/direct: Fix updates to fields being ignored for database_catalogs, synced_database_tables (([#4388](https://github.com/databricks/cli/pull/4388)))

acceptance/bin/read_id.py

@@ -28,7 +28,8 @@ def get_id_terraform(filename, name):
         if r_name == name:
             for inst in r["instances"]:
                 attribute_values = inst.get("attributes") or {}
-                return attribute_values.get("id")
+                # Try "id" first, fall back to "name" for postgres resources
+                return attribute_values.get("id") or attribute_values.get("name")
 
     print(f"Cannot find resource with {name=}. Available: {available}", file=sys.stderr)
 

acceptance/bundle/invariant/configs/postgres_branch.yml.tmpl

@@ -0,0 +1,14 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  postgres_projects:
+    project:
+      project_id: test-pg-project-$UNIQUE_NAME
+      display_name: Test Postgres Project
+
+  postgres_branches:
+    foo:
+      parent: ${resources.postgres_projects.project.name}
+      branch_id: test-branch-$UNIQUE_NAME
+      no_expiry: true

acceptance/bundle/invariant/configs/postgres_endpoint.yml.tmpl

@@ -0,0 +1,20 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  postgres_projects:
+    project:
+      project_id: test-pg-project-$UNIQUE_NAME
+      display_name: Test Postgres Project
+
+  postgres_branches:
+    branch:
+      parent: ${resources.postgres_projects.project.name}
+      branch_id: test-branch-$UNIQUE_NAME
+      no_expiry: true
+
+  postgres_endpoints:
+    foo:
+      parent: ${resources.postgres_branches.branch.name}
+      endpoint_id: test-endpoint-$UNIQUE_NAME
+      endpoint_type: ENDPOINT_TYPE_READ_WRITE

acceptance/bundle/invariant/configs/postgres_project.yml.tmpl

@@ -0,0 +1,8 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  postgres_projects:
+    foo:
+      project_id: test-pg-project-$UNIQUE_NAME
+      display_name: Test Postgres Project

acceptance/bundle/invariant/test.toml

@@ -30,6 +30,9 @@ EnvMatrix.INPUT_CONFIG = [
     "model.yml.tmpl",
     "model_serving_endpoint.yml.tmpl",
     "pipeline.yml.tmpl",
+    "postgres_branch.yml.tmpl",
+    "postgres_endpoint.yml.tmpl",
+    "postgres_project.yml.tmpl",
     "registered_model.yml.tmpl",
     "schema.yml.tmpl",
     "secret_scope.yml.tmpl",

acceptance/bundle/refschema/out.fields.txt

@@ -3593,6 +3593,131 @@ resources.pipelines.*.permissions.permissions[*].group_name	string	ALL
 resources.pipelines.*.permissions.permissions[*].permission_level	iam.PermissionLevel	ALL
 resources.pipelines.*.permissions.permissions[*].service_principal_name	string	ALL
 resources.pipelines.*.permissions.permissions[*].user_name	string	ALL
+resources.postgres_branches.*.branch_id	string	INPUT	STATE
+resources.postgres_branches.*.create_time	*time.Time	REMOTE
+resources.postgres_branches.*.expire_time	*time.Time	INPUT	STATE
+resources.postgres_branches.*.id	string	INPUT
+resources.postgres_branches.*.is_protected	bool	INPUT	STATE
+resources.postgres_branches.*.lifecycle	resources.Lifecycle	INPUT
+resources.postgres_branches.*.lifecycle.prevent_destroy	bool	INPUT
+resources.postgres_branches.*.modified_status	string	INPUT
+resources.postgres_branches.*.name	string	REMOTE
+resources.postgres_branches.*.no_expiry	bool	INPUT	STATE
+resources.postgres_branches.*.parent	string	ALL
+resources.postgres_branches.*.source_branch	string	INPUT	STATE
+resources.postgres_branches.*.source_branch_lsn	string	INPUT	STATE
+resources.postgres_branches.*.source_branch_time	*time.Time	INPUT	STATE
+resources.postgres_branches.*.spec	*postgres.BranchSpec	REMOTE
+resources.postgres_branches.*.spec.expire_time	*time.Time	REMOTE
+resources.postgres_branches.*.spec.is_protected	bool	REMOTE
+resources.postgres_branches.*.spec.no_expiry	bool	REMOTE
+resources.postgres_branches.*.spec.source_branch	string	REMOTE
+resources.postgres_branches.*.spec.source_branch_lsn	string	REMOTE
+resources.postgres_branches.*.spec.source_branch_time	*time.Time	REMOTE
+resources.postgres_branches.*.spec.ttl	*duration.Duration	REMOTE
+resources.postgres_branches.*.status	*postgres.BranchStatus	REMOTE
+resources.postgres_branches.*.status.current_state	postgres.BranchStatusState	REMOTE
+resources.postgres_branches.*.status.default	bool	REMOTE
+resources.postgres_branches.*.status.expire_time	*time.Time	REMOTE
+resources.postgres_branches.*.status.is_protected	bool	REMOTE
+resources.postgres_branches.*.status.logical_size_bytes	int64	REMOTE
+resources.postgres_branches.*.status.pending_state	postgres.BranchStatusState	REMOTE
+resources.postgres_branches.*.status.source_branch	string	REMOTE
+resources.postgres_branches.*.status.source_branch_lsn	string	REMOTE
+resources.postgres_branches.*.status.source_branch_time	*time.Time	REMOTE
+resources.postgres_branches.*.status.state_change_time	*time.Time	REMOTE
+resources.postgres_branches.*.ttl	*duration.Duration	INPUT	STATE
+resources.postgres_branches.*.uid	string	REMOTE
+resources.postgres_branches.*.update_time	*time.Time	REMOTE
+resources.postgres_branches.*.url	string	INPUT
+resources.postgres_endpoints.*.autoscaling_limit_max_cu	float64	INPUT	STATE
+resources.postgres_endpoints.*.autoscaling_limit_min_cu	float64	INPUT	STATE
+resources.postgres_endpoints.*.create_time	*time.Time	REMOTE
+resources.postgres_endpoints.*.disabled	bool	INPUT	STATE
+resources.postgres_endpoints.*.endpoint_id	string	INPUT	STATE
+resources.postgres_endpoints.*.endpoint_type	postgres.EndpointType	INPUT	STATE
+resources.postgres_endpoints.*.id	string	INPUT
+resources.postgres_endpoints.*.lifecycle	resources.Lifecycle	INPUT
+resources.postgres_endpoints.*.lifecycle.prevent_destroy	bool	INPUT
+resources.postgres_endpoints.*.modified_status	string	INPUT
+resources.postgres_endpoints.*.name	string	REMOTE
+resources.postgres_endpoints.*.no_suspension	bool	INPUT	STATE
+resources.postgres_endpoints.*.parent	string	ALL
+resources.postgres_endpoints.*.settings	*postgres.EndpointSettings	INPUT	STATE
+resources.postgres_endpoints.*.settings.pg_settings	map[string]string	INPUT	STATE
+resources.postgres_endpoints.*.settings.pg_settings.*	string	INPUT	STATE
+resources.postgres_endpoints.*.spec	*postgres.EndpointSpec	REMOTE
+resources.postgres_endpoints.*.spec.autoscaling_limit_max_cu	float64	REMOTE
+resources.postgres_endpoints.*.spec.autoscaling_limit_min_cu	float64	REMOTE
+resources.postgres_endpoints.*.spec.disabled	bool	REMOTE
+resources.postgres_endpoints.*.spec.endpoint_type	postgres.EndpointType	REMOTE
+resources.postgres_endpoints.*.spec.no_suspension	bool	REMOTE
+resources.postgres_endpoints.*.spec.settings	*postgres.EndpointSettings	REMOTE
+resources.postgres_endpoints.*.spec.settings.pg_settings	map[string]string	REMOTE
+resources.postgres_endpoints.*.spec.settings.pg_settings.*	string	REMOTE
+resources.postgres_endpoints.*.spec.suspend_timeout_duration	*duration.Duration	REMOTE
+resources.postgres_endpoints.*.status	*postgres.EndpointStatus	REMOTE
+resources.postgres_endpoints.*.status.autoscaling_limit_max_cu	float64	REMOTE
+resources.postgres_endpoints.*.status.autoscaling_limit_min_cu	float64	REMOTE
+resources.postgres_endpoints.*.status.current_state	postgres.EndpointStatusState	REMOTE
+resources.postgres_endpoints.*.status.disabled	bool	REMOTE
+resources.postgres_endpoints.*.status.endpoint_type	postgres.EndpointType	REMOTE
+resources.postgres_endpoints.*.status.hosts	*postgres.EndpointHosts	REMOTE
+resources.postgres_endpoints.*.status.hosts.host	string	REMOTE
+resources.postgres_endpoints.*.status.pending_state	postgres.EndpointStatusState	REMOTE
+resources.postgres_endpoints.*.status.settings	*postgres.EndpointSettings	REMOTE
+resources.postgres_endpoints.*.status.settings.pg_settings	map[string]string	REMOTE
+resources.postgres_endpoints.*.status.settings.pg_settings.*	string	REMOTE
+resources.postgres_endpoints.*.status.suspend_timeout_duration	*duration.Duration	REMOTE
+resources.postgres_endpoints.*.suspend_timeout_duration	*duration.Duration	INPUT	STATE
+resources.postgres_endpoints.*.uid	string	REMOTE
+resources.postgres_endpoints.*.update_time	*time.Time	REMOTE
+resources.postgres_endpoints.*.url	string	INPUT
+resources.postgres_projects.*.create_time	*time.Time	REMOTE
+resources.postgres_projects.*.default_endpoint_settings	*postgres.ProjectDefaultEndpointSettings	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.autoscaling_limit_max_cu	float64	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.autoscaling_limit_min_cu	float64	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.no_suspension	bool	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.pg_settings	map[string]string	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.pg_settings.*	string	INPUT	STATE
+resources.postgres_projects.*.default_endpoint_settings.suspend_timeout_duration	*duration.Duration	INPUT	STATE
+resources.postgres_projects.*.display_name	string	INPUT	STATE
+resources.postgres_projects.*.history_retention_duration	*duration.Duration	INPUT	STATE
+resources.postgres_projects.*.id	string	INPUT
+resources.postgres_projects.*.lifecycle	resources.Lifecycle	INPUT
+resources.postgres_projects.*.lifecycle.prevent_destroy	bool	INPUT
+resources.postgres_projects.*.modified_status	string	INPUT
+resources.postgres_projects.*.name	string	REMOTE
+resources.postgres_projects.*.pg_version	int	INPUT	STATE
+resources.postgres_projects.*.project_id	string	INPUT	STATE
+resources.postgres_projects.*.spec	*postgres.ProjectSpec	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings	*postgres.ProjectDefaultEndpointSettings	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.autoscaling_limit_max_cu	float64	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.autoscaling_limit_min_cu	float64	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.no_suspension	bool	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.pg_settings	map[string]string	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.pg_settings.*	string	REMOTE
+resources.postgres_projects.*.spec.default_endpoint_settings.suspend_timeout_duration	*duration.Duration	REMOTE
+resources.postgres_projects.*.spec.display_name	string	REMOTE
+resources.postgres_projects.*.spec.history_retention_duration	*duration.Duration	REMOTE
+resources.postgres_projects.*.spec.pg_version	int	REMOTE
+resources.postgres_projects.*.status	*postgres.ProjectStatus	REMOTE
+resources.postgres_projects.*.status.branch_logical_size_limit_bytes	int64	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings	*postgres.ProjectDefaultEndpointSettings	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.autoscaling_limit_max_cu	float64	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.autoscaling_limit_min_cu	float64	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.no_suspension	bool	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.pg_settings	map[string]string	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.pg_settings.*	string	REMOTE
+resources.postgres_projects.*.status.default_endpoint_settings.suspend_timeout_duration	*duration.Duration	REMOTE
+resources.postgres_projects.*.status.display_name	string	REMOTE
+resources.postgres_projects.*.status.history_retention_duration	*duration.Duration	REMOTE
+resources.postgres_projects.*.status.owner	string	REMOTE
+resources.postgres_projects.*.status.pg_version	int	REMOTE
+resources.postgres_projects.*.status.synthetic_storage_size_bytes	int64	REMOTE
+resources.postgres_projects.*.uid	string	REMOTE
+resources.postgres_projects.*.update_time	*time.Time	REMOTE
+resources.postgres_projects.*.url	string	INPUT
 resources.quality_monitors.*.assets_dir	string	ALL
 resources.quality_monitors.*.baseline_table_name	string	ALL
 resources.quality_monitors.*.custom_metrics	[]catalog.MonitorMetric	ALL

acceptance/bundle/resources/postgres_branches/basic/databricks.yml.tmpl

@@ -0,0 +1,23 @@
+bundle:
+  name: deploy-postgres-branch-$UNIQUE_NAME
+
+sync:
+  paths: []
+
+resources:
+  postgres_projects:
+    my_project:
+      project_id: test-pg-proj-$UNIQUE_NAME
+      display_name: "Test Project for Branch"
+      pg_version: 16
+      history_retention_duration: "604800s"
+      default_endpoint_settings:
+        autoscaling_limit_min_cu: 0.5
+        autoscaling_limit_max_cu: 4
+        suspend_timeout_duration: "300s"
+
+  postgres_branches:
+    main:
+      parent: ${resources.postgres_projects.my_project.id}
+      branch_id: main
+      no_expiry: true

acceptance/bundle/resources/postgres_branches/basic/out.requests.direct.json

@@ -0,0 +1,35 @@
+{
+  "method": "POST",
+  "path": "/api/2.0/postgres/projects",
+  "q": {
+    "project_id": "test-pg-proj-[UNIQUE_NAME]"
+  },
+  "body": {
+    "spec": {
+      "default_endpoint_settings": {
+        "autoscaling_limit_max_cu": 4,
+        "autoscaling_limit_min_cu": 0.5,
+        "suspend_timeout_duration": "300s"
+      },
+      "display_name": "Test Project for Branch",
+      "history_retention_duration": "604800s",
+      "pg_version": 16
+    }
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/postgres/projects/test-pg-proj-[UNIQUE_NAME]/branches",
+  "q": {
+    "branch_id": "main"
+  },
+  "body": {
+    "spec": {
+      "no_expiry": true
+    }
+  }
+}
+{
+  "method": "GET",
+  "path": "/api/2.0/postgres/projects/test-pg-proj-[UNIQUE_NAME]/branches/main"
+}

acceptance/bundle/resources/postgres_branches/basic/out.requests.terraform.json

@@ -0,0 +1,36 @@
+{
+  "method": "POST",
+  "path": "/api/2.0/postgres/projects",
+  "q": {
+    "project_id": "test-pg-proj-[UNIQUE_NAME]"
+  },
+  "body": {
+    "spec": {
+      "default_endpoint_settings": {
+        "autoscaling_limit_max_cu": 4,
+        "autoscaling_limit_min_cu": 0.5,
+        "suspend_timeout_duration": "300s"
+      },
+      "display_name": "Test Project for Branch",
+      "history_retention_duration": "604800s",
+      "pg_version": 16
+    }
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/postgres/projects/test-pg-proj-[UNIQUE_NAME]/branches",
+  "q": {
+    "branch_id": "main"
+  },
+  "body": {
+    "parent": "projects/test-pg-proj-[UNIQUE_NAME]",
+    "spec": {
+      "no_expiry": true
+    }
+  }
+}
+{
+  "method": "GET",
+  "path": "/api/2.0/postgres/projects/test-pg-proj-[UNIQUE_NAME]/branches/main"
+}