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 cb513ea8873ba..5614fd403dd5f 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py @@ -97,8 +97,6 @@ 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 @@ -131,6 +129,8 @@ 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,14 +225,6 @@ 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 @@ -248,6 +240,14 @@ 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,10 +297,6 @@ 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 @@ -314,6 +310,10 @@ 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 398050cf1ff8f..19d0adccf558a 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_repos.py @@ -77,24 +77,26 @@ 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): - """ - Tests handling of incorrect parameters passed to ``__init__`` - """ + """Incorrect parameter combinations are rejected at execute, after rendering.""" + op = DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="abc", repo_path="path", branch="abc") with pytest.raises( AirflowException, match="Only one of repo_id or repo_path should be provided, but not both" ): - DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="abc", repo_path="path", branch="abc") + op.execute(None) + op = DatabricksReposUpdateOperator(task_id=TASK_ID, branch="abc") with pytest.raises(AirflowException, match="One of repo_id or repo_path should be provided"): - DatabricksReposUpdateOperator(task_id=TASK_ID, branch="abc") + op.execute(None) + 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" ): - DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123", branch="123", tag="123") + op.execute(None) + op = DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123") with pytest.raises(AirflowException, match="One of branch or tag should be provided"): - DatabricksReposUpdateOperator(task_id=TASK_ID, repo_id="123") + op.execute(None) class TestDatabricksReposDeleteOperator: @@ -140,16 +142,16 @@ def test_delete_with_path(self, db_mock_class): db_mock.delete_repo.assert_called_once_with("123") def test_init_exception(self): - """ - Tests handling of incorrect parameters passed to ``__init__`` - """ + """Incorrect parameter combinations are rejected at execute, after rendering.""" + op = DatabricksReposDeleteOperator(task_id=TASK_ID, repo_id="abc", repo_path="path") with pytest.raises( AirflowException, match="Only one of repo_id or repo_path should be provided, but not both" ): - DatabricksReposDeleteOperator(task_id=TASK_ID, repo_id="abc", repo_path="path") + op.execute(None) + op = DatabricksReposDeleteOperator(task_id=TASK_ID) with pytest.raises(AirflowException, match="One of repo_id repo_path tag should be provided"): - DatabricksReposDeleteOperator(task_id=TASK_ID) + op.execute(None) class TestDatabricksReposCreateOperator: @@ -286,7 +288,8 @@ 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" ): - DatabricksReposCreateOperator(task_id=TASK_ID, git_url=git_url, branch="123", tag="123") + op.execute(None) diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index bfb3791cc10c7..ec735e38dde53 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -23,9 +23,6 @@ providers/amazon/src/airflow/providers/amazon/aws/transfers/s3_to_redshift.py::S providers/anthropic/src/airflow/providers/anthropic/operators/agent.py::AnthropicAgentSessionOperator providers/apache/kafka/src/airflow/providers/apache/kafka/operators/produce.py::ProduceToTopicOperator providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator -providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposCreateOperator -providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposDeleteOperator -providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposUpdateOperator providers/docker/src/airflow/providers/docker/operators/docker.py::DockerOperator providers/google/src/airflow/providers/google/cloud/operators/bigquery.py::BigQueryInsertJobOperator providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator