amoghrajesh commented on code in PR #76918:
URL: https://github.com/apache/airflow/pull/76918#discussion_r4546400603
##########
providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py:
##########
@@ -1058,6 +1063,150 @@ def test_run_job_operator_link(
)
+STORED_RUN_ID = 9999
+STORED_RUN_URL = EXPECTED_JOB_RUN_OP_EXTRA_LINK.format(
+ account_id=DEFAULT_ACCOUNT_ID, project_id=PROJECT_ID, run_id=STORED_RUN_ID
+)
+STORED_RUN_API_URL =
f"https://cloud.getdbt.com/api/v2/accounts/{DEFAULT_ACCOUNT_ID}/runs/{STORED_RUN_ID}/";
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
[email protected](DbtCloudHook, "wait_for_job_run_status", autospec=True,
return_value=True)
[email protected](
+ DbtCloudHook,
+ "trigger_job_run",
+ autospec=True,
+ return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE),
+)
+class TestDbtCloudRunJobOperatorDurable:
+ def setup_method(self):
+ self.dag = DAG("test_dbt_cloud_job_run_op_durable", schedule=None,
start_date=DEFAULT_DATE)
+ self.ti = MagicMock(stats_tags={}, try_number=1)
+ self.task_store = MagicMock(spec_set=["get", "set"])
+ self.task_store.get.return_value = None
+ self.context = {"ti": self.ti, "task_state_store": self.task_store}
+
+ def _operator(self, **kwargs):
+ return DbtCloudRunJobOperator(
+ task_id=TASK_ID, dbt_cloud_conn_id=ACCOUNT_ID_CONN, job_id=JOB_ID,
dag=self.dag, **kwargs
+ )
+
+ def _assert_run_xcoms(self, run_id, run_url):
+ assert self.ti.xcom_push.call_args_list == [
+ call(key="job_run_url", value=run_url),
+ call(key="job_run_id", value=run_id),
+ ]
+
+ def test_fresh_submit_persists_run_id(self, mock_trigger, mock_wait):
+ assert self._operator().execute(self.context) == RUN_ID
+
+ mock_trigger.assert_called_once()
+ mock_wait.assert_called_once()
+ self.task_store.set.assert_called_once_with("dbt_cloud_job_run_id",
RUN_ID)
+ self._assert_run_xcoms(RUN_ID,
DEFAULT_ACCOUNT_JOB_RUN_RESPONSE["data"]["href"])
+
+ @pytest.mark.parametrize(
+ ("status", "waits"),
+ [
+ (DbtCloudJobRunStatus.QUEUED, True),
+ (DbtCloudJobRunStatus.STARTING, True),
+ (DbtCloudJobRunStatus.RUNNING, True),
+ (DbtCloudJobRunStatus.SUCCESS, False),
+ ],
+ )
+ def test_stored_run_is_reused(self, mock_trigger, mock_wait,
requests_mock, status, waits):
+ self.task_store.get.return_value = STORED_RUN_ID
+ requests_mock.get(
+ STORED_RUN_API_URL,
+ json={"data": {"id": STORED_RUN_ID, "status": status.value,
"href": STORED_RUN_URL}},
+ )
+ operator = self._operator()
+
+ assert operator.execute(self.context) == STORED_RUN_ID
+
+ mock_trigger.assert_not_called()
+ assert mock_wait.called is waits
+ self.task_store.set.assert_not_called()
+ assert operator.run_id == STORED_RUN_ID
+ self._assert_run_xcoms(STORED_RUN_ID, STORED_RUN_URL)
+
+ @pytest.mark.parametrize(
+ "response",
+ [
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.ERROR.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.CANCELLED.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {"status_code": 404, "reason": "Not Found"},
+ ],
+ ids=["error", "cancelled", "not_found"],
+ )
+ def test_stored_run_not_reusable_submits_fresh(self, mock_trigger,
mock_wait, requests_mock, response):
+ self.task_store.get.return_value = STORED_RUN_ID
+ requests_mock.get(STORED_RUN_API_URL, **response)
+
+ assert self._operator().execute(self.context) == RUN_ID
+
+ mock_trigger.assert_called_once()
+ self.task_store.set.assert_called_once_with("dbt_cloud_job_run_id",
RUN_ID)
+ self._assert_run_xcoms(RUN_ID,
DEFAULT_ACCOUNT_JOB_RUN_RESPONSE["data"]["href"])
+
+ @pytest.mark.parametrize(
+ ("stored_run_id", "expected_run_id"),
+ [(None, RUN_ID), (STORED_RUN_ID, STORED_RUN_ID)],
+ ids=["fresh", "reconnect"],
+ )
+ @patch.object(DbtCloudRunJobOperator, "defer", autospec=True)
+ @patch.object(
+ DbtCloudHook, "get_job_run_status", autospec=True,
return_value=DbtCloudJobRunStatus.RUNNING.value
+ )
+ def test_deferrable_defers_on_the_durable_run(
+ self, mock_status, mock_defer, mock_trigger, mock_wait, requests_mock,
stored_run_id, expected_run_id
+ ):
+ self.task_store.get.return_value = stored_run_id
+ requests_mock.get(
+ STORED_RUN_API_URL,
+ json={
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.RUNNING.value,
+ "href": STORED_RUN_URL,
+ }
+ },
+ )
+
+ self._operator(deferrable=True).execute(self.context)
+
+ assert mock_trigger.called is (stored_run_id is None)
+ mock_wait.assert_not_called()
+ assert mock_defer.call_args.kwargs["trigger"].run_id == expected_run_id
+ if stored_run_id is None:
Review Comment:
Handled in [review comments on dbt cloud
run](https://github.com/apache/airflow/pull/76918/commits/f472198725f7dd951c9e9c2f8bfe052b0f45e67b)
##########
providers/dbt/cloud/tests/unit/dbt/cloud/operators/test_dbt.py:
##########
@@ -1058,6 +1063,150 @@ def test_run_job_operator_link(
)
+STORED_RUN_ID = 9999
+STORED_RUN_URL = EXPECTED_JOB_RUN_OP_EXTRA_LINK.format(
+ account_id=DEFAULT_ACCOUNT_ID, project_id=PROJECT_ID, run_id=STORED_RUN_ID
+)
+STORED_RUN_API_URL =
f"https://cloud.getdbt.com/api/v2/accounts/{DEFAULT_ACCOUNT_ID}/runs/{STORED_RUN_ID}/";
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
[email protected](DbtCloudHook, "wait_for_job_run_status", autospec=True,
return_value=True)
[email protected](
+ DbtCloudHook,
+ "trigger_job_run",
+ autospec=True,
+ return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE),
+)
+class TestDbtCloudRunJobOperatorDurable:
+ def setup_method(self):
+ self.dag = DAG("test_dbt_cloud_job_run_op_durable", schedule=None,
start_date=DEFAULT_DATE)
+ self.ti = MagicMock(stats_tags={}, try_number=1)
+ self.task_store = MagicMock(spec_set=["get", "set"])
+ self.task_store.get.return_value = None
+ self.context = {"ti": self.ti, "task_state_store": self.task_store}
+
+ def _operator(self, **kwargs):
+ return DbtCloudRunJobOperator(
+ task_id=TASK_ID, dbt_cloud_conn_id=ACCOUNT_ID_CONN, job_id=JOB_ID,
dag=self.dag, **kwargs
+ )
+
+ def _assert_run_xcoms(self, run_id, run_url):
+ assert self.ti.xcom_push.call_args_list == [
+ call(key="job_run_url", value=run_url),
+ call(key="job_run_id", value=run_id),
+ ]
+
+ def test_fresh_submit_persists_run_id(self, mock_trigger, mock_wait):
+ assert self._operator().execute(self.context) == RUN_ID
+
+ mock_trigger.assert_called_once()
+ mock_wait.assert_called_once()
+ self.task_store.set.assert_called_once_with("dbt_cloud_job_run_id",
RUN_ID)
+ self._assert_run_xcoms(RUN_ID,
DEFAULT_ACCOUNT_JOB_RUN_RESPONSE["data"]["href"])
+
+ @pytest.mark.parametrize(
+ ("status", "waits"),
+ [
+ (DbtCloudJobRunStatus.QUEUED, True),
+ (DbtCloudJobRunStatus.STARTING, True),
+ (DbtCloudJobRunStatus.RUNNING, True),
+ (DbtCloudJobRunStatus.SUCCESS, False),
+ ],
+ )
+ def test_stored_run_is_reused(self, mock_trigger, mock_wait,
requests_mock, status, waits):
+ self.task_store.get.return_value = STORED_RUN_ID
+ requests_mock.get(
+ STORED_RUN_API_URL,
+ json={"data": {"id": STORED_RUN_ID, "status": status.value,
"href": STORED_RUN_URL}},
+ )
+ operator = self._operator()
+
+ assert operator.execute(self.context) == STORED_RUN_ID
+
+ mock_trigger.assert_not_called()
+ assert mock_wait.called is waits
+ self.task_store.set.assert_not_called()
+ assert operator.run_id == STORED_RUN_ID
+ self._assert_run_xcoms(STORED_RUN_ID, STORED_RUN_URL)
+
+ @pytest.mark.parametrize(
+ "response",
+ [
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.ERROR.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {
+ "json": {
+ "data": {
+ "id": STORED_RUN_ID,
+ "status": DbtCloudJobRunStatus.CANCELLED.value,
+ "href": STORED_RUN_URL,
+ }
+ }
+ },
+ {"status_code": 404, "reason": "Not Found"},
Review Comment:
Handled in [review comments on dbt cloud
run](https://github.com/apache/airflow/pull/76918/commits/f472198725f7dd951c9e9c2f8bfe052b0f45e67b)
##########
providers/dbt/cloud/src/airflow/providers/dbt/cloud/operators/dbt.py:
##########
@@ -94,6 +114,9 @@ class DbtCloudRunJobOperator(BaseOperator):
reading dbt Cloud runs can parse this to link the dbt Cloud run back
to the Airflow task that
triggered it. Note this overwrites the ``trigger_reason``. Defaults to
the
``openlineage.spark_inject_parent_job_info`` config value.
+ :param durable: When ``True`` (the default), the dbt Cloud run id is
persisted to task state before
Review Comment:
Handled in [review comments on dbt cloud
run](https://github.com/apache/airflow/pull/76918/commits/f472198725f7dd951c9e9c2f8bfe052b0f45e67b)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
