Prepare for environments with serverless notebooks (#938)

Author: benc-db

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))

dbt/adapters/databricks/api_client.py

@@ -326,6 +326,10 @@ def __init__(self, session: Session, host: str, polling_interval: int, timeout:
     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}
         )

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):

dbt/adapters/databricks/python_models/python_submissions.py

@@ -209,6 +209,8 @@ def __init__(
         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 @@ def compile(self, path: str) -> PythonJobDetails:
                 "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
         )

tests/functional/adapter/python_model/fixtures.py

@@ -42,6 +42,32 @@ def model(dbt, spark):
         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:

tests/functional/adapter/python_model/test_python_model.py

@@ -113,6 +113,21 @@ def models(self):
         }
 
 
+@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")

tests/unit/python/test_python_job_support.py

@@ -146,8 +146,16 @@ def run_name(self, parsed_model):
         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 @@ def test_compile__empty_configs(self, client, permission_builder, parsed_model,
         }
         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 @@ def test_compile__nonempty_configs(self, client, permission_builder, parsed_mode
         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 @@ def test_compile__nonempty_configs(self, client, permission_builder, parsed_mode
             "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"]},
+                }
+            ],
+        }