Repository: databricks/dbt-databricks
Pull request #938 (2025-02-11 17:55:35)

Prepare for environments with serverless notebooks
<!---
  Include the number of the issue addressed by this PR above if applicable.

  Example:
    resolves #1234

  Please review our pull request review process in CONTRIBUTING.md before your proceed.
-->

### Description

Learned from working with jobs team today that to use environments, we need the tasks to take an environment key...but also learned this doesn't work yet with serverless notebooks.  This code is to prepare for that feature to come to Databricks, to minimize time without this feature for users.  Normally I hate doing that, but there are user workloads blocked on this.

### Checklist

- [x] I have run this code in development and it appears to resolve the stated issue
- [x] This PR includes tests, or tests are not required/relevant for this PR
- [x] I have updated the `CHANGELOG.md` and added information about my change to the "dbt-databricks next" section.



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c1c2d60b..40570724 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,13 +3,13 @@
 ### Features
 
 - Add `auto_liquid_cluster` config to enable Auto Liquid Clustering for Delta-based dbt models ([935](https://github.com/databricks/dbt-databricks/pull/935))
+- Prepare for environments for python models with serverless clusters ([938](https://github.com/databricks/dbt-databricks/pull/938))
 
 ### Fixes
 
 - table_format: iceberg is unblocked for snapshots ([930](https://github.com/databricks/dbt-databricks/pull/930))
 - Fix for regression in glue table listing behavior ([934](https://github.com/databricks/dbt-databricks/pull/934))
 
-
 ### Under the Hood
 
 - Collapsing to a single connection manager (since the old one no longer works) ([910](https://github.com/databricks/dbt-databricks/pull/910))
diff --git a/dbt/adapters/databricks/api_client.py b/dbt/adapters/databricks/api_client.py
index 57bc4ca3..251da5a6 100644
--- a/dbt/adapters/databricks/api_client.py
+++ b/dbt/adapters/databricks/api_client.py
@@ -326,6 +326,10 @@ class JobRunsApi(PollableApi):
     def submit(
         self, run_name: str, job_spec: dict[str, Any], **additional_job_settings: dict[str, Any]
     ) -> str:
+        logger.debug(
+            f"Submitting job with run_name={run_name} and job_spec={job_spec}"
+            " and additional_job_settings={additional_job_settings}"
+        )
         submit_response = self.session.post(
             "/submit", json={"run_name": run_name, "tasks": [job_spec], **additional_job_settings}
         )
diff --git a/dbt/adapters/databricks/python_models/python_config.py b/dbt/adapters/databricks/python_models/python_config.py
index 29aa44ef..5a39a5c9 100644
--- a/dbt/adapters/databricks/python_models/python_config.py
+++ b/dbt/adapters/databricks/python_models/python_config.py
@@ -36,6 +36,8 @@ class PythonModelConfig(BaseModel):
     cluster_id: Optional[str] = None
     http_path: Optional[str] = None
     create_notebook: bool = False
+    environment_key: Optional[str] = None
+    environment_dependencies: list[str] = Field(default_factory=list)
 
 
 class ParsedPythonModel(BaseModel):
diff --git a/dbt/adapters/databricks/python_models/python_submissions.py b/dbt/adapters/databricks/python_models/python_submissions.py
index afcb383c..f98db234 100644
--- a/dbt/adapters/databricks/python_models/python_submissions.py
+++ b/dbt/adapters/databricks/python_models/python_submissions.py
@@ -209,6 +209,8 @@ class PythonJobConfigCompiler:
         self.job_grants = parsed_model.config.python_job_config.grants
         self.acls = parsed_model.config.access_control_list
         self.additional_job_settings = parsed_model.config.python_job_config.dict()
+        self.environment_key = parsed_model.config.environment_key
+        self.environment_deps = parsed_model.config.environment_dependencies
 
     def compile(self, path: str) -> PythonJobDetails:
         job_spec: dict[str, Any] = {
@@ -217,9 +219,20 @@ class PythonJobConfigCompiler:
                 "notebook_path": path,
             },
         }
-        job_spec.update(self.cluster_spec)  # updates 'new_cluster' config
 
         additional_job_config = self.additional_job_settings
+
+        if self.environment_key:
+            job_spec["environment_key"] = self.environment_key
+            if self.environment_deps and not self.additional_job_settings.get("environments"):
+                additional_job_config["environments"] = [
+                    {
+                        "environment_key": self.environment_key,
+                        "spec": {"client": "2", "dependencies": self.environment_deps},
+                    }
+                ]
+        job_spec.update(self.cluster_spec)  # updates 'new_cluster' config
+
         access_control_list = self.permission_builder.build_job_permissions(
             self.job_grants, self.acls
         )


Test patch:
diff --git a/tests/functional/adapter/python_model/fixtures.py b/tests/functional/adapter/python_model/fixtures.py
index 127bcf74..d1a4dd4a 100644
--- a/tests/functional/adapter/python_model/fixtures.py
+++ b/tests/functional/adapter/python_model/fixtures.py
@@ -42,6 +42,32 @@ sources:
         identifier: source
 """
 
+serverless_schema_with_environment = """version: 2
+
+models:
+  - name: my_versioned_sql_model
+    versions:
+      - v: 1
+  - name: my_python_model
+    config:
+      submission_method: serverless_cluster
+      create_notebook: true
+      environment_key: "test_key"
+      environment_dependencies: ["requests"]
+
+sources:
+  - name: test_source
+    loader: custom
+    schema: "{{ var(env_var('DBT_TEST_SCHEMA_NAME_VARIABLE')) }}"
+    quoting:
+      identifier: True
+    tags:
+      - my_test_source_tag
+    tables:
+      - name: test_table
+        identifier: source
+"""
+
 workflow_schema = """version: 2
 
 models:
diff --git a/tests/functional/adapter/python_model/test_python_model.py b/tests/functional/adapter/python_model/test_python_model.py
index 726791df..ce490c0f 100644
--- a/tests/functional/adapter/python_model/test_python_model.py
+++ b/tests/functional/adapter/python_model/test_python_model.py
@@ -113,6 +113,21 @@ class TestServerlessCluster(BasePythonModelTests):
         }
 
 
+@pytest.mark.python
+# @pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster")
+@pytest.mark.skip("Not available in Databricks yet")
+class TestServerlessClusterWithEnvironment(BasePythonModelTests):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "schema.yml": override_fixtures.serverless_schema_with_environment,
+            "my_sql_model.sql": fixtures.basic_sql,
+            "my_versioned_sql_model_v1.sql": fixtures.basic_sql,
+            "my_python_model.py": fixtures.basic_python,
+            "second_sql_model.sql": fixtures.second_sql,
+        }
+
+
 @pytest.mark.python
 @pytest.mark.external
 @pytest.mark.skip_profile("databricks_cluster", "databricks_uc_sql_endpoint")
diff --git a/tests/unit/python/test_python_job_support.py b/tests/unit/python/test_python_job_support.py
index 41f48041..fb996efa 100644
--- a/tests/unit/python/test_python_job_support.py
+++ b/tests/unit/python/test_python_job_support.py
@@ -146,8 +146,16 @@ class TestPythonJobConfigCompiler:
         parsed_model.config.additional_libs = []
         return run_name
 
+    @pytest.fixture
+    def environment_key(self, parsed_model):
+        environment_key = "test_key"
+        parsed_model.config.environment_key = environment_key
+        parsed_model.config.environment_dependencies = ["requests"]
+        return environment_key
+
     def test_compile__empty_configs(self, client, permission_builder, parsed_model, run_name):
         parsed_model.config.python_job_config.dict.return_value = {}
+        parsed_model.config.environment_key = None
         compiler = PythonJobConfigCompiler(client, permission_builder, parsed_model, {})
         permission_builder.build_job_permissions.return_value = []
         details = compiler.compile("path")
@@ -162,7 +170,9 @@ class TestPythonJobConfigCompiler:
         }
         assert details.additional_job_config == {}
 
-    def test_compile__nonempty_configs(self, client, permission_builder, parsed_model, run_name):
+    def test_compile__nonempty_configs(
+        self, client, permission_builder, parsed_model, run_name, environment_key
+    ):
         parsed_model.config.packages = ["foo"]
         parsed_model.config.index_url = None
         parsed_model.config.python_job_config.dict.return_value = {"foo": "bar"}
@@ -176,6 +186,7 @@ class TestPythonJobConfigCompiler:
         details = compiler.compile("path")
         assert details.run_name == run_name
         assert details.job_spec == {
+            "environment_key": environment_key,
             "task_key": "inner_notebook",
             "notebook_task": {
                 "notebook_path": "path",
@@ -185,4 +196,12 @@ class TestPythonJobConfigCompiler:
             "access_control_list": [{"user_name": "user", "permission_level": "IS_OWNER"}],
             "queue": {"enabled": True},
         }
-        assert details.additional_job_config == {"foo": "bar"}
+        assert details.additional_job_config == {
+            "foo": "bar",
+            "environments": [
+                {
+                    "environment_key": environment_key,
+                    "spec": {"client": "2", "dependencies": ["requests"]},
+                }
+            ],
+        }