Enable access control for job runs (#518)

.gitignore

@@ -16,4 +16,5 @@ test.env
 .vscode
 *.log
 logs/
-.venv*
+.venv*
+*.sublime*

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Adding capability to specify compute on a per model basis ([488](https://github.com/databricks/dbt-databricks/pull/488))
 - Selectively persist column docs that have changed between runs of incremental ([513](https://github.com/databricks/dbt-databricks/pull/513))
+- Enabling access control list for job runs (thanks @srggrs!)([518](https://github.com/databricks/dbt-databricks/pull/518))
 
 ## dbt-databricks 1.7.1 (Nov 13, 2023)
 

dbt/adapters/databricks/python_submissions.py

@@ -65,6 +65,12 @@ def _create_work_dir(self, path: str) -> None:
                 f"Error creating work_dir for python notebooks\n {response.content!r}"
             )
 
+    def _update_with_acls(self, cluster_dict: dict) -> dict:
+        acl = self.parsed_model["config"].get("access_control_list", None)
+        if acl:
+            cluster_dict.update({"access_control_list": acl})
+        return cluster_dict
+
     def _upload_notebook(self, path: str, compiled_code: str) -> None:
         b64_encoded_content = base64.b64encode(compiled_code.encode()).decode()
         response = requests.post(
@@ -206,7 +212,7 @@ def check_credentials(self) -> None:
 
     def submit(self, compiled_code: str) -> None:
         cluster_spec = {"new_cluster": self.parsed_model["config"]["job_cluster_config"]}
-        self._submit_through_notebook(compiled_code, cluster_spec)
+        self._submit_through_notebook(compiled_code, self._update_with_acls(cluster_spec))
 
 
 class DBContext:
@@ -374,7 +380,8 @@ def check_credentials(self) -> None:
 
     def submit(self, compiled_code: str) -> None:
         if self.parsed_model["config"].get("create_notebook", False):
-            self._submit_through_notebook(compiled_code, {"existing_cluster_id": self.cluster_id})
+            config = {"existing_cluster_id": self.cluster_id}
+            self._submit_through_notebook(compiled_code, self._update_with_acls(config))
         else:
             context = DBContext(self.credentials, self.cluster_id, self.auth_header)
             command = DBCommand(self.credentials, self.cluster_id, self.auth_header)

tests/unit/python/test_python_submissions.py

@@ -1,7 +1,9 @@
 import unittest
 from unittest.mock import patch, Mock
 
-from dbt.adapters.databricks.python_submissions import DBContext
+from dbt.adapters.databricks.connections import DatabricksCredentials
+
+from dbt.adapters.databricks.python_submissions import DBContext, BaseDatabricksHelper
 
 
 class TestDatabricksPythonSubmissions(unittest.TestCase):
@@ -18,3 +20,40 @@ def test_start_cluster_returns_on_receiving_running_state(self, mock_post, mock_
         context.start_cluster()
 
         mock_get.assert_called_once()
+
+
+class DatabricksTestHelper(BaseDatabricksHelper):
+    def __init__(self, parsed_model: dict, credentials: DatabricksCredentials):
+        self.parsed_model = parsed_model
+        self.credentials = credentials
+
+
+class TestAclUpdate:
+    def test_empty_acl_empty_config(self):
+        helper = DatabricksTestHelper({"config": {}}, DatabricksCredentials())
+        assert helper._update_with_acls({}) == {}
+
+    def test_empty_acl_non_empty_config(self):
+        helper = DatabricksTestHelper({"config": {}}, DatabricksCredentials())
+        assert helper._update_with_acls({"a": "b"}) == {"a": "b"}
+
+    def test_non_empty_acl_empty_config(self):
+        expected_access_control = {
+            "access_control_list": [
+                {"user_name": "user2", "permission_level": "CAN_VIEW"},
+            ]
+        }
+        helper = DatabricksTestHelper({"config": expected_access_control}, DatabricksCredentials())
+        assert helper._update_with_acls({}) == expected_access_control
+
+    def test_non_empty_acl_non_empty_config(self):
+        expected_access_control = {
+            "access_control_list": [
+                {"user_name": "user2", "permission_level": "CAN_VIEW"},
+            ]
+        }
+        helper = DatabricksTestHelper({"config": expected_access_control}, DatabricksCredentials())
+        assert helper._update_with_acls({"a": "b"}) == {
+            "a": "b",
+            "access_control_list": expected_access_control["access_control_list"],
+        }