diff --git a/airflow/providers/dbt/cloud/operators/dbt.py b/airflow/providers/dbt/cloud/operators/dbt.py index a4c04a09190a2..5d8c6baca641e 100644 --- a/airflow/providers/dbt/cloud/operators/dbt.py +++ b/airflow/providers/dbt/cloud/operators/dbt.py @@ -117,7 +117,7 @@ def __init__( self.timeout = timeout self.check_interval = check_interval self.additional_run_config = additional_run_config or {} - self.run_id: int + self.run_id: int | None = None self.deferrable = deferrable def execute(self, context: Context): @@ -135,12 +135,13 @@ def execute(self, context: Context): additional_run_config=self.additional_run_config, ) self.run_id = trigger_job_response.json()["data"]["id"] + print(self.run_id) job_run_url = trigger_job_response.json()["data"]["href"] # Push the ``job_run_url`` value to XCom regardless of what happens during execution so that the job # run can be monitored via the operator link. context["ti"].xcom_push(key="job_run_url", value=job_run_url) - if self.wait_for_termination: + if self.wait_for_termination and isinstance(self.run_id, int): if self.deferrable is False: self.log.info("Waiting for job run %s to terminate.", str(self.run_id)) @@ -197,6 +198,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> int: if event["status"] == "error": raise AirflowException(event["message"]) self.log.info(event["message"]) + self.run_id = event["run_id"] return int(event["run_id"]) def on_kill(self) -> None: @@ -225,7 +227,7 @@ def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage: """ from airflow.providers.openlineage.extractors import OperatorLineage - if self.wait_for_termination is True: + if isinstance(self.run_id, int) and self.wait_for_termination is True: return generate_openlineage_events_from_dbt_cloud_run(operator=self, task_instance=task_instance) return OperatorLineage() diff --git a/tests/providers/dbt/cloud/operators/test_dbt_cloud.py b/tests/providers/dbt/cloud/operators/test_dbt_cloud.py index 4400ebd4abb46..17a929c66fd7b 100644 --- a/tests/providers/dbt/cloud/operators/test_dbt_cloud.py +++ b/tests/providers/dbt/cloud/operators/test_dbt_cloud.py @@ -63,6 +63,12 @@ } +def mock_response_json(response: dict): + run_response = MagicMock(**response) + run_response.json.return_value = response + return run_response + + def setup_module(): # Connection with ``account_id`` specified conn_account_id = Connection( @@ -125,7 +131,10 @@ def test_execute_succeeded_before_getting_deferred( ) @patch("airflow.providers.dbt.cloud.operators.dbt.DbtCloudRunJobOperator.defer") @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_connection") - @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.trigger_job_run") + @patch( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.trigger_job_run", + return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE), + ) def test_execute_failed_before_getting_deferred( self, mock_trigger_job_run, mock_dbt_hook, mock_defer, mock_job_run_status ): @@ -154,7 +163,10 @@ def test_execute_failed_before_getting_deferred( "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_run_status", ) @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_connection") - @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.trigger_job_run") + @patch( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.trigger_job_run", + return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE), + ) def test_dbt_run_job_op_async(self, mock_trigger_job_run, mock_dbt_hook, mock_job_run_status, status): """ Asserts that a task is deferred and an DbtCloudRunJobTrigger will be fired @@ -174,7 +186,9 @@ def test_dbt_run_job_op_async(self, mock_trigger_job_run, mock_dbt_hook, mock_jo dbt_op.execute(MagicMock()) assert isinstance(exc.value.trigger, DbtCloudRunJobTrigger), "Trigger is not a DbtCloudRunJobTrigger" - @patch.object(DbtCloudHook, "trigger_job_run", return_value=MagicMock(**DEFAULT_ACCOUNT_JOB_RUN_RESPONSE)) + @patch.object( + DbtCloudHook, "trigger_job_run", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE) + ) @pytest.mark.parametrize( "job_run_status, expected_output", [ diff --git a/tests/providers/dbt/cloud/utils/test_openlineage.py b/tests/providers/dbt/cloud/utils/test_openlineage.py index 0d93d5a2e140d..ed9c7bfa40939 100644 --- a/tests/providers/dbt/cloud/utils/test_openlineage.py +++ b/tests/providers/dbt/cloud/utils/test_openlineage.py @@ -22,6 +22,7 @@ from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook from airflow.providers.dbt.cloud.operators.dbt import DbtCloudRunJobOperator from airflow.providers.dbt.cloud.utils.openlineage import generate_openlineage_events_from_dbt_cloud_run +from airflow.providers.openlineage.extractors import OperatorLineage TASK_ID = "dbt_test" DAG_ID = "dbt_dag" @@ -130,7 +131,10 @@ def test_generate_events( ) mock_build_task_instance_run_id.return_value = TASK_UUID - generate_openlineage_events_from_dbt_cloud_run(mock_operator, task_instance=mock_task_instance) - assert mock_client.emit.call_count == 4 + + def test_do_not_raise_error_if_runid_not_set_on_operator(self): + operator = DbtCloudRunJobOperator(task_id="dbt-job-runid-taskid", job_id=1500) + assert operator.run_id is None + assert operator.get_openlineage_facets_on_complete(MagicMock()) == OperatorLineage()