diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py index fe028f881e2de..647cd9628b3f0 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -194,6 +194,15 @@ def my_after_func(retry_state): def databricks_conn(self) -> Connection: return self.get_connection(self.databricks_conn_id) # type: ignore[return-value] + async def aget_connection(self, conn_id: str) -> Connection: + """Async version of get_connection for use in async contexts.""" + return await BaseHook.aget_connection(conn_id) + + @property + async def a_databricks_conn(self) -> Connection: + """Async version of databricks_conn for use in async contexts (triggerer).""" + return await self.aget_connection(self.databricks_conn_id) + def get_conn(self) -> Connection: return self.databricks_conn @@ -263,6 +272,13 @@ def _get_connection_attr(self, attr_name: str) -> str: raise ValueError(f"`{attr_name}` must be present in Connection") return attr + async def _a_get_connection_attr(self, attr_name: str) -> str: + """Async version of _get_connection_attr for use in async contexts.""" + conn = await self.a_databricks_conn + if not (attr := getattr(conn, attr_name)): + raise ValueError(f"`{attr_name}` must be present in Connection") + return attr + @cached_property def proxies(self) -> dict[str, str] | None: """Return validated proxy configuration from connection extras.""" @@ -367,14 +383,13 @@ async def _a_get_sp_token(self, resource: str) -> str: return sp_token["access_token"] self.log.info("Existing Service Principal token is expired, or going to expire soon. Refreshing...") + conn = await self.a_databricks_conn try: async for attempt in self._a_get_retry_object(): with attempt: async with self._session.post( resource, - auth=aiohttp.BasicAuth( - self._get_connection_attr("login"), self.databricks_conn.password - ), + auth=aiohttp.BasicAuth(await self._a_get_connection_attr("login"), conn.password), data="grant_type=client_credentials&scope=all-apis", headers={ **self.user_agent_header, @@ -475,12 +490,11 @@ async def _a_get_aad_token(self, resource: str) -> str: ManagedIdentityCredential as AsyncManagedIdentityCredential, ) + conn = await self.a_databricks_conn async for attempt in self._a_get_retry_object(): with attempt: - if self.databricks_conn.extra_dejson.get("use_azure_managed_identity", False): - client_id = self.databricks_conn.extra_dejson.get( - "azure_managed_identity_client_id", None - ) + if conn.extra_dejson.get("use_azure_managed_identity", False): + client_id = conn.extra_dejson.get("azure_managed_identity_client_id", None) # Managed identity authenticates against the link-local IMDS endpoint # (169.254.169.254), which must be reached directly and is unsupported behind a # proxy, so the `proxies` extra is intentionally not forwarded here (unlike the @@ -490,9 +504,9 @@ async def _a_get_aad_token(self, resource: str) -> str: token = await credential.get_token(f"{resource}/.default") else: async with AsyncClientSecretCredential( - client_id=self._get_connection_attr("login"), - client_secret=self.databricks_conn.password, - tenant_id=self.databricks_conn.extra_dejson["azure_tenant_id"], + client_id=await self._a_get_connection_attr("login"), + client_secret=conn.password, + tenant_id=conn.extra_dejson["azure_tenant_id"], **self._get_azure_credential_kwargs(), ) as credential: token = await credential.get_token(f"{resource}/.default") @@ -576,7 +590,7 @@ async def _a_get_aad_token_for_default_az_credential(self, resource: str) -> str DefaultAzureCredential as AsyncDefaultAzureCredential, ) - for attempt in self._get_retry_object(): + async for attempt in self._a_get_retry_object(): with attempt: # This only works in an Azure Kubernetes Service Cluster given the following environment variables: # AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_FEDERATED_TOKEN_FILE @@ -599,9 +613,8 @@ async def _a_get_aad_token_for_default_az_credential(self, resource: str) -> str raise AirflowOptionalProviderFeatureException(e) except RetryError: raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.") - except requests_exceptions.HTTPError as e: - msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}" - raise AirflowException(msg) + except aiohttp.ClientResponseError as err: + raise AirflowException(f"Response: {err.message}, Status Code: {err.status}") return token.token @@ -627,11 +640,10 @@ async def _a_get_aad_headers(self) -> dict: :return: dictionary with filled AAD headers """ headers = {} - if "azure_resource_id" in self.databricks_conn.extra_dejson: + conn = await self.a_databricks_conn + if "azure_resource_id" in conn.extra_dejson: mgmt_token = await self._a_get_aad_token(AZURE_MANAGEMENT_ENDPOINT) - headers["X-Databricks-Azure-Workspace-Resource-Id"] = self.databricks_conn.extra_dejson[ - "azure_resource_id" - ] + headers["X-Databricks-Azure-Workspace-Resource-Id"] = conn.extra_dejson["azure_resource_id"] headers["X-Databricks-Azure-SP-Management-Token"] = mgmt_token return headers @@ -668,7 +680,8 @@ def _get_k8s_jwt_token(self) -> str: async def _a_get_k8s_jwt_token(self) -> str: """Async version of _get_k8s_jwt_token().""" - if "k8s_projected_volume_token_path" in self.databricks_conn.extra_dejson: + conn = await self.a_databricks_conn + if "k8s_projected_volume_token_path" in conn.extra_dejson: self.log.info("Using Kubernetes projected volume token") return await self._a_get_k8s_projected_volume_token() @@ -725,8 +738,9 @@ def _get_aiofiles(): async def _a_get_k8s_projected_volume_token(self) -> str: """Async version of _get_k8s_projected_volume_token().""" aiofiles = self._get_aiofiles() + conn = await self.a_databricks_conn - projected_token_path: str = self.databricks_conn.extra_dejson["k8s_projected_volume_token_path"] + projected_token_path: str = conn.extra_dejson["k8s_projected_volume_token_path"] try: async with aiofiles.open(projected_token_path) as f: @@ -836,15 +850,12 @@ def _get_k8s_token_request_api(self) -> str: async def _a_get_k8s_token_request_api(self) -> str: """Async version of _get_k8s_token_request_api().""" aiofiles = self._get_aiofiles() + conn = await self.a_databricks_conn - audience = self.databricks_conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) - expiration_seconds = self.databricks_conn.extra_dejson.get("expiration_seconds", 3600) - token_path = self.databricks_conn.extra_dejson.get( - "k8s_token_path", DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH - ) - namespace_path = self.databricks_conn.extra_dejson.get( - "k8s_namespace_path", DEFAULT_K8S_NAMESPACE_PATH - ) + audience = conn.extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) + expiration_seconds = conn.extra_dejson.get("expiration_seconds", 3600) + token_path = conn.extra_dejson.get("k8s_token_path", DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH) + namespace_path = conn.extra_dejson.get("k8s_namespace_path", DEFAULT_K8S_NAMESPACE_PATH) try: async with aiofiles.open(token_path) as f: @@ -923,6 +934,24 @@ def _get_required_client_id(self) -> str: ) return client_id + async def _a_get_required_client_id(self) -> str: + """ + Async version of _get_required_client_id for use in async contexts. + + :return: Service principal client ID + :raises AirflowException: If client_id is not provided + """ + conn = await self.a_databricks_conn + client_id = conn.extra_dejson.get("client_id") + if not client_id: + raise AirflowException( + "client_id is required for Kubernetes OIDC token federation. " + "Kubernetes service account tokens do not support custom claims, " + "so service principal-level federation must be used. " + "Please provide client_id in the connection extra parameters." + ) + return client_id + def _get_federation_subject_token(self) -> tuple[str, str | None]: """ Resolve the OIDC JWT to exchange for a Databricks token (RFC 8693 ``subject_token``). @@ -960,18 +989,19 @@ def _get_federation_subject_token(self) -> tuple[str, str | None]: async def _a_get_federation_subject_token(self) -> tuple[str, str | None]: """Async version of :meth:`_get_federation_subject_token`.""" - provider = self.databricks_conn.extra_dejson.get("federated_token_provider") + conn = await self.a_databricks_conn + provider = conn.extra_dejson.get("federated_token_provider") if provider: # The provider is a synchronous callable that typically makes a blocking network call to # mint the token. Offload it to a worker thread so it can't stall the triggerer event loop. loop = asyncio.get_running_loop() subject_token = await loop.run_in_executor(None, self._resolve_supplied_subject_token, provider) - return subject_token, self.databricks_conn.extra_dejson.get("client_id") - if self._is_aws_federation(): + return subject_token, conn.extra_dejson.get("client_id") + if await self._a_is_aws_federation(): loop = asyncio.get_running_loop() subject_token = await loop.run_in_executor(None, self._get_aws_subject_token) - return subject_token, self.databricks_conn.extra_dejson.get("client_id") - client_id = self._get_required_client_id() + return subject_token, conn.extra_dejson.get("client_id") + client_id = await self._a_get_required_client_id() return await self._a_get_k8s_jwt_token(), client_id def _is_aws_federation(self) -> bool: @@ -980,6 +1010,11 @@ def _is_aws_federation(self) -> bool: "federated_aws", False ) + async def _a_is_aws_federation(self) -> bool: + """Async version of _is_aws_federation for use in async contexts.""" + conn = await self.a_databricks_conn + return conn.login == "federated_aws" or conn.extra_dejson.get("federated_aws", False) + def _get_aws_subject_token(self) -> str: """ Mint an AWS-signed OIDC JWT for the exchange via AWS STS ``GetWebIdentityToken``. @@ -1248,41 +1283,40 @@ def _get_token(self, raise_error: bool = False) -> str | None: return None async def _a_get_token(self, raise_error: bool = False) -> str | None: - if "token" in self.databricks_conn.extra_dejson: + conn = await self.a_databricks_conn + if "token" in conn.extra_dejson: self.log.info( "Using token auth. For security reasons, please set token in Password field instead of extra" ) - return self.databricks_conn.extra_dejson["token"] - if not self.databricks_conn.login and self.databricks_conn.password: + return conn.extra_dejson["token"] + if not conn.login and conn.password: self.log.debug("Using token auth.") - return self.databricks_conn.password - if "azure_tenant_id" in self.databricks_conn.extra_dejson: - if self.databricks_conn.login == "" or self.databricks_conn.password == "": + return conn.password + if "azure_tenant_id" in conn.extra_dejson: + if conn.login == "" or conn.password == "": raise AirflowException("Azure SPN credentials aren't provided") self.log.debug("Using AAD Token for SPN.") return await self._a_get_aad_token(DEFAULT_DATABRICKS_SCOPE) - if self.databricks_conn.extra_dejson.get("use_azure_managed_identity", False): + if conn.extra_dejson.get("use_azure_managed_identity", False): self.log.debug("Using AAD Token for managed identity.") await self._a_check_azure_metadata_service() return await self._a_get_aad_token(DEFAULT_DATABRICKS_SCOPE) - if self.databricks_conn.extra_dejson.get(DEFAULT_AZURE_CREDENTIAL_SETTING_KEY, False): + if conn.extra_dejson.get(DEFAULT_AZURE_CREDENTIAL_SETTING_KEY, False): self.log.debug("Using AzureDefaultCredential for authentication.") return await self._a_get_aad_token_for_default_az_credential(DEFAULT_DATABRICKS_SCOPE) - if self.databricks_conn.extra_dejson.get("service_principal_oauth", False): - if self.databricks_conn.login == "" or self.databricks_conn.password == "": + if conn.extra_dejson.get("service_principal_oauth", False): + if conn.login == "" or conn.password == "": raise AirflowException("Service Principal credentials aren't provided") self.log.debug("Using Service Principal Token.") return await self._a_get_sp_token(self._get_oidc_token_service_url()) - if self.databricks_conn.extra_dejson.get("federated_token_provider"): + if conn.extra_dejson.get("federated_token_provider"): self.log.debug("Using OIDC token federation with a supplied token provider.") return await self._a_get_federated_databricks_token(self._get_oidc_token_service_url()) - if self._is_aws_federation(): + if await self._a_is_aws_federation(): self.log.debug("Using AWS IAM OIDC token federation.") return await self._a_get_federated_databricks_token(self._get_oidc_token_service_url()) - if self.databricks_conn.login == "federated_k8s" or self.databricks_conn.extra_dejson.get( - "federated_k8s", False - ): + if conn.login == "federated_k8s" or conn.extra_dejson.get("federated_k8s", False): self.log.debug("Using Kubernetes OIDC token federation.") return await self._a_get_federated_databricks_token(self._get_oidc_token_service_url()) if raise_error: @@ -1306,6 +1340,13 @@ def _endpoint_url(self, endpoint): schema = self.databricks_conn.schema or "https" return f"{schema}://{self.host}{port}/{endpoint}" + async def _a_endpoint_url(self, endpoint): + """Async version of _endpoint_url for use in async contexts.""" + conn = await self.a_databricks_conn + port = f":{conn.port}" if conn.port else "" + schema = conn.schema or "https" + return f"{schema}://{self.host}{port}/{endpoint}" + def _do_api_call( self, endpoint_info: tuple[str, str], @@ -1396,7 +1437,7 @@ async def _a_do_api_call(self, endpoint_info: tuple[str, str], json: dict[str, A method, endpoint = endpoint_info full_endpoint = f"api/{endpoint}" - url = self._endpoint_url(full_endpoint) + url = await self._a_endpoint_url(full_endpoint) aad_headers = await self._a_get_aad_headers() headers = {**self.user_agent_header, **aad_headers} @@ -1407,7 +1448,8 @@ async def _a_do_api_call(self, endpoint_info: tuple[str, str], json: dict[str, A auth = BearerAuth(token) else: self.log.info("Using basic auth.") - auth = aiohttp.BasicAuth(self._get_connection_attr("login"), self.databricks_conn.password) + conn = await self.a_databricks_conn + auth = aiohttp.BasicAuth(await self._a_get_connection_attr("login"), conn.password or "") request_func: Any if method == "GET": diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py index be87a82483c5f..9f5ba62195eaf 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py @@ -2522,3 +2522,268 @@ def test_sqlstatementstate_from_json(self): assert obj.state == "FAILED" assert obj.error_code == "123" assert obj.error_message == "Error occurred" + + +@pytest.mark.db_test +class TestDatabricksHookAsyncConnection: + """ + Tests for async connection retrieval in BaseDatabricksHook. + + These tests verify that the async connection methods work correctly + and do not raise AsyncToSync RuntimeError when called from within + an async event loop (e.g., in the triggerer). + """ + + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + extra=None, + ) + ) + self.hook = DatabricksHook(retry_delay=0) + + @pytest.mark.asyncio + async def test_aget_connection_returns_valid_connection(self): + """Test that aget_connection returns a valid Connection object.""" + conn = await self.hook.aget_connection(DEFAULT_CONN_ID) + assert conn is not None + assert conn.host == HOST + assert conn.login == LOGIN + assert conn.password == PASSWORD + + @pytest.mark.asyncio + async def test_aget_connection_with_token(self, create_connection_without_db): + """Test that aget_connection works with token authentication.""" + create_connection_without_db( + Connection( + conn_id="databricks_token", + conn_type="databricks", + host=HOST, + login=None, + password=TOKEN, + extra=json.dumps({"token": TOKEN}), + ) + ) + hook = DatabricksHook(databricks_conn_id="databricks_token", retry_delay=0) + conn = await hook.aget_connection("databricks_token") + assert conn.password == TOKEN + + @pytest.mark.asyncio + async def test_a_databricks_conn_is_available(self): + """Test that a_databricks_conn async property returns connection.""" + async with self.hook: + conn = await self.hook.a_databricks_conn + assert conn is not None + assert conn.host == HOST + assert conn.login == LOGIN + + @pytest.mark.asyncio + async def test_a_do_api_call_uses_async_connection(self): + """Test that _a_do_api_call works with async connection.""" + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + result = await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + + assert result == GET_RUN_RESPONSE + mock_get.assert_called_once() + + @pytest.mark.asyncio + async def test_a_do_api_call_in_running_event_loop(self): + """ + Test that _a_do_api_call works when called from within a running event loop. + + This simulates the triggerer environment where an async event loop is already + running. This test ensures no AsyncToSync RuntimeError is raised. + """ + + async def run_in_existing_loop(): + """Run async API call in an existing event loop - simulates triggerer.""" + async with self.hook: + # This should NOT raise: RuntimeError: You cannot use AsyncToSync... + result = await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + return result + + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + # Create a new event loop and run the async function + result = await run_in_existing_loop() + assert result == GET_RUN_RESPONSE + + @pytest.mark.asyncio + async def test_a_get_run_state_uses_async_connection(self): + """Test that a_get_run_state works with async connection.""" + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + run_state = await self.hook.a_get_run_state(RUN_ID) + + assert run_state.life_cycle_state == LIFE_CYCLE_STATE + assert run_state.state_message == STATE_MESSAGE + + @pytest.mark.asyncio + async def test_a_get_run_state_in_triggerer_context(self): + """ + Test that a_get_run_state works when called from the triggerer context. + + This is the actual code path used by DatabricksExecutionTrigger. + """ + + async def triggerer_like_polling(): + """Simulate the triggerer's polling loop.""" + async with self.hook: + while True: + run_state = await self.hook.a_get_run_state(RUN_ID) + if run_state.is_terminal: + return run_state + # In real triggerer, would sleep here + + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + result = await triggerer_like_polling() + assert result.life_cycle_state == LIFE_CYCLE_STATE + + @pytest.mark.asyncio + async def test_multiple_async_calls_in_same_loop(self): + """ + Test that multiple async API calls work in the same event loop. + + This verifies that the async connection caching works correctly across + multiple calls within the same async context. + """ + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + # Make multiple API calls in sequence + result1 = await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + result2 = await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID + 1}) + result3 = await self.hook.a_get_run_state(RUN_ID) + + assert result1 == GET_RUN_RESPONSE + assert result2 == GET_RUN_RESPONSE + assert isinstance(result3, RunState) + + @pytest.mark.asyncio + async def test_async_connection_not_blocking_event_loop(self): + """ + Test that async connection methods do not block the event loop. + + If sync connection was used incorrectly, this would block the event loop. + """ + import asyncio + + async def timed_async_call(): + """Make async call and track if it was truly async.""" + async with self.hook: + await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_response = mock.MagicMock() + mock_response.json = AsyncMock(return_value=GET_RUN_RESPONSE) + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_get.return_value = mock_response + + # Run with timeout to detect blocking + try: + result = await asyncio.wait_for(timed_async_call(), timeout=5.0) + assert result is None # No exception means success + except asyncio.TimeoutError: + pytest.fail("Async call blocked the event loop - possible sync call inside async context") + + +@pytest.mark.db_test +class TestDatabricksHookAsyncConnectionWithProxies: + """ + Tests for async connection with proxy configuration. + """ + + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + extra=json.dumps({"proxies": PROXIES}), + ) + ) + self.hook = DatabricksHook(retry_delay=0) + + @pytest.mark.asyncio + async def test_a_do_api_call_uses_proxies(self): + """Test that async API call respects proxy configuration.""" + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + + # Verify proxy was used + call_kwargs = mock_get.call_args.kwargs + assert "proxy" in call_kwargs + assert call_kwargs["proxy"] == PROXIES["https"] + + +@pytest.mark.db_test +class TestDatabricksHookAsyncConnectionWithTokenAuth: + """ + Tests for async connection with token authentication. + """ + + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=TOKEN, + extra=None, + ) + ) + self.hook = DatabricksHook(retry_delay=0) + + @pytest.mark.asyncio + async def test_a_do_api_call_with_token_auth(self): + """Test that async API call works with token authentication.""" + with mock.patch( + "airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get" + ) as mock_get: + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + result = await self.hook._a_do_api_call(("GET", "2.0/jobs/runs/get"), {"run_id": RUN_ID}) + + assert result == GET_RUN_RESPONSE + # Verify Bearer auth was used + call_kwargs = mock_get.call_args.kwargs + assert isinstance(call_kwargs["auth"], BearerAuth) + assert call_kwargs["auth"].token == TOKEN