diff --git a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py index 89e512c51fb02..b4ced28bbb7d5 100644 --- a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from requests.models import PreparedRequest, Response - from airflow.models import Connection + from airflow.providers.common.compat.sdk import Connection DBT_CAUSE_MAX_LENGTH = 255 @@ -254,12 +254,13 @@ def get_request_url_params( async def get_headers_tenants_from_connection(self) -> tuple[dict[str, Any], str]: """Get Headers, tenants from the connection details.""" + conn = await self._resolve_connection_async() headers: dict[str, Any] = {} - tenant = self._get_tenant_domain(self.connection) + tenant = self._get_tenant_domain(conn) package_name, provider_version = _get_provider_info() headers["User-Agent"] = f"{package_name}-v{provider_version}" headers["Content-Type"] = "application/json" - headers["Authorization"] = f"Token {self.connection.password}" + headers["Authorization"] = f"Token {conn.password}" return headers, tenant def _log_request_error(self, attempt_num: int, error: str) -> None: @@ -307,7 +308,8 @@ async def get_job_details( endpoint = f"{account_id}/runs/{run_id}/" headers, tenant = await self.get_headers_tenants_from_connection() url, params = self.get_request_url_params(tenant, endpoint, include_related) - proxies = self._get_proxies(self.connection) or {} + conn = await self._resolve_connection_async() + proxies = self._get_proxies(conn) or {} proxy = proxies.get("https") if proxies and url.startswith("https") else proxies.get("http") extra_request_args = {} @@ -341,13 +343,40 @@ async def get_job_status( job_run_status: int = response["data"]["status"] return job_run_status + @staticmethod + def _require_password(conn: Connection) -> Connection: + if not conn.password: + raise AirflowException("An API token is required to connect to dbt Cloud.") + return conn + @cached_property def connection(self) -> Connection: - _connection = self.get_connection(self.dbt_cloud_conn_id) - if not _connection.password: - raise AirflowException("An API token is required to connect to dbt Cloud.") + """ + Resolve and cache the dbt Cloud connection (sync). - return _connection # type: ignore[return-value] + Do not read this property from async code running inside the triggerer's + event loop — it calls the synchronous ``get_connection()``, whose secret-masking + step raises ``RuntimeError`` when invoked from a thread that's already running an + event loop. Use ``_resolve_connection_async()`` instead; it shares this property's + cache slot so the connection is still only looked up once per hook instance + regardless of which path is used first. + """ + return self._require_password(self.get_connection(self.dbt_cloud_conn_id)) + + async def _resolve_connection_async(self) -> Connection: + """ + Resolve and cache the dbt Cloud connection (async). + + Shares the ``connection`` cached_property's cache slot so a connection + fetched from either the sync or async path is not looked up twice on + the same hook instance, and so the async path never touches the sync + ``get_connection()``/``mask_secret()`` path from inside a running + event loop (which raises in the triggerer). + """ + if "connection" not in self.__dict__: + conn = await get_async_connection(self.dbt_cloud_conn_id) + self.__dict__["connection"] = self._require_password(conn) + return self.__dict__["connection"] def get_conn(self, *args, **kwargs) -> Session: tenant = self._get_tenant_domain(self.connection) diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py b/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py index 3ba9653bb0be6..951e850549609 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py +++ b/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py @@ -356,6 +356,107 @@ async def test_account_id_cache_shared_between_sync_and_async(self): assert mock_get_connection.call_count == 1 assert mock_get_async_connection.call_count == 0 + @pytest.mark.asyncio + async def test_get_headers_tenants_from_connection_does_not_use_sync_get_connection(self): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + + with ( + patch.object(DbtCloudHook, "get_connection") as mock_get_connection, + patch( + "airflow.providers.dbt.cloud.hooks.dbt.get_async_connection", + new=AsyncMock( + return_value=Connection( + conn_id=ACCOUNT_ID_CONN, + conn_type=DbtCloudHook.conn_type, + login=str(DEFAULT_ACCOUNT_ID), + password=TOKEN, + host=SINGLE_TENANT_DOMAIN, + ) + ), + ) as mock_get_async_connection, + ): + headers, tenant = await hook.get_headers_tenants_from_connection() + + assert tenant == SINGLE_TENANT_DOMAIN + assert headers["Authorization"] == f"Token {TOKEN}" + mock_get_connection.assert_not_called() + assert mock_get_async_connection.call_count == 1 + + @pytest.mark.asyncio + async def test_resolve_connection_cached_async(self): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + + with patch( + "airflow.providers.dbt.cloud.hooks.dbt.get_async_connection", + new=AsyncMock( + return_value=Connection( + conn_id=ACCOUNT_ID_CONN, + conn_type=DbtCloudHook.conn_type, + login=str(DEFAULT_ACCOUNT_ID), + password=TOKEN, + ) + ), + ) as mock_get_async_connection: + first_call = await hook._resolve_connection_async() + second_call = await hook._resolve_connection_async() + + assert first_call.password == TOKEN + assert second_call.password == TOKEN + assert mock_get_async_connection.call_count == 1 + + @pytest.mark.asyncio + async def test_connection_cache_shared_between_sync_and_async(self): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + + with ( + patch.object( + DbtCloudHook, + "get_connection", + return_value=Connection( + conn_id=ACCOUNT_ID_CONN, + conn_type=DbtCloudHook.conn_type, + login=str(DEFAULT_ACCOUNT_ID), + password=TOKEN, + ), + ) as mock_get_connection, + patch( + "airflow.providers.dbt.cloud.hooks.dbt.get_async_connection", + new=AsyncMock( + return_value=Connection( + conn_id=ACCOUNT_ID_CONN, + conn_type=DbtCloudHook.conn_type, + login=str(DEFAULT_ACCOUNT_ID), + password=TOKEN, + ) + ), + ) as mock_get_async_connection, + ): + sync_conn = hook.connection + async_conn = await hook._resolve_connection_async() + + assert sync_conn.password == TOKEN + assert async_conn.password == TOKEN + + assert mock_get_connection.call_count == 1 + assert mock_get_async_connection.call_count == 0 + + @pytest.mark.asyncio + async def test_resolve_connection_async_requires_password(self): + hook = DbtCloudHook(ACCOUNT_ID_CONN) + + with patch( + "airflow.providers.dbt.cloud.hooks.dbt.get_async_connection", + new=AsyncMock( + return_value=Connection( + conn_id=ACCOUNT_ID_CONN, + conn_type=DbtCloudHook.conn_type, + login=str(DEFAULT_ACCOUNT_ID), + ) + ), + ): + with pytest.raises(AirflowException, match="An API token is required"): + await hook._resolve_connection_async() + @pytest.mark.parametrize( argnames=("conn_id", "account_id"), argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)],