From 6722407d2d4083bb83c8206af9ba8d8ed49b86ce Mon Sep 17 00:00:00 2001 From: haseebmalik18 Date: Mon, 27 Jul 2026 13:45:05 -0400 Subject: [PATCH] Restore Dag-parse-time validation for Databricks Repos operator arguments #70503 --- .../databricks/operators/databricks_repos.py | 28 +++++++++--------- .../operators/test_databricks_repos.py | 29 +++++++++---------- 2 files changed, 27 insertions(+), 30 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py index 5614fd403dd5f..cb513ea8873ba 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py @@ -97,6 +97,8 @@ def __init__( else: self.git_provider = git_provider self.repo_path = repo_path + if branch is not None and tag is not None: + raise AirflowException("Only one of branch or tag should be provided, but not both") self.branch = branch self.tag = tag @@ -129,8 +131,6 @@ def execute(self, context: Context): :param context: context :return: Repo ID """ - if self.branch is not None and self.tag is not None: - raise AirflowException("Only one of branch or tag should be provided, but not both") payload = { "url": self.git_url, "provider": self.git_provider, @@ -225,6 +225,14 @@ def __init__( self.databricks_conn_id = databricks_conn_id self.databricks_retry_limit = databricks_retry_limit self.databricks_retry_delay = databricks_retry_delay + if branch is not None and tag is not None: + raise AirflowException("Only one of branch or tag should be provided, but not both") + if branch is None and tag is None: + raise AirflowException("One of branch or tag should be provided") + if repo_id is not None and repo_path is not None: + raise AirflowException("Only one of repo_id or repo_path should be provided, but not both") + if repo_id is None and repo_path is None: + raise AirflowException("One of repo_id or repo_path should be provided") self.repo_path = repo_path self.repo_id = repo_id self.branch = branch @@ -240,14 +248,6 @@ def _hook(self) -> DatabricksHook: ) def execute(self, context: Context): - if self.branch is not None and self.tag is not None: - raise AirflowException("Only one of branch or tag should be provided, but not both") - if self.branch is None and self.tag is None: - raise AirflowException("One of branch or tag should be provided") - if self.repo_id is not None and self.repo_path is not None: - raise AirflowException("Only one of repo_id or repo_path should be provided, but not both") - if self.repo_id is None and self.repo_path is None: - raise AirflowException("One of repo_id or repo_path should be provided") if self.repo_path is not None: self.repo_id = self._hook.get_repo_by_path(self.repo_path) if self.repo_id is None: @@ -297,6 +297,10 @@ def __init__( self.databricks_conn_id = databricks_conn_id self.databricks_retry_limit = databricks_retry_limit self.databricks_retry_delay = databricks_retry_delay + if repo_id is not None and repo_path is not None: + raise AirflowException("Only one of repo_id or repo_path should be provided, but not both") + if repo_id is None and repo_path is None: + raise AirflowException("One of repo_id repo_path tag should be provided") self.repo_path = repo_path self.repo_id = repo_id @@ -310,10 +314,6 @@ def _hook(self) -> DatabricksHook: ) def execute(self, context: Context): - if self.repo_id is not None and self.repo_path is not None: - raise AirflowException("Only one of repo_id or repo_path should be provided, but not both") - if self.repo_id is None and self.repo_path is None: - raise AirflowException("One of repo_id repo_path tag should be provided") if self.repo_path is not None: self.repo_id = self._hook.get_repo_by_path(self.repo_path) if self.repo_id is None: diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py index 19d0adccf558a..398050cf1ff8f 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py @@ -77,26 +77,24 @@ def test_update_with_path(self, db_mock_class): db_mock.update_repo.assert_called_once_with("123", {"tag": "v1.0.0"}) def test_init_exception(self): - """Incorrect parameter combinations are rejected at execute, after rendering.""" - op = DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="abc", repo_path="path", branch="abc") + """ + Tests handling of incorrect parameters passed to ``__init__`` + """ with pytest.raises( AirflowException, match="Only one of repo_id or repo_path should be provided, but not both" ): - op.execute(None) + DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="abc", repo_path="path", branch="abc") - op = DatabricksReposUpdateOperator(task_id=TASK_ID, branch="abc") with pytest.raises(AirflowException, match="One of repo_id or repo_path should be provided"): - op.execute(None) + DatabricksReposUpdateOperator(task_id=TASK_ID, branch="abc") - op = DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123", branch="123", tag="123") with pytest.raises( AirflowException, match="Only one of branch or tag should be provided, but not both" ): - op.execute(None) + DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123", branch="123", tag="123") - op = DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123") with pytest.raises(AirflowException, match="One of branch or tag should be provided"): - op.execute(None) + DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123") class TestDatabricksReposDeleteOperator: @@ -142,16 +140,16 @@ def test_delete_with_path(self, db_mock_class): db_mock.delete_repo.assert_called_once_with("123") def test_init_exception(self): - """Incorrect parameter combinations are rejected at execute, after rendering.""" - op = DatabricksReposDeleteOperator(task_id=TASK_ID, repo_id="abc", repo_path="path") + """ + Tests handling of incorrect parameters passed to ``__init__`` + """ with pytest.raises( AirflowException, match="Only one of repo_id or repo_path should be provided, but not both" ): - op.execute(None) + DatabricksReposDeleteOperator(task_id=TASK_ID, repo_id="abc", repo_path="path") - op = DatabricksReposDeleteOperator(task_id=TASK_ID) with pytest.raises(AirflowException, match="One of repo_id repo_path tag should be provided"): - op.execute(None) + DatabricksReposDeleteOperator(task_id=TASK_ID) class TestDatabricksReposCreateOperator: @@ -288,8 +286,7 @@ def test_init_exception(self): with pytest.raises(AirflowException, match=exception_message): op.execute(None) - op = DatabricksReposCreateOperator(task_id=TASK_ID, git_url=git_url, branch="123", tag="123") with pytest.raises( AirflowException, match="Only one of branch or tag should be provided, but not both" ): - op.execute(None) + DatabricksReposCreateOperator(task_id=TASK_ID, git_url=git_url, branch="123", tag="123")