kaxil commented on code in PR #76918:
URL: https://github.com/apache/airflow/pull/76918#discussion_r4540753421
##########
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:
On 3.3+ this makes durable the default for dbt, so should `operators.rst`
get a durable execution section like the glue, databricks and snowflake pages
have? The case I'd want spelled out is Clear Run on a whole Dag run. Per
`resumable-tasks.rst` that path still keeps task state (#72929), so after
fixing upstream data and clearing the run, this task finds the stored run in
`SUCCESS` and returns without rebuilding anything. Only clearing individual
task instances discards the run id.
##########
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:
The reconnect case never asserts that `task_store.set` wasn't called, which
the sync test does. Could this be split into two tests so both assertions are
unconditional?
##########
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:
Could you add a 5xx case that asserts the operator raises and that
`trigger_job_run` isn't called? Nothing pins the `== 404` in `_is_not_found`
today. Widening it to `>= 400` keeps this file green, and then a 500 on the
status lookup would start the duplicate run this PR is meant to prevent.
--
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]
