From eed00769340a2195cb8ac1c45fc6e8ce89f71ec3 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Thu, 9 Jan 2025 12:08:28 +0100 Subject: [PATCH 01/10] Update DbtCloudRunJobOperator and corresponding tests --- .../airflow/providers/dbt/cloud/hooks/dbt.py | 98 +++++++++++- .../providers/dbt/cloud/operators/dbt.py | 29 +++- providers/tests/dbt/cloud/hooks/test_dbt.py | 141 +++++++++++++++++- .../tests/dbt/cloud/operators/test_dbt.py | 101 ++++++++++++- 4 files changed, 361 insertions(+), 8 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index 815f6779dbeaa..fe8fb1f11e782 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -356,14 +356,23 @@ def get_account(self, account_id: int | None = None) -> Response: return self._run_and_get_response(endpoint=f"{account_id}/") @fallback_to_default_account - def list_projects(self, account_id: int | None = None) -> list[Response]: + def list_projects( + self, name_contains: str | None = None, account_id: int | None = None + ) -> list[Response]: """ Retrieve metadata for all projects tied to a specified dbt Cloud account. + :param name_contains: Optional. The case-insensitive substring of a dbt Cloud project name to filter by. :param account_id: Optional. The ID of a dbt Cloud account. :return: List of request responses. """ - return self._run_and_get_response(endpoint=f"{account_id}/projects/", paginate=True, api_version="v3") + payload = {"name__icontains": name_contains} if name_contains else None + return self._run_and_get_response( + endpoint=f"{account_id}/projects/", + payload=payload, + paginate=True, + api_version="v3", + ) @fallback_to_default_account def get_project(self, project_id: int, account_id: int | None = None) -> Response: @@ -376,27 +385,64 @@ def get_project(self, project_id: int, account_id: int | None = None) -> Respons """ return self._run_and_get_response(endpoint=f"{account_id}/projects/{project_id}/", api_version="v3") + @fallback_to_default_account + def list_environments(self, project_id: int, account_id: int | None = None) -> list[Response]: + """ + Retrieve metadata for all environments tied to a specified dbt Cloud project. + + :param project_id: The ID of a dbt Cloud project. + :param account_id: Optional. The ID of a dbt Cloud account. + :return: List of request responses. + """ + return self._run_and_get_response( + endpoint=f"{account_id}/projects/{project_id}/environments/", + paginate=True, + api_version="v3", + ) + + @fallback_to_default_account + def get_environment( + self, project_id: int, environment_id: int, account_id: int | None = None + ) -> Response: + """ + Retrieve metadata for a specific project's environment. + + :param project_id: The ID of a dbt Cloud project. + :param environment_id: The ID of a dbt Cloud environment. + :param account_id: Optional. The ID of a dbt Cloud account. + :return: The request response. + """ + return self._run_and_get_response( + endpoint=f"{account_id}/projects/{project_id}/environments/{environment_id}/", api_version="v3" + ) + @fallback_to_default_account def list_jobs( self, account_id: int | None = None, order_by: str | None = None, project_id: int | None = None, + environment_id: int | None = None, ) -> list[Response]: """ Retrieve metadata for all jobs tied to a specified dbt Cloud account. If a ``project_id`` is supplied, only jobs pertaining to this project will be retrieved. + If an ``environment_id`` is supplied, only jobs pertaining to this environment will be retrieved. :param account_id: Optional. The ID of a dbt Cloud account. :param order_by: Optional. Field to order the result by. Use '-' to indicate reverse order. For example, to use reverse order by the run ID use ``order_by=-id``. - :param project_id: The ID of a dbt Cloud project. + :param project_id: Optional. The ID of a dbt Cloud project. + :param environment_id: Optional. The ID of a dbt Cloud environment. :return: List of request responses. """ + payload = {"order_by": order_by, "project_id": project_id} + if environment_id: + payload["environment_id"] = environment_id return self._run_and_get_response( endpoint=f"{account_id}/jobs/", - payload={"order_by": order_by, "project_id": project_id}, + payload=payload, paginate=True, ) @@ -411,6 +457,50 @@ def get_job(self, job_id: int, account_id: int | None = None) -> Response: """ return self._run_and_get_response(endpoint=f"{account_id}/jobs/{job_id}") + @fallback_to_default_account + def get_job_by_name( + self, project_name: str, environment_name: str, job_name: str, account_id: int | None = None + ) -> Response: + """ + Retrieve metadata for a specific job by combination of project, environment, and job name. + + Raises AirflowException if the job is not found or cannot be uniquely identified by provided parameters. + + :param project_name: The name of a dbt Cloud project. + :param environment_name: The name of a dbt Cloud environment. + :param job_name: The name of a dbt Cloud job. + :param account_id: Optional. The ID of a dbt Cloud account. + :return: The request response. + """ + # get project_id using project_name + projects = self.list_projects(name_contains=project_name, account_id=account_id)[0].json()["data"] + projects = [project for project in projects if project["name"] == project_name] + if len(projects) != 1: + raise AirflowException(f"Found {len(projects)} projects with name `{project_name}`.") + project_id = projects[0]["id"] + + # get environment_id using project_id and environment_name + environments = self.list_environments(project_id=project_id, account_id=account_id)[0].json()["data"] + environments = [env for env in environments if env["name"] == environment_name] + if len(environments) != 1: + raise AirflowException( + f"Found {len(environments)} environments with name `{environment_name}` in project `{project_name}`." + ) + environment_id = environments[0]["id"] + + # get job using project_id, environment_id and job_name + list_jobs_responses = self.list_jobs( + project_id=project_id, environment_id=environment_id, account_id=account_id + ) + jobs = list_jobs_responses[0].json()["data"] + jobs = [job for job in jobs if job["name"] == job_name] + if len(jobs) != 1: + raise AirflowException( + f"Found {len(jobs)} jobs with name `{job_name}` in project `{project_name}` and environment `{environment_name}`." + ) + + return list_jobs_responses[0] + @fallback_to_default_account def trigger_job_run( self, diff --git a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py index 8795ebf0ca714..65aa8b1ab784f 100644 --- a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Any from airflow.configuration import conf +from airflow.exceptions import AirflowException from airflow.models import BaseOperator, BaseOperatorLink, XCom from airflow.providers.dbt.cloud.hooks.dbt import ( DbtCloudHook, @@ -57,7 +58,10 @@ class DbtCloudRunJobOperator(BaseOperator): :ref:`howto/operator:DbtCloudRunJobOperator` :param dbt_cloud_conn_id: The connection ID for connecting to dbt Cloud. - :param job_id: The ID of a dbt Cloud job. + :param job_id: The ID of a dbt Cloud job. Required if project_name, environment_name, and job_name are not provided. + :param project_name: Optional. The name of a dbt Cloud project. Used only if ``job_id`` is None. + :param environment_name: Optional. The name of a dbt Cloud environment. Used only if ``job_id`` is None. + :param job_name: Optional. The name of a dbt Cloud job. Used only if ``job_id`` is None. :param account_id: Optional. The ID of a dbt Cloud account. :param trigger_reason: Optional. Description of the reason to trigger the job. Defaults to "Triggered via Apache Airflow by task in the DAG." @@ -86,6 +90,9 @@ class DbtCloudRunJobOperator(BaseOperator): template_fields = ( "dbt_cloud_conn_id", "job_id", + "project_name", + "environment_name", + "job_name", "account_id", "trigger_reason", "steps_override", @@ -99,7 +106,10 @@ def __init__( self, *, dbt_cloud_conn_id: str = DbtCloudHook.default_conn_name, - job_id: int, + job_id: int | None = None, + project_name: str | None = None, + environment_name: str | None = None, + job_name: str | None = None, account_id: int | None = None, trigger_reason: str | None = None, steps_override: list[str] | None = None, @@ -117,6 +127,9 @@ def __init__( self.dbt_cloud_conn_id = dbt_cloud_conn_id self.account_id = account_id self.job_id = job_id + self.project_name = project_name + self.environment_name = environment_name + self.job_name = job_name self.trigger_reason = trigger_reason self.steps_override = steps_override self.schema_override = schema_override @@ -135,6 +148,18 @@ def execute(self, context: Context): f"Triggered via Apache Airflow by task {self.task_id!r} in the {self.dag.dag_id} DAG." ) + if self.job_id is None: + if not all([self.project_name, self.environment_name, self.job_name]): + raise AirflowException( + "Either job_id or project_name, environment_name, and job_name must be provided." + ) + self.job_id = self.hook.get_job_by_name( + account_id=self.account_id, + project_name=self.project_name, + environment_name=self.environment_name, + job_name=self.job_name, + ).json()["data"]["id"] + non_terminal_runs = None if self.reuse_existing_run: non_terminal_runs = self.hook.get_job_runs( diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index 590f1b677f10c..930c3a1a81648 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -19,9 +19,10 @@ import json from datetime import timedelta from typing import Any -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest +from requests.models import Response from airflow.exceptions import AirflowException from airflow.models.connection import Connection @@ -47,12 +48,48 @@ EXTRA_PROXIES = {"proxies": {"https": "http://myproxy:1234"}} TOKEN = "token" PROJECT_ID = 33333 +PROJECT_NAME = "project_name" +ENVIRONMENT_ID = 44444 +ENVIRONMENT_NAME = "environment_name" JOB_ID = 4444 +JOB_NAME = "job_name" RUN_ID = 5555 BASE_URL = "https://cloud.getdbt.com/" SINGLE_TENANT_URL = "https://single.tenant.getdbt.com/" +DEFAULT_LIST_PROJECTS_RESPONSE = { + "data": [ + { + "id": PROJECT_ID, + "name": PROJECT_NAME, + } + ] +} +DEFAULT_LIST_ENVIRONMENTS_RESPONSE = { + "data": [ + { + "id": ENVIRONMENT_ID, + "name": ENVIRONMENT_NAME, + } + ] +} +DEFAULT_LIST_JOBS_RESPONSE = { + "data": [ + { + "id": JOB_ID, + "name": JOB_NAME, + } + ] +} +EMPTY_RESPONSE = {"data": []} + + +def mock_response_json(response: dict): + run_response = MagicMock(**response, spec=Response) + run_response.json.return_value = response + return run_response + class TestDbtCloudJobRunStatus: valid_job_run_statuses = [ @@ -269,6 +306,48 @@ def test_get_project(self, mock_http_run, mock_paginate, conn_id, account_id): ) hook._paginate.assert_not_called() + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "run") + @patch.object(DbtCloudHook, "_paginate") + def test_list_environments(self, mock_http_run, mock_paginate, conn_id, account_id): + hook = DbtCloudHook(conn_id) + hook.list_environments(project_id=PROJECT_ID, account_id=account_id) + + assert hook.method == "GET" + + _account_id = account_id or DEFAULT_ACCOUNT_ID + hook.run.assert_not_called() + hook._paginate.assert_called_once_with( + endpoint=f"api/v3/accounts/{_account_id}/projects/{PROJECT_ID}/environments/", + payload=None, + proxies=None, + ) + + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "run") + @patch.object(DbtCloudHook, "_paginate") + def test_get_environment(self, mock_http_run, mock_paginate, conn_id, account_id): + hook = DbtCloudHook(conn_id) + hook.get_environment(project_id=PROJECT_ID, environment_id=ENVIRONMENT_ID, account_id=account_id) + + assert hook.method == "GET" + + _account_id = account_id or DEFAULT_ACCOUNT_ID + hook.run.assert_called_once_with( + endpoint=f"api/v3/accounts/{_account_id}/projects/{PROJECT_ID}/environments/{ENVIRONMENT_ID}/", + data=None, + extra_options=None, + ) + hook._paginate.assert_not_called() + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], @@ -330,6 +409,66 @@ def test_get_job(self, mock_http_run, mock_paginate, conn_id, account_id): ) hook._paginate.assert_not_called() + @patch.object(DbtCloudHook, "list_jobs", return_value=[mock_response_json(DEFAULT_LIST_JOBS_RESPONSE)]) + @patch.object( + DbtCloudHook, + "list_environments", + return_value=[mock_response_json(DEFAULT_LIST_ENVIRONMENTS_RESPONSE)], + ) + @patch.object( + DbtCloudHook, "list_projects", return_value=[mock_response_json(DEFAULT_LIST_PROJECTS_RESPONSE)] + ) + def test_get_job_by_name_returns_response( + self, mock_list_projects, mock_list_environments, mock_list_jobs + ): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + response = hook.get_job_by_name( + project_name=PROJECT_NAME, + environment_name=ENVIRONMENT_NAME, + job_name=JOB_NAME, + account_id=None, + ) + + assert isinstance(response, Response) + + @pytest.mark.parametrize( + argnames="project_name, environment_name, job_name", + argvalues=[ + ("dummy_name", ENVIRONMENT_NAME, JOB_NAME), + (PROJECT_NAME, "dummy_name", JOB_NAME), + (PROJECT_NAME, ENVIRONMENT_NAME, JOB_NAME.upper()), + (None, ENVIRONMENT_NAME, JOB_NAME), + (PROJECT_NAME, "", JOB_NAME), + ("", "", ""), + ], + ) + @patch.object(DbtCloudHook, "list_jobs", return_value=[mock_response_json(EMPTY_RESPONSE)]) + @patch.object( + DbtCloudHook, + "list_environments", + return_value=[mock_response_json(DEFAULT_LIST_ENVIRONMENTS_RESPONSE)], + ) + @patch.object( + DbtCloudHook, "list_projects", return_value=[mock_response_json(DEFAULT_LIST_PROJECTS_RESPONSE)] + ) + def test_get_job_by_incorrect_name_raises_exception( + self, + mock_list_projects, + mock_list_environments, + mock_list_jobs, + project_name, + environment_name, + job_name, + ): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + with pytest.raises(AirflowException, match="Found 0"): + hook.get_job_by_name( + project_name=project_name, + environment_name=environment_name, + job_name=job_name, + account_id=None, + ) + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], diff --git a/providers/tests/dbt/cloud/operators/test_dbt.py b/providers/tests/dbt/cloud/operators/test_dbt.py index 315e770d84dbd..05e269548c24f 100644 --- a/providers/tests/dbt/cloud/operators/test_dbt.py +++ b/providers/tests/dbt/cloud/operators/test_dbt.py @@ -22,7 +22,7 @@ import pytest -from airflow.exceptions import TaskDeferred +from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import DAG, Connection from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus from airflow.providers.dbt.cloud.operators.dbt import ( @@ -43,7 +43,11 @@ ACCOUNT_ID = 22222 TOKEN = "token" PROJECT_ID = 33333 +PROJECT_NAME = "project_name" +ENVIRONMENT_ID = 44444 +ENVIRONMENT_NAME = "environment_name" JOB_ID = 4444 +JOB_NAME = "job_name" RUN_ID = 5555 EXPECTED_JOB_RUN_OP_EXTRA_LINK = ( "https://cloud.getdbt.com/#/accounts/{account_id}/projects/{project_id}/runs/{run_id}/" @@ -75,6 +79,12 @@ } ] } +DEFAULT_ACCOUNT_JOB_RESPONSE = { + "data": { + "id": JOB_ID, + "account_id": DEFAULT_ACCOUNT_ID, + } +} def mock_response_json(response: dict): @@ -200,6 +210,95 @@ 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( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_by_name", + return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RESPONSE), + ) + @patch( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_run_status", + return_value=DbtCloudJobRunStatus.SUCCESS.value, + ) + @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_connection") + @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_by_name( + self, mock_trigger_job_run, mock_dbt_hook, mock_job_run_status, mock_job_by_name + ): + """ + Test alternative way to run a job by project, + environment and job name instead of job id. + """ + dbt_op = DbtCloudRunJobOperator( + dbt_cloud_conn_id=ACCOUNT_ID_CONN, + task_id=TASK_ID, + project_name=PROJECT_NAME, + environment_name=ENVIRONMENT_NAME, + job_name=JOB_NAME, + check_interval=1, + timeout=3, + dag=self.dag, + ) + dbt_op.execute(MagicMock()) + mock_trigger_job_run.assert_called_once() + + @pytest.mark.parametrize( + argnames="project_name, environment_name, job_name", + argvalues=[ + (None, ENVIRONMENT_NAME, JOB_NAME), + (PROJECT_NAME, "", JOB_NAME), + (PROJECT_NAME, ENVIRONMENT_NAME, None), + ("", "", ""), + ], + ) + @patch( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_by_name", + return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RESPONSE), + ) + @patch( + "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_run_status", + return_value=DbtCloudJobRunStatus.SUCCESS.value, + ) + @patch("airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_connection") + @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_by_incorrect_name_raises_exception( + self, + mock_trigger_job_run, + mock_dbt_hook, + mock_job_run_status, + mock_job_by_name, + project_name, + environment_name, + job_name, + ): + """ + Test alternative way to run a job by project, + environment and job name instead of job id. + + This test is to check if the operator raises an exception + when the project, environment or job name is missing. + """ + dbt_op = DbtCloudRunJobOperator( + dbt_cloud_conn_id=ACCOUNT_ID_CONN, + task_id=TASK_ID, + project_name=project_name, + environment_name=environment_name, + job_name=job_name, + check_interval=1, + timeout=3, + dag=self.dag, + ) + with pytest.raises( + AirflowException, + match="Either job_id or project_name, environment_name, and job_name must be provided.", + ): + dbt_op.execute(MagicMock()) + mock_trigger_job_run.assert_not_called() + @patch.object( DbtCloudHook, "trigger_job_run", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE) ) From e50d6f3d5e52e0048562cd8872d33a807869f6e8 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Thu, 9 Jan 2025 14:54:47 +0100 Subject: [PATCH 02/10] update documentation --- docs/apache-airflow-providers-dbt-cloud/operators.rst | 11 +++++++++++ providers/tests/system/dbt/cloud/example_dbt_cloud.py | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/docs/apache-airflow-providers-dbt-cloud/operators.rst b/docs/apache-airflow-providers-dbt-cloud/operators.rst index eaa285f6d4082..8d069bb869bcd 100644 --- a/docs/apache-airflow-providers-dbt-cloud/operators.rst +++ b/docs/apache-airflow-providers-dbt-cloud/operators.rst @@ -82,6 +82,17 @@ via the ``additional_run_config`` dictionary. :start-after: [START howto_operator_dbt_cloud_run_job_async] :end-before: [END howto_operator_dbt_cloud_run_job_async] +You can also trigger a dbt Cloud job without providing the ``job_id``. Instead, you can identify the job +by providing the ``project_name``, ``environment_name``, and ``job_name``. +Please note that it will only work if the above three parameters uniquely identify a job in your account +(i.e. you cannot have two jobs with the same name in the same project and environment). + +.. exampleinclude:: /../../providers/tests/system/dbt/cloud/example_dbt_cloud.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_dbt_cloud_run_job_without_job_id] + :end-before: [END howto_operator_dbt_cloud_run_job_without_job_id] + .. _howto/operator:DbtCloudJobRunSensor: Poll for status of a dbt Cloud Job run diff --git a/providers/tests/system/dbt/cloud/example_dbt_cloud.py b/providers/tests/system/dbt/cloud/example_dbt_cloud.py index b36f7439fd698..603b3100b9891 100644 --- a/providers/tests/system/dbt/cloud/example_dbt_cloud.py +++ b/providers/tests/system/dbt/cloud/example_dbt_cloud.py @@ -67,6 +67,17 @@ ) # [END howto_operator_dbt_cloud_run_job_async] + # [START howto_operator_dbt_cloud_run_job_without_job_id] + trigger_job_run3 = DbtCloudRunJobOperator( + task_id="trigger_job_run3", + project_name="my_dbt_project", + environment_name="prod", + job_name="my_dbt_job", + check_interval=10, + timeout=300, + ) + # [END howto_operator_dbt_cloud_run_job_without_job_id] + # [START howto_operator_dbt_cloud_run_job_sensor] job_run_sensor = DbtCloudJobRunSensor( task_id="job_run_sensor", run_id=trigger_job_run2.output, timeout=20 From 196b08507d82cb86e614324b499a734fbcb28354 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Fri, 10 Jan 2025 17:27:56 +0100 Subject: [PATCH 03/10] Fixes after manual tests --- providers/src/airflow/providers/dbt/cloud/hooks/dbt.py | 6 +++--- providers/src/airflow/providers/dbt/cloud/operators/dbt.py | 2 +- providers/tests/dbt/cloud/hooks/test_dbt.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index fe8fb1f11e782..b133f7575feba 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -460,7 +460,7 @@ def get_job(self, job_id: int, account_id: int | None = None) -> Response: @fallback_to_default_account def get_job_by_name( self, project_name: str, environment_name: str, job_name: str, account_id: int | None = None - ) -> Response: + ) -> dict: """ Retrieve metadata for a specific job by combination of project, environment, and job name. @@ -470,7 +470,7 @@ def get_job_by_name( :param environment_name: The name of a dbt Cloud environment. :param job_name: The name of a dbt Cloud job. :param account_id: Optional. The ID of a dbt Cloud account. - :return: The request response. + :return: The details of a job. """ # get project_id using project_name projects = self.list_projects(name_contains=project_name, account_id=account_id)[0].json()["data"] @@ -499,7 +499,7 @@ def get_job_by_name( f"Found {len(jobs)} jobs with name `{job_name}` in project `{project_name}` and environment `{environment_name}`." ) - return list_jobs_responses[0] + return jobs[0] @fallback_to_default_account def trigger_job_run( diff --git a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py index 65aa8b1ab784f..9429f24118886 100644 --- a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py @@ -158,7 +158,7 @@ def execute(self, context: Context): project_name=self.project_name, environment_name=self.environment_name, job_name=self.job_name, - ).json()["data"]["id"] + )["id"] non_terminal_runs = None if self.reuse_existing_run: diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index 930c3a1a81648..a75d1c9b6672a 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -422,14 +422,14 @@ def test_get_job_by_name_returns_response( self, mock_list_projects, mock_list_environments, mock_list_jobs ): hook = DbtCloudHook(ACCOUNT_ID_CONN) - response = hook.get_job_by_name( + job_details = hook.get_job_by_name( project_name=PROJECT_NAME, environment_name=ENVIRONMENT_NAME, job_name=JOB_NAME, account_id=None, ) - assert isinstance(response, Response) + assert job_details == DEFAULT_LIST_JOBS_RESPONSE["data"][0] @pytest.mark.parametrize( argnames="project_name, environment_name, job_name", From a328dcc8e3687847cc722f8ba09ea3f9304a38dc Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Wed, 15 Jan 2025 11:33:57 +0100 Subject: [PATCH 04/10] add unit test for duplicate names --- providers/tests/dbt/cloud/hooks/test_dbt.py | 56 ++++++++++++++++++++- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index a75d1c9b6672a..31a98f9ac0e59 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -17,6 +17,7 @@ from __future__ import annotations import json +from copy import deepcopy from datetime import timedelta from typing import Any from unittest.mock import MagicMock, patch @@ -82,7 +83,6 @@ } ] } -EMPTY_RESPONSE = {"data": []} def mock_response_json(response: dict): @@ -442,7 +442,7 @@ def test_get_job_by_name_returns_response( ("", "", ""), ], ) - @patch.object(DbtCloudHook, "list_jobs", return_value=[mock_response_json(EMPTY_RESPONSE)]) + @patch.object(DbtCloudHook, "list_jobs", return_value=[mock_response_json(DEFAULT_LIST_JOBS_RESPONSE)]) @patch.object( DbtCloudHook, "list_environments", @@ -469,6 +469,58 @@ def test_get_job_by_incorrect_name_raises_exception( account_id=None, ) + @pytest.mark.parametrize("duplicated", ["projects", "environments", "jobs"]) + def test_get_job_by_duplicate_name_raises_exception(self, duplicated): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + mock_list_jobs_response = deepcopy(DEFAULT_LIST_JOBS_RESPONSE) + mock_list_environments_response = deepcopy(DEFAULT_LIST_ENVIRONMENTS_RESPONSE) + mock_list_projects_response = deepcopy(DEFAULT_LIST_PROJECTS_RESPONSE) + + if duplicated == "projects": + mock_list_projects_response["data"].append( + { + "id": PROJECT_ID + 1, + "name": PROJECT_NAME, + } + ) + elif duplicated == "environments": + mock_list_environments_response["data"].append( + { + "id": ENVIRONMENT_ID + 1, + "name": ENVIRONMENT_NAME, + } + ) + elif duplicated == "jobs": + mock_list_jobs_response["data"].append( + { + "id": JOB_ID + 1, + "name": JOB_NAME, + } + ) + + with ( + patch.object( + DbtCloudHook, "list_jobs", return_value=[mock_response_json(mock_list_jobs_response)] + ), + patch.object( + DbtCloudHook, + "list_environments", + return_value=[mock_response_json(mock_list_environments_response)], + ), + patch.object( + DbtCloudHook, + "list_projects", + return_value=[mock_response_json(mock_list_projects_response)], + ), + ): + with pytest.raises(AirflowException, match=f"Found 2 {duplicated}"): + hook.get_job_by_name( + project_name=PROJECT_NAME, + environment_name=ENVIRONMENT_NAME, + job_name=JOB_NAME, + account_id=None, + ) + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], From e9765e356244be465afdb181c50a4c68696d0cbd Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Thu, 16 Jan 2025 15:49:39 +0100 Subject: [PATCH 05/10] Enhance DbtCloudHook to support name filtering in list_environments and list_jobs methods --- .../airflow/providers/dbt/cloud/hooks/dbt.py | 47 +++++++++++--- providers/tests/dbt/cloud/hooks/test_dbt.py | 65 ++++++++++++++++++- 2 files changed, 101 insertions(+), 11 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index b133f7575feba..d06c3e0fd338c 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -386,16 +386,21 @@ def get_project(self, project_id: int, account_id: int | None = None) -> Respons return self._run_and_get_response(endpoint=f"{account_id}/projects/{project_id}/", api_version="v3") @fallback_to_default_account - def list_environments(self, project_id: int, account_id: int | None = None) -> list[Response]: + def list_environments( + self, project_id: int, name_contains: str | None = None, account_id: int | None = None + ) -> list[Response]: """ Retrieve metadata for all environments tied to a specified dbt Cloud project. :param project_id: The ID of a dbt Cloud project. + :param name_contains: Optional. The case-insensitive substring of a dbt Cloud environment name to filter by. :param account_id: Optional. The ID of a dbt Cloud account. :return: List of request responses. """ + payload = {"name__icontains": name_contains} if name_contains else None return self._run_and_get_response( endpoint=f"{account_id}/projects/{project_id}/environments/", + payload=payload, paginate=True, api_version="v3", ) @@ -423,6 +428,7 @@ def list_jobs( order_by: str | None = None, project_id: int | None = None, environment_id: int | None = None, + name_contains: str | None = None, ) -> list[Response]: """ Retrieve metadata for all jobs tied to a specified dbt Cloud account. @@ -435,11 +441,14 @@ def list_jobs( For example, to use reverse order by the run ID use ``order_by=-id``. :param project_id: Optional. The ID of a dbt Cloud project. :param environment_id: Optional. The ID of a dbt Cloud environment. + :param name_contains: Optional. The case-insensitive substring of a dbt Cloud job name to filter by. :return: List of request responses. """ payload = {"order_by": order_by, "project_id": project_id} if environment_id: payload["environment_id"] = environment_id + if name_contains: + payload["name__icontains"] = name_contains return self._run_and_get_response( endpoint=f"{account_id}/jobs/", payload=payload, @@ -473,15 +482,29 @@ def get_job_by_name( :return: The details of a job. """ # get project_id using project_name - projects = self.list_projects(name_contains=project_name, account_id=account_id)[0].json()["data"] - projects = [project for project in projects if project["name"] == project_name] + projects = self.list_projects(name_contains=project_name, account_id=account_id) + # flatten & filter the list of responses + projects = [ + project + for response in projects + for project in response.json()["data"] + if project["name"] == project_name + ] if len(projects) != 1: raise AirflowException(f"Found {len(projects)} projects with name `{project_name}`.") project_id = projects[0]["id"] # get environment_id using project_id and environment_name - environments = self.list_environments(project_id=project_id, account_id=account_id)[0].json()["data"] - environments = [env for env in environments if env["name"] == environment_name] + environments = self.list_environments( + project_id=project_id, name_contains=environment_name, account_id=account_id + ) + # flatten & filter the list of responses + environments = [ + env + for response in environments + for env in response.json()["data"] + if env["name"] == environment_name + ] if len(environments) != 1: raise AirflowException( f"Found {len(environments)} environments with name `{environment_name}` in project `{project_name}`." @@ -490,10 +513,18 @@ def get_job_by_name( # get job using project_id, environment_id and job_name list_jobs_responses = self.list_jobs( - project_id=project_id, environment_id=environment_id, account_id=account_id + project_id=project_id, + environment_id=environment_id, + name_contains=job_name, + account_id=account_id, ) - jobs = list_jobs_responses[0].json()["data"] - jobs = [job for job in jobs if job["name"] == job_name] + # flatten & filter the list of responses + jobs = [ + job + for response in list_jobs_responses + for job in response.json()["data"] + if job["name"] == job_name + ] if len(jobs) != 1: raise AirflowException( f"Found {len(jobs)} jobs with name `{job_name}` in project `{project_name}` and environment `{environment_name}`." diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index 31a98f9ac0e59..4ca518dca3c93 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -284,7 +284,32 @@ def test_list_projects(self, mock_http_run, mock_paginate, conn_id, account_id): _account_id = account_id or DEFAULT_ACCOUNT_ID hook.run.assert_not_called() hook._paginate.assert_called_once_with( - endpoint=f"api/v3/accounts/{_account_id}/projects/", payload=None, proxies=None + endpoint=f"api/v3/accounts/{_account_id}/projects/", + payload=None, + proxies=None, + ) + + @pytest.mark.parametrize( + argnames="conn_id, account_id, name_contains", + argvalues=[(ACCOUNT_ID_CONN, None, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID, PROJECT_NAME)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "run") + @patch.object(DbtCloudHook, "_paginate") + def test_list_projects_with_payload( + self, mock_http_run, mock_paginate, conn_id, account_id, name_contains + ): + hook = DbtCloudHook(conn_id) + hook.list_projects(account_id=account_id, name_contains=name_contains) + + assert hook.method == "GET" + + _account_id = account_id or DEFAULT_ACCOUNT_ID + hook.run.assert_not_called() + hook._paginate.assert_called_once_with( + endpoint=f"api/v3/accounts/{_account_id}/projects/", + payload={"name__icontains": name_contains} if name_contains else None, + proxies=None, ) @pytest.mark.parametrize( @@ -327,6 +352,29 @@ def test_list_environments(self, mock_http_run, mock_paginate, conn_id, account_ proxies=None, ) + @pytest.mark.parametrize( + argnames="conn_id, account_id, name_contains", + argvalues=[(ACCOUNT_ID_CONN, None, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID, ENVIRONMENT_NAME)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "run") + @patch.object(DbtCloudHook, "_paginate") + def test_list_environments_with_payload( + self, mock_http_run, mock_paginate, conn_id, account_id, name_contains + ): + hook = DbtCloudHook(conn_id) + hook.list_environments(project_id=PROJECT_ID, account_id=account_id, name_contains=name_contains) + + assert hook.method == "GET" + + _account_id = account_id or DEFAULT_ACCOUNT_ID + hook.run.assert_not_called() + hook._paginate.assert_called_once_with( + endpoint=f"api/v3/accounts/{_account_id}/projects/{PROJECT_ID}/environments/", + payload={"name__icontains": name_contains} if name_contains else None, + proxies=None, + ) + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], @@ -378,14 +426,25 @@ def test_list_jobs(self, mock_http_run, mock_paginate, conn_id, account_id): @patch.object(DbtCloudHook, "_paginate") def test_list_jobs_with_payload(self, mock_http_run, mock_paginate, conn_id, account_id): hook = DbtCloudHook(conn_id) - hook.list_jobs(project_id=PROJECT_ID, account_id=account_id, order_by="-id") + hook.list_jobs( + project_id=PROJECT_ID, + account_id=account_id, + order_by="-id", + environment_id=ENVIRONMENT_ID, + name_contains=JOB_NAME, + ) assert hook.method == "GET" _account_id = account_id or DEFAULT_ACCOUNT_ID hook._paginate.assert_called_once_with( endpoint=f"api/v2/accounts/{_account_id}/jobs/", - payload={"order_by": "-id", "project_id": PROJECT_ID}, + payload={ + "order_by": "-id", + "project_id": PROJECT_ID, + "environment_id": ENVIRONMENT_ID, + "name__icontains": JOB_NAME, + }, proxies=None, ) hook.run.assert_not_called() From 465ce9ad79881a010da0140c03b7bce86a722e2f Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Fri, 17 Jan 2025 08:44:48 +0100 Subject: [PATCH 06/10] Improve clarity by renaming variables and enhancing comments --- .../src/airflow/providers/dbt/cloud/hooks/dbt.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index d06c3e0fd338c..009fac7ea3d35 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -482,11 +482,11 @@ def get_job_by_name( :return: The details of a job. """ # get project_id using project_name - projects = self.list_projects(name_contains=project_name, account_id=account_id) - # flatten & filter the list of responses + list_projects_responses = self.list_projects(name_contains=project_name, account_id=account_id) + # flatten & filter the list of responses to find the exact match projects = [ project - for response in projects + for response in list_projects_responses for project in response.json()["data"] if project["name"] == project_name ] @@ -495,13 +495,13 @@ def get_job_by_name( project_id = projects[0]["id"] # get environment_id using project_id and environment_name - environments = self.list_environments( + list_environments_responses = self.list_environments( project_id=project_id, name_contains=environment_name, account_id=account_id ) - # flatten & filter the list of responses + # flatten & filter the list of responses to find the exact match environments = [ env - for response in environments + for response in list_environments_responses for env in response.json()["data"] if env["name"] == environment_name ] @@ -518,7 +518,7 @@ def get_job_by_name( name_contains=job_name, account_id=account_id, ) - # flatten & filter the list of responses + # flatten & filter the list of responses to find the exact match jobs = [ job for response in list_jobs_responses @@ -527,7 +527,7 @@ def get_job_by_name( ] if len(jobs) != 1: raise AirflowException( - f"Found {len(jobs)} jobs with name `{job_name}` in project `{project_name}` and environment `{environment_name}`." + f"Found {len(jobs)} jobs with name `{job_name}` in environment `{environment_name}` in project `{project_name}`." ) return jobs[0] From a2457f0ab0e91093f2c3800dfd298b4977899284 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Mon, 27 Jan 2025 10:24:01 +0100 Subject: [PATCH 07/10] Refactor dbt Cloud hook and operator error handling & update method signatures to use keyword-only arguments --- .../airflow/providers/dbt/cloud/hooks/dbt.py | 20 +++++++++++-------- .../providers/dbt/cloud/operators/dbt.py | 3 +-- providers/tests/dbt/cloud/hooks/test_dbt.py | 5 +++-- .../tests/dbt/cloud/operators/test_dbt.py | 4 ++-- 4 files changed, 18 insertions(+), 14 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index 009fac7ea3d35..f9bb4e05b0178 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -135,6 +135,10 @@ class DbtCloudJobRunException(AirflowException): """An exception that indicates a job run failed to complete.""" +class DbtCloudResourceLookupError(AirflowException): + """Exception raised when a dbt Cloud resource cannot be uniquely identified.""" + + T = TypeVar("T", bound=Any) @@ -357,13 +361,13 @@ def get_account(self, account_id: int | None = None) -> Response: @fallback_to_default_account def list_projects( - self, name_contains: str | None = None, account_id: int | None = None + self, account_id: int | None = None, name_contains: str | None = None ) -> list[Response]: """ Retrieve metadata for all projects tied to a specified dbt Cloud account. - :param name_contains: Optional. The case-insensitive substring of a dbt Cloud project name to filter by. :param account_id: Optional. The ID of a dbt Cloud account. + :param name_contains: Optional. The case-insensitive substring of a dbt Cloud project name to filter by. :return: List of request responses. """ payload = {"name__icontains": name_contains} if name_contains else None @@ -387,7 +391,7 @@ def get_project(self, project_id: int, account_id: int | None = None) -> Respons @fallback_to_default_account def list_environments( - self, project_id: int, name_contains: str | None = None, account_id: int | None = None + self, project_id: int, *, name_contains: str | None = None, account_id: int | None = None ) -> list[Response]: """ Retrieve metadata for all environments tied to a specified dbt Cloud project. @@ -407,7 +411,7 @@ def list_environments( @fallback_to_default_account def get_environment( - self, project_id: int, environment_id: int, account_id: int | None = None + self, project_id: int, environment_id: int, *, account_id: int | None = None ) -> Response: """ Retrieve metadata for a specific project's environment. @@ -468,7 +472,7 @@ def get_job(self, job_id: int, account_id: int | None = None) -> Response: @fallback_to_default_account def get_job_by_name( - self, project_name: str, environment_name: str, job_name: str, account_id: int | None = None + self, *, project_name: str, environment_name: str, job_name: str, account_id: int | None = None ) -> dict: """ Retrieve metadata for a specific job by combination of project, environment, and job name. @@ -491,7 +495,7 @@ def get_job_by_name( if project["name"] == project_name ] if len(projects) != 1: - raise AirflowException(f"Found {len(projects)} projects with name `{project_name}`.") + raise DbtCloudResourceLookupError(f"Found {len(projects)} projects with name `{project_name}`.") project_id = projects[0]["id"] # get environment_id using project_id and environment_name @@ -506,7 +510,7 @@ def get_job_by_name( if env["name"] == environment_name ] if len(environments) != 1: - raise AirflowException( + raise DbtCloudResourceLookupError( f"Found {len(environments)} environments with name `{environment_name}` in project `{project_name}`." ) environment_id = environments[0]["id"] @@ -526,7 +530,7 @@ def get_job_by_name( if job["name"] == job_name ] if len(jobs) != 1: - raise AirflowException( + raise DbtCloudResourceLookupError( f"Found {len(jobs)} jobs with name `{job_name}` in environment `{environment_name}` in project `{project_name}`." ) diff --git a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py index 9429f24118886..d9aab8e9e2f67 100644 --- a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any from airflow.configuration import conf -from airflow.exceptions import AirflowException from airflow.models import BaseOperator, BaseOperatorLink, XCom from airflow.providers.dbt.cloud.hooks.dbt import ( DbtCloudHook, @@ -150,7 +149,7 @@ def execute(self, context: Context): if self.job_id is None: if not all([self.project_name, self.environment_name, self.job_name]): - raise AirflowException( + raise ValueError( "Either job_id or project_name, environment_name, and job_name must be provided." ) self.job_id = self.hook.get_job_by_name( diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index 4ca518dca3c93..1a11662019b46 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -32,6 +32,7 @@ DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus, + DbtCloudResourceLookupError, TokenAuth, fallback_to_default_account, ) @@ -520,7 +521,7 @@ def test_get_job_by_incorrect_name_raises_exception( job_name, ): hook = DbtCloudHook(ACCOUNT_ID_CONN) - with pytest.raises(AirflowException, match="Found 0"): + with pytest.raises(DbtCloudResourceLookupError, match="Found 0"): hook.get_job_by_name( project_name=project_name, environment_name=environment_name, @@ -572,7 +573,7 @@ def test_get_job_by_duplicate_name_raises_exception(self, duplicated): return_value=[mock_response_json(mock_list_projects_response)], ), ): - with pytest.raises(AirflowException, match=f"Found 2 {duplicated}"): + with pytest.raises(DbtCloudResourceLookupError, match=f"Found 2 {duplicated}"): hook.get_job_by_name( project_name=PROJECT_NAME, environment_name=ENVIRONMENT_NAME, diff --git a/providers/tests/dbt/cloud/operators/test_dbt.py b/providers/tests/dbt/cloud/operators/test_dbt.py index 05e269548c24f..8791a09a5ba71 100644 --- a/providers/tests/dbt/cloud/operators/test_dbt.py +++ b/providers/tests/dbt/cloud/operators/test_dbt.py @@ -22,7 +22,7 @@ import pytest -from airflow.exceptions import AirflowException, TaskDeferred +from airflow.exceptions import TaskDeferred from airflow.models import DAG, Connection from airflow.providers.dbt.cloud.hooks.dbt import DbtCloudHook, DbtCloudJobRunException, DbtCloudJobRunStatus from airflow.providers.dbt.cloud.operators.dbt import ( @@ -293,7 +293,7 @@ def test_dbt_run_job_by_incorrect_name_raises_exception( dag=self.dag, ) with pytest.raises( - AirflowException, + ValueError, match="Either job_id or project_name, environment_name, and job_name must be provided.", ): dbt_op.execute(MagicMock()) From e426f654e680fc5e226d7b6494f0c4984c0e117d Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Tue, 28 Jan 2025 17:10:35 +0100 Subject: [PATCH 08/10] Correct get_job_by_name docstring --- providers/src/airflow/providers/dbt/cloud/hooks/dbt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index f9bb4e05b0178..66967654fb3ea 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -477,7 +477,7 @@ def get_job_by_name( """ Retrieve metadata for a specific job by combination of project, environment, and job name. - Raises AirflowException if the job is not found or cannot be uniquely identified by provided parameters. + Raises DbtCloudResourceLookupError if the job is not found or cannot be uniquely identified by provided parameters. :param project_name: The name of a dbt Cloud project. :param environment_name: The name of a dbt Cloud environment. From 29c6a6df4fcd46250c3d5345b2dd280c24963472 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Tue, 28 Jan 2025 16:51:44 +0100 Subject: [PATCH 09/10] Enhance dbt Cloud hook and operator with flexible job lookup methods --- .../airflow/providers/dbt/cloud/hooks/dbt.py | 110 ++++++++---- .../providers/dbt/cloud/operators/dbt.py | 43 ++++- providers/tests/dbt/cloud/hooks/test_dbt.py | 170 ++++++++++++++++-- .../tests/dbt/cloud/operators/test_dbt.py | 67 ++++--- 4 files changed, 317 insertions(+), 73 deletions(-) diff --git a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py index 66967654fb3ea..77310aabd5595 100644 --- a/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -389,6 +389,29 @@ def get_project(self, project_id: int, account_id: int | None = None) -> Respons """ return self._run_and_get_response(endpoint=f"{account_id}/projects/{project_id}/", api_version="v3") + @fallback_to_default_account + def get_project_by_name(self, project_name: str, account_id: int | None = None) -> dict: + """ + Retrieve metadata for a specific project using project_name. + + Raises DbtCloudResourceLookupError if the project is not found or cannot be uniquely identified by provided parameters. + + :param project_name: The name of a dbt Cloud project. + :param account_id: Optional. The ID of a dbt Cloud account. + :return: The details of a project. + """ + list_projects_responses = self.list_projects(name_contains=project_name, account_id=account_id) + # flatten & filter the list of responses to find the exact match + projects = [ + project + for response in list_projects_responses + for project in response.json()["data"] + if project["name"] == project_name + ] + if len(projects) != 1: + raise DbtCloudResourceLookupError(f"Found {len(projects)} projects with name `{project_name}`.") + return projects[0] + @fallback_to_default_account def list_environments( self, project_id: int, *, name_contains: str | None = None, account_id: int | None = None @@ -425,6 +448,36 @@ def get_environment( endpoint=f"{account_id}/projects/{project_id}/environments/{environment_id}/", api_version="v3" ) + @fallback_to_default_account + def get_environment_by_name( + self, project_id: int, environment_name: str, *, account_id: int | None = None + ) -> dict: + """ + Retrieve metadata for a specific project's environment using project_id and environment_name. + + Raises DbtCloudResourceLookupError if the environment is not found or cannot be uniquely identified by provided parameters. + + :param project_id: The ID of a dbt Cloud project. + :param environment_name: The name of a dbt Cloud environment. + :param account_id: Optional. The ID of a dbt Cloud account. + :return: The details of an environment. + """ + list_environments_responses = self.list_environments( + project_id=project_id, name_contains=environment_name, account_id=account_id + ) + # flatten & filter the list of responses to find the exact match + environments = [ + env + for response in list_environments_responses + for env in response.json()["data"] + if env["name"] == environment_name + ] + if len(environments) != 1: + raise DbtCloudResourceLookupError( + f"Found {len(environments)} environments with name `{environment_name}` in project `{project_id}`." + ) + return environments[0] + @fallback_to_default_account def list_jobs( self, @@ -472,48 +525,39 @@ def get_job(self, job_id: int, account_id: int | None = None) -> Response: @fallback_to_default_account def get_job_by_name( - self, *, project_name: str, environment_name: str, job_name: str, account_id: int | None = None + self, + *, + project_id: int | None = None, + project_name: str | None = None, + environment_id: int | None = None, + environment_name: str | None = None, + job_name: str, + account_id: int | None = None, ) -> dict: """ Retrieve metadata for a specific job by combination of project, environment, and job name. Raises DbtCloudResourceLookupError if the job is not found or cannot be uniquely identified by provided parameters. - :param project_name: The name of a dbt Cloud project. - :param environment_name: The name of a dbt Cloud environment. - :param job_name: The name of a dbt Cloud job. + :param project_id: The ID of the dbt Cloud project. Can be used interchangeably with + project_name. Either project_id or project_name must be provided. + :param project_name: The name of the dbt Cloud project. Can be used interchangeably with + project_id. Either project_id or project_name must be provided. + :param environment_id: The ID of the dbt Cloud environment. Can be used interchangeably with + environment_name. Either environment_id or environment_name must be provided. + :param environment_name: The name of the dbt Cloud environment. Can be used interchangeably with + environment_id. Either environment_id or environment_name must be provided. + :param job_name: The name of the dbt Cloud job to look up. Required. :param account_id: Optional. The ID of a dbt Cloud account. :return: The details of a job. """ - # get project_id using project_name - list_projects_responses = self.list_projects(name_contains=project_name, account_id=account_id) - # flatten & filter the list of responses to find the exact match - projects = [ - project - for response in list_projects_responses - for project in response.json()["data"] - if project["name"] == project_name - ] - if len(projects) != 1: - raise DbtCloudResourceLookupError(f"Found {len(projects)} projects with name `{project_name}`.") - project_id = projects[0]["id"] + if not project_id: + project_id = self.get_project_by_name(project_name=project_name, account_id=account_id)["id"] - # get environment_id using project_id and environment_name - list_environments_responses = self.list_environments( - project_id=project_id, name_contains=environment_name, account_id=account_id - ) - # flatten & filter the list of responses to find the exact match - environments = [ - env - for response in list_environments_responses - for env in response.json()["data"] - if env["name"] == environment_name - ] - if len(environments) != 1: - raise DbtCloudResourceLookupError( - f"Found {len(environments)} environments with name `{environment_name}` in project `{project_name}`." - ) - environment_id = environments[0]["id"] + if not environment_id: + environment_id = self.get_environment_by_name( + project_id=project_id, environment_name=environment_name, account_id=account_id + )["id"] # get job using project_id, environment_id and job_name list_jobs_responses = self.list_jobs( @@ -531,7 +575,7 @@ def get_job_by_name( ] if len(jobs) != 1: raise DbtCloudResourceLookupError( - f"Found {len(jobs)} jobs with name `{job_name}` in environment `{environment_name}` in project `{project_name}`." + f"Found {len(jobs)} jobs with name `{job_name}` in environment `{(environment_id or environment_name)}` in project `{(project_id or project_name)}`." ) return jobs[0] diff --git a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py index d9aab8e9e2f67..9c6da171aabe1 100644 --- a/providers/src/airflow/providers/dbt/cloud/operators/dbt.py +++ b/providers/src/airflow/providers/dbt/cloud/operators/dbt.py @@ -57,10 +57,19 @@ class DbtCloudRunJobOperator(BaseOperator): :ref:`howto/operator:DbtCloudRunJobOperator` :param dbt_cloud_conn_id: The connection ID for connecting to dbt Cloud. - :param job_id: The ID of a dbt Cloud job. Required if project_name, environment_name, and job_name are not provided. - :param project_name: Optional. The name of a dbt Cloud project. Used only if ``job_id`` is None. - :param environment_name: Optional. The name of a dbt Cloud environment. Used only if ``job_id`` is None. - :param job_name: Optional. The name of a dbt Cloud job. Used only if ``job_id`` is None. + :param job_id: The ID of a dbt Cloud job. You can either provide this ID directly, or let + the operator look it up using a combination of project, environment, and job name. + :param job_name: The name of the dbt Cloud job to run. Only used for job lookup when + job_id is not provided. Must be used together with either project_id/project_name + and environment_id/environment_name. + :param project_id: The ID of the dbt Cloud project. Only used for job lookup when + job_id is not provided. Can be used interchangeably with project_name. + :param project_name: The name of the dbt Cloud project. Only used for job lookup when + job_id is not provided. Can be used interchangeably with project_id. + :param environment_id: The ID of the dbt Cloud environment. Only used for job lookup when + job_id is not provided. Can be used interchangeably with environment_name. + :param environment_name: The name of the dbt Cloud environment. Only used for job lookup when + job_id is not provided. Can be used interchangeably with environment_id. :param account_id: Optional. The ID of a dbt Cloud account. :param trigger_reason: Optional. Description of the reason to trigger the job. Defaults to "Triggered via Apache Airflow by task in the DAG." @@ -89,9 +98,11 @@ class DbtCloudRunJobOperator(BaseOperator): template_fields = ( "dbt_cloud_conn_id", "job_id", + "job_name", + "project_id", "project_name", + "environment_id", "environment_name", - "job_name", "account_id", "trigger_reason", "steps_override", @@ -106,9 +117,11 @@ def __init__( *, dbt_cloud_conn_id: str = DbtCloudHook.default_conn_name, job_id: int | None = None, + job_name: str | None = None, + project_id: int | None = None, project_name: str | None = None, + environment_id: int | None = None, environment_name: str | None = None, - job_name: str | None = None, account_id: int | None = None, trigger_reason: str | None = None, steps_override: list[str] | None = None, @@ -126,9 +139,11 @@ def __init__( self.dbt_cloud_conn_id = dbt_cloud_conn_id self.account_id = account_id self.job_id = job_id + self.job_name = job_name + self.project_id = project_id self.project_name = project_name + self.environment_id = environment_id self.environment_name = environment_name - self.job_name = job_name self.trigger_reason = trigger_reason self.steps_override = steps_override self.schema_override = schema_override @@ -148,9 +163,19 @@ def execute(self, context: Context): ) if self.job_id is None: - if not all([self.project_name, self.environment_name, self.job_name]): + if not ( + (self.project_id or self.project_name) + and (self.environment_id or self.environment_name) + and self.job_name + ): raise ValueError( - "Either job_id or project_name, environment_name, and job_name must be provided." + "Not enough information to lookup the job_id. " + "To identify a job, you must provide either:\n" + "1. A job_id, or\n" + "2. All of the following:\n" + " - project_id or project_name\n" + " - environment_id or environment_name\n" + " - job_name" ) self.job_id = self.hook.get_job_by_name( account_id=self.account_id, diff --git a/providers/tests/dbt/cloud/hooks/test_dbt.py b/providers/tests/dbt/cloud/hooks/test_dbt.py index 1a11662019b46..f5fa49f619ee1 100644 --- a/providers/tests/dbt/cloud/hooks/test_dbt.py +++ b/providers/tests/dbt/cloud/hooks/test_dbt.py @@ -332,6 +332,57 @@ def test_get_project(self, mock_http_run, mock_paginate, conn_id, account_id): ) hook._paginate.assert_not_called() + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object( + DbtCloudHook, "list_projects", return_value=[mock_response_json(DEFAULT_LIST_PROJECTS_RESPONSE)] + ) + def test_get_project_by_name(self, mock_list_projects, conn_id, account_id): + """Test retrieving a project by name returns the correct project details.""" + hook = DbtCloudHook(conn_id) + project_details = hook.get_project_by_name(project_name=PROJECT_NAME, account_id=account_id) + + assert project_details == DEFAULT_LIST_PROJECTS_RESPONSE["data"][0] + _account_id = account_id or DEFAULT_ACCOUNT_ID + mock_list_projects.assert_called_once_with(name_contains=PROJECT_NAME, account_id=_account_id) + + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "list_projects") + def test_get_project_by_name_not_found(self, mock_list_projects, conn_id, account_id): + """Test retrieving a non-existent project by name raises an error.""" + hook = DbtCloudHook(conn_id) + mock_list_projects.return_value = [mock_response_json({"data": []})] + + with pytest.raises(DbtCloudResourceLookupError, match="Found 0 projects"): + hook.get_project_by_name(project_name="non_existent_project", account_id=account_id) + + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "list_projects") + def test_get_project_by_name_multiple_found(self, mock_list_projects, conn_id, account_id): + """Test retrieving a project by name when multiple matches exist raises an error.""" + hook = DbtCloudHook(conn_id) + duplicate_projects_response = { + "data": [ + {"id": PROJECT_ID, "name": PROJECT_NAME}, + {"id": PROJECT_ID + 1, "name": PROJECT_NAME}, + ] + } + mock_list_projects.return_value = [mock_response_json(duplicate_projects_response)] + + with pytest.raises(DbtCloudResourceLookupError, match="Found 2 projects"): + hook.get_project_by_name(project_name=PROJECT_NAME, account_id=account_id) + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], @@ -397,6 +448,75 @@ def test_get_environment(self, mock_http_run, mock_paginate, conn_id, account_id ) hook._paginate.assert_not_called() + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object( + DbtCloudHook, + "list_environments", + return_value=[mock_response_json(DEFAULT_LIST_ENVIRONMENTS_RESPONSE)], + ) + def test_get_environment_by_name(self, mock_list_environments, conn_id, account_id): + """Test retrieving an environment by name returns the correct environment details.""" + hook = DbtCloudHook(conn_id) + environment_details = hook.get_environment_by_name( + project_id=PROJECT_ID, + environment_name=ENVIRONMENT_NAME, + account_id=account_id, + ) + + assert environment_details == DEFAULT_LIST_ENVIRONMENTS_RESPONSE["data"][0] + _account_id = account_id or DEFAULT_ACCOUNT_ID + mock_list_environments.assert_called_once_with( + project_id=PROJECT_ID, + name_contains=ENVIRONMENT_NAME, + account_id=_account_id, + ) + + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "list_environments") + def test_get_environment_by_name_not_found(self, mock_list_environments, conn_id, account_id): + """Test retrieving a non-existent environment by name raises an error.""" + hook = DbtCloudHook(conn_id) + mock_list_environments.return_value = [mock_response_json({"data": []})] + + with pytest.raises(DbtCloudResourceLookupError, match="Found 0 environments"): + hook.get_environment_by_name( + project_id=PROJECT_ID, + environment_name="non_existent_environment", + account_id=account_id, + ) + + @pytest.mark.parametrize( + argnames="conn_id, account_id", + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], + ids=["default_account", "explicit_account"], + ) + @patch.object(DbtCloudHook, "list_environments") + def test_get_environment_by_name_multiple_found(self, mock_list_environments, conn_id, account_id): + """Test retrieving an environment by name when multiple matches exist raises an error.""" + hook = DbtCloudHook(conn_id) + duplicate_environments_response = { + "data": [ + {"id": ENVIRONMENT_ID, "name": ENVIRONMENT_NAME}, + {"id": ENVIRONMENT_ID + 1, "name": ENVIRONMENT_NAME}, + ] + } + mock_list_environments.return_value = [mock_response_json(duplicate_environments_response)] + + with pytest.raises(DbtCloudResourceLookupError, match="Found 2 environments"): + hook.get_environment_by_name( + project_id=PROJECT_ID, + environment_name=ENVIRONMENT_NAME, + account_id=account_id, + ) + @pytest.mark.parametrize( argnames="conn_id, account_id", argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], @@ -469,27 +589,57 @@ def test_get_job(self, mock_http_run, mock_paginate, conn_id, account_id): ) hook._paginate.assert_not_called() + @pytest.mark.parametrize( + argnames="project_id, project_name, environment_id, environment_name, job_name", + argvalues=[ + (PROJECT_ID, None, ENVIRONMENT_ID, None, JOB_NAME), + (None, PROJECT_NAME, None, ENVIRONMENT_NAME, JOB_NAME), + (PROJECT_ID, None, None, ENVIRONMENT_NAME, JOB_NAME), + (None, PROJECT_NAME, ENVIRONMENT_ID, None, JOB_NAME), + ], + ) @patch.object(DbtCloudHook, "list_jobs", return_value=[mock_response_json(DEFAULT_LIST_JOBS_RESPONSE)]) @patch.object( DbtCloudHook, - "list_environments", - return_value=[mock_response_json(DEFAULT_LIST_ENVIRONMENTS_RESPONSE)], - ) - @patch.object( - DbtCloudHook, "list_projects", return_value=[mock_response_json(DEFAULT_LIST_PROJECTS_RESPONSE)] + "get_environment_by_name", + return_value=DEFAULT_LIST_ENVIRONMENTS_RESPONSE["data"][0], ) + @patch.object(DbtCloudHook, "get_project_by_name", return_value=DEFAULT_LIST_PROJECTS_RESPONSE["data"][0]) def test_get_job_by_name_returns_response( - self, mock_list_projects, mock_list_environments, mock_list_jobs + self, + mock_get_project_by_name, + mock_get_environment_by_name, + mock_list_jobs, + project_id, + project_name, + environment_id, + environment_name, + job_name, ): + """Test retrieving a job by project, environment, and job name.""" hook = DbtCloudHook(ACCOUNT_ID_CONN) job_details = hook.get_job_by_name( - project_name=PROJECT_NAME, - environment_name=ENVIRONMENT_NAME, - job_name=JOB_NAME, - account_id=None, + project_id=project_id, + project_name=project_name, + environment_id=environment_id, + environment_name=environment_name, + job_name=job_name, + account_id=DEFAULT_ACCOUNT_ID, ) assert job_details == DEFAULT_LIST_JOBS_RESPONSE["data"][0] + if project_id is None: + mock_get_project_by_name.assert_called_once_with( + project_name=project_name, account_id=DEFAULT_ACCOUNT_ID + ) + else: + mock_get_project_by_name.assert_not_called() + if environment_id is None: + mock_get_environment_by_name.assert_called_once_with( + project_id=PROJECT_ID, environment_name=environment_name, account_id=DEFAULT_ACCOUNT_ID + ) + else: + mock_get_environment_by_name.assert_not_called() @pytest.mark.parametrize( argnames="project_name, environment_name, job_name", diff --git a/providers/tests/dbt/cloud/operators/test_dbt.py b/providers/tests/dbt/cloud/operators/test_dbt.py index 8791a09a5ba71..f0e9444979b6c 100644 --- a/providers/tests/dbt/cloud/operators/test_dbt.py +++ b/providers/tests/dbt/cloud/operators/test_dbt.py @@ -210,6 +210,15 @@ 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" + @pytest.mark.parametrize( + argnames="project_id, project_name, environment_id, environment_name, job_name", + argvalues=[ + (PROJECT_ID, None, ENVIRONMENT_ID, None, JOB_NAME), + (None, PROJECT_NAME, None, ENVIRONMENT_NAME, JOB_NAME), + (PROJECT_ID, None, None, ENVIRONMENT_NAME, JOB_NAME), + (None, PROJECT_NAME, ENVIRONMENT_ID, None, JOB_NAME), + ], + ) @patch( "airflow.providers.dbt.cloud.hooks.dbt.DbtCloudHook.get_job_by_name", return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RESPONSE), @@ -224,18 +233,26 @@ def test_dbt_run_job_op_async(self, mock_trigger_job_run, mock_dbt_hook, mock_jo return_value=mock_response_json(DEFAULT_ACCOUNT_JOB_RUN_RESPONSE), ) def test_dbt_run_job_by_name( - self, mock_trigger_job_run, mock_dbt_hook, mock_job_run_status, mock_job_by_name + self, + mock_trigger_job_run, + mock_dbt_hook, + mock_job_run_status, + mock_job_by_name, + project_id, + project_name, + environment_id, + environment_name, + job_name, ): - """ - Test alternative way to run a job by project, - environment and job name instead of job id. - """ + """Test running a job by project, environment, and job name.""" dbt_op = DbtCloudRunJobOperator( dbt_cloud_conn_id=ACCOUNT_ID_CONN, task_id=TASK_ID, - project_name=PROJECT_NAME, - environment_name=ENVIRONMENT_NAME, - job_name=JOB_NAME, + project_id=project_id, + project_name=project_name, + environment_id=environment_id, + environment_name=environment_name, + job_name=job_name, check_interval=1, timeout=3, dag=self.dag, @@ -244,12 +261,22 @@ def test_dbt_run_job_by_name( mock_trigger_job_run.assert_called_once() @pytest.mark.parametrize( - argnames="project_name, environment_name, job_name", + argnames="project_id, project_name, environment_id, environment_name, job_name", argvalues=[ - (None, ENVIRONMENT_NAME, JOB_NAME), - (PROJECT_NAME, "", JOB_NAME), - (PROJECT_NAME, ENVIRONMENT_NAME, None), - ("", "", ""), + # Missing all required fields + (None, None, None, None, None), + # Missing project (both id and name) + (None, None, ENVIRONMENT_ID, None, JOB_NAME), + (None, None, None, ENVIRONMENT_NAME, JOB_NAME), + (None, "", None, ENVIRONMENT_NAME, JOB_NAME), + # Missing environment (both id and name) + (PROJECT_ID, None, None, None, JOB_NAME), + (None, PROJECT_NAME, None, None, JOB_NAME), + (None, PROJECT_NAME, None, "", JOB_NAME), + # Missing job name + (PROJECT_ID, None, ENVIRONMENT_ID, None, None), + (None, PROJECT_NAME, None, ENVIRONMENT_NAME, None), + (None, PROJECT_NAME, ENVIRONMENT_ID, None, ""), ], ) @patch( @@ -271,21 +298,19 @@ def test_dbt_run_job_by_incorrect_name_raises_exception( mock_dbt_hook, mock_job_run_status, mock_job_by_name, + project_id, project_name, + environment_id, environment_name, job_name, ): - """ - Test alternative way to run a job by project, - environment and job name instead of job id. - - This test is to check if the operator raises an exception - when the project, environment or job name is missing. - """ + """Test that the operator raises an exception when job lookup information is invalid.""" dbt_op = DbtCloudRunJobOperator( dbt_cloud_conn_id=ACCOUNT_ID_CONN, task_id=TASK_ID, + project_id=project_id, project_name=project_name, + environment_id=environment_id, environment_name=environment_name, job_name=job_name, check_interval=1, @@ -294,7 +319,7 @@ def test_dbt_run_job_by_incorrect_name_raises_exception( ) with pytest.raises( ValueError, - match="Either job_id or project_name, environment_name, and job_name must be provided.", + match="Not enough information to lookup the job_id.", ): dbt_op.execute(MagicMock()) mock_trigger_job_run.assert_not_called() From 1e9251ba93fabc6f72a529bf4da81462d9fc5d67 Mon Sep 17 00:00:00 2001 From: Marcin Sitnik Date: Tue, 28 Jan 2025 17:02:09 +0100 Subject: [PATCH 10/10] Update dbt Cloud job lookup documentation --- .../operators.rst | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/docs/apache-airflow-providers-dbt-cloud/operators.rst b/docs/apache-airflow-providers-dbt-cloud/operators.rst index 8d069bb869bcd..ba6f06aef5775 100644 --- a/docs/apache-airflow-providers-dbt-cloud/operators.rst +++ b/docs/apache-airflow-providers-dbt-cloud/operators.rst @@ -82,10 +82,21 @@ via the ``additional_run_config`` dictionary. :start-after: [START howto_operator_dbt_cloud_run_job_async] :end-before: [END howto_operator_dbt_cloud_run_job_async] -You can also trigger a dbt Cloud job without providing the ``job_id``. Instead, you can identify the job -by providing the ``project_name``, ``environment_name``, and ``job_name``. -Please note that it will only work if the above three parameters uniquely identify a job in your account -(i.e. you cannot have two jobs with the same name in the same project and environment). +You can trigger a dbt Cloud job in two ways: + +1. Directly using the ``job_id`` parameter +2. Looking up the job using a combination of identifiers: + + * Project: either ``project_id`` or ``project_name`` + * Environment: either ``environment_id`` or ``environment_name`` + * Job: ``job_name`` + +When using the lookup method, you must provide either IDs or names for both the project and environment, +along with the job name. For example, you could use ``project_name``, ``environment_id``, and ``job_name``, +or ``project_id``, ``environment_name``, and ``job_name``. + +Please note that the lookup method will only work if the provided combination uniquely identifies a job +in your account. The job name must be unique within the specified project and environment. .. exampleinclude:: /../../providers/tests/system/dbt/cloud/example_dbt_cloud.py :language: python