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 ad48f09a9360b..ba2ed493b3ffd 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -196,6 +196,32 @@ def databricks_conn(self) -> Connection: def get_conn(self) -> Connection: return self.databricks_conn + + async def adatabricks_conn(self): + if not hasattr(self, "_adatabricks_conn"): + self._adatabricks_conn = await self.aget_connection(self.databricks_conn_id) + return self._adatabricks_conn + + async def ahost(self): + conn = await self.adatabricks_conn() + host = None + if "host" in conn.extra_dejson: + host = self._parse_host(conn.extra_dejson["host"]) + elif conn.host: + host = self._parse_host(conn.host) + return host + + async def _a_endpoint_url(self, endpoint: str) -> str: + conn = await self.adatabricks_conn() + port = f":{conn.port}" if conn.port else "" + schema = conn.schema or "https" + host = await self.ahost() + return f"{schema}://{host}{port}/{endpoint}" + + async def _a_get_oidc_token_service_url(self) -> str: + host = await self.ahost() + return OIDC_TOKEN_SERVICE_URL.format(f"https://{host}") + @cached_property def user_agent_header(self) -> dict[str, str]: return {"user-agent": self.user_agent_value} @@ -372,7 +398,7 @@ async def _a_get_sp_token(self, resource: str) -> str: async with self._session.post( resource, auth=aiohttp.BasicAuth( - self._get_connection_attr("login"), self.databricks_conn.password + self._get_connection_attr("login"), (await self.adatabricks_conn()).password ), data="grant_type=client_credentials&scope=all-apis", headers={ @@ -476,8 +502,8 @@ async def _a_get_aad_token(self, resource: str) -> str: 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( + if (await self.adatabricks_conn()).extra_dejson.get("use_azure_managed_identity", False): + client_id = (await self.adatabricks_conn()).extra_dejson.get( "azure_managed_identity_client_id", None ) # Managed identity authenticates against the link-local IMDS endpoint @@ -490,8 +516,8 @@ async def _a_get_aad_token(self, resource: str) -> str: 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_secret=(await self.adatabricks_conn()).password, + tenant_id=(await self.adatabricks_conn()).extra_dejson["azure_tenant_id"], **self._get_azure_credential_kwargs(), ) as credential: token = await credential.get_token(f"{resource}/.default") @@ -626,9 +652,9 @@ 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: + if "azure_resource_id" in (await self.adatabricks_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[ + headers["X-Databricks-Azure-Workspace-Resource-Id"] = (await self.adatabricks_conn()).extra_dejson[ "azure_resource_id" ] headers["X-Databricks-Azure-SP-Management-Token"] = mgmt_token @@ -667,7 +693,7 @@ 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: + if "k8s_projected_volume_token_path" in (await self.adatabricks_conn()).extra_dejson: self.log.info("Using Kubernetes projected volume token") return await self._a_get_k8s_projected_volume_token() @@ -725,7 +751,7 @@ async def _a_get_k8s_projected_volume_token(self) -> str: """Async version of _get_k8s_projected_volume_token().""" aiofiles = self._get_aiofiles() - projected_token_path: str = self.databricks_conn.extra_dejson["k8s_projected_volume_token_path"] + projected_token_path: str = (await self.adatabricks_conn()).extra_dejson["k8s_projected_volume_token_path"] try: async with aiofiles.open(projected_token_path) as f: @@ -836,12 +862,12 @@ async def _a_get_k8s_token_request_api(self) -> str: """Async version of _get_k8s_token_request_api().""" aiofiles = self._get_aiofiles() - 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( + audience = (await self.adatabricks_conn()).extra_dejson.get("audience", DEFAULT_K8S_AUDIENCE) + expiration_seconds = (await self.adatabricks_conn()).extra_dejson.get("expiration_seconds", 3600) + token_path = (await self.adatabricks_conn()).extra_dejson.get( "k8s_token_path", DEFAULT_K8S_SERVICE_ACCOUNT_TOKEN_PATH ) - namespace_path = self.databricks_conn.extra_dejson.get( + namespace_path = (await self.adatabricks_conn()).extra_dejson.get( "k8s_namespace_path", DEFAULT_K8S_NAMESPACE_PATH ) @@ -959,17 +985,17 @@ 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") + provider = (await self.adatabricks_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") + return subject_token, (await self.adatabricks_conn()).extra_dejson.get("client_id") if self._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") + return subject_token, (await self.adatabricks_conn()).extra_dejson.get("client_id") client_id = self._get_required_client_id() return await self._a_get_k8s_jwt_token(), client_id @@ -1247,43 +1273,43 @@ 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: + if "token" in (await self.adatabricks_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 (await self.adatabricks_conn()).extra_dejson["token"] + if not (await self.adatabricks_conn()).login and (await self.adatabricks_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 (await self.adatabricks_conn()).password + if "azure_tenant_id" in (await self.adatabricks_conn()).extra_dejson: + if (await self.adatabricks_conn()).login == "" or (await self.adatabricks_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 (await self.adatabricks_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 (await self.adatabricks_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 (await self.adatabricks_conn()).extra_dejson.get("service_principal_oauth", False): + if (await self.adatabricks_conn()).login == "" or (await self.adatabricks_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"): + return await self._a_get_sp_token(await self._a_get_oidc_token_service_url()) + if (await self.adatabricks_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()) + return await self._a_get_federated_databricks_token(await self._a_get_oidc_token_service_url()) if self._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( + return await self._a_get_federated_databricks_token(await self._a_get_oidc_token_service_url()) + if (await self.adatabricks_conn()).login == "federated_k8s" or (await self.adatabricks_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()) + return await self._a_get_federated_databricks_token(await self._a_get_oidc_token_service_url()) if raise_error: raise AirflowException("Token authentication isn't configured") @@ -1395,7 +1421,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} @@ -1406,7 +1432,7 @@ 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) + auth = aiohttp.BasicAuth(self._get_connection_attr("login"), (await self.adatabricks_conn()).password) request_func: Any if method == "GET": diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py index 9b0e340ef12e7..1d1e319043c0f 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py @@ -304,6 +304,7 @@ async def test_a_get_sp_token(self, mock_post): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 async with aiohttp.ClientSession() as session: @@ -343,6 +344,7 @@ async def test_a_get_sp_token_retry_error(self, mock_time): mock_conn.password = "client_secret" hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 hook.retry_limit = 3 @@ -635,10 +637,7 @@ def test_get_token_not_configured_raises(self, mock_conn): hook._get_token(raise_error=True) @pytest.mark.asyncio - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_from_extra_dejson(self, mock_conn): extra = {"token": "test_token"} mock_conn.return_value = Connection(extra=extra) @@ -651,10 +650,7 @@ async def test_a_get_token_from_extra_dejson(self, mock_conn): ) @pytest.mark.asyncio - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_token_from_password_when_login_missing(self, mock_conn): mock_conn.return_value = Connection(login=None, password="pw-token") hook = BaseDatabricksHook() @@ -668,10 +664,7 @@ async def test_a_token_from_password_when_login_missing(self, mock_conn): "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_sp_token", new_callable=mock.AsyncMock, ) - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_service_principal_oauth_success(self, mock_conn, mock_get_sp_token): mock_conn.return_value = Connection( host="example.databricks.com", @@ -692,10 +685,7 @@ async def test_a_get_token_service_principal_oauth_success(self, mock_conn, mock "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_aad_token", new_callable=mock.AsyncMock, ) - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_azure_spn_success(self, mock_conn, mock_get_aad_token): extra = {"azure_tenant_id": "tenant_id"} mock_conn.return_value = Connection(login="spn_client_id", password="spn_client_secret", extra=extra) @@ -708,10 +698,7 @@ async def test_a_get_token_azure_spn_success(self, mock_conn, mock_get_aad_token mock_log_debug.assert_called_once_with("Using AAD Token for SPN.") @pytest.mark.asyncio - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_azure_spn_missing_credentials_raises(self, mock_conn): mock_conn.return_value = Connection(login="", password="", extra={"azure_tenant_id": "tenant_id"}) hook = BaseDatabricksHook() @@ -727,10 +714,7 @@ async def test_a_get_token_azure_spn_missing_credentials_raises(self, mock_conn) "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_aad_token", new_callable=mock.AsyncMock, ) - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_managed_identity(self, mock_conn, mock_get_aad_token, mock_check_metadata): mock_conn.return_value = Connection(extra={"use_azure_managed_identity": True}) mock_get_aad_token.return_value = "mi_token" @@ -747,10 +731,7 @@ async def test_a_get_token_managed_identity(self, mock_conn, mock_get_aad_token, "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_aad_token_for_default_az_credential", new_callable=mock.AsyncMock, ) - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_default_azure_credential(self, mock_conn, mock_get_default_cred_token): extra = {DEFAULT_AZURE_CREDENTIAL_SETTING_KEY: True} mock_conn.return_value = Connection(extra=extra) @@ -763,10 +744,7 @@ async def test_a_get_token_default_azure_credential(self, mock_conn, mock_get_de mock_log_debug.assert_called_once_with("Using AzureDefaultCredential for authentication.") @pytest.mark.asyncio - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_service_principal_oauth_missing_credentials(self, mock_conn): mock_conn.return_value = Connection( host="host", login="", password="", extra={"service_principal_oauth": True} @@ -776,10 +754,7 @@ async def test_a_get_token_service_principal_oauth_missing_credentials(self, moc await hook._a_get_token() @pytest.mark.asyncio - @mock.patch( - "airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn", - new_callable=mock.PropertyMock, - ) + @mock.patch("airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.adatabricks_conn") async def test_a_get_token_not_configured_raises(self, mock_conn): mock_conn.return_value = Connection( host="host", @@ -1422,6 +1397,7 @@ async def test_a_get_token_with_supplied_provider(self, mock_post): ) hook = BaseDatabricksHook() hook.databricks_conn = conn + hook._adatabricks_conn = conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 @@ -1879,6 +1855,7 @@ async def test_a_get_token_with_federated_aws(self, mock_post, mock_sts_hook): ) hook = BaseDatabricksHook() hook.databricks_conn = conn + hook._adatabricks_conn = conn hook.user_agent_header = {"User-Agent": "test-agent"} async with aiohttp.ClientSession() as session: @@ -1935,6 +1912,7 @@ async def test_a_get_federated_token(self, _mock_ssl_ctx, mock_post): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 @@ -1963,6 +1941,7 @@ async def test_a_get_federated_token_cached_valid(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn resource = f"https://{mock_conn.host}/oidc/v1/token" # Set expiration far in the future @@ -1988,6 +1967,7 @@ async def test_a_get_federated_token_k8s_not_available(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn resource = f"https://{mock_conn.host}/oidc/v1/token" with mock.patch("aiofiles.open", side_effect=FileNotFoundError()): @@ -2004,6 +1984,7 @@ async def test_a_get_federated_token_missing_client_id(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn resource = f"https://{mock_conn.host}/oidc/v1/token" with pytest.raises( @@ -2056,6 +2037,7 @@ async def test_a_get_federated_token_databricks_error(self, _mock_ssl_ctx, mock_ hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 @@ -2077,6 +2059,7 @@ async def test_a_get_k8s_projected_volume_token_success(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn # Mock aiofiles.open mock_file = mock.AsyncMock() @@ -2097,6 +2080,7 @@ async def test_a_get_k8s_projected_volume_token_file_not_found(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn with mock.patch("aiofiles.open", side_effect=FileNotFoundError()): with pytest.raises( @@ -2113,6 +2097,7 @@ async def test_a_get_k8s_projected_volume_token_permission_denied(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn with mock.patch("aiofiles.open", side_effect=PermissionError()): with pytest.raises( @@ -2129,6 +2114,7 @@ async def test_a_get_k8s_projected_volume_token_empty_file(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn # Mock aiofiles.open with empty content mock_file = mock.AsyncMock() @@ -2150,6 +2136,7 @@ async def test_a_get_k8s_jwt_token_uses_projected_volume_when_configured(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn with mock.patch.object( hook, @@ -2174,6 +2161,7 @@ async def test_a_get_k8s_jwt_token_uses_token_request_api_when_no_projected_path hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn with mock.patch.object( hook, "_a_get_k8s_projected_volume_token", new_callable=mock.AsyncMock @@ -2203,6 +2191,7 @@ async def test_a_get_k8s_token_request_api_uses_ca_cert_for_tls(self, mock_ssl_c mock_conn.extra_dejson = {} hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn mock_response = mock.AsyncMock() mock_response.json = mock.AsyncMock(return_value={"status": {"token": "jwt_token"}}) @@ -2252,6 +2241,7 @@ async def test_a_get_federated_token_with_projected_volume(self): hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} # Mock aiofiles.open for projected volume @@ -2335,6 +2325,7 @@ async def test_a_get_token_with_federated_k8s_login(self, _mock_ssl_ctx, mock_po hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10 @@ -2389,6 +2380,7 @@ async def test_a_get_token_with_federated_k8s_extra(self, _mock_ssl_ctx, mock_po hook = BaseDatabricksHook() hook.databricks_conn = mock_conn + hook._adatabricks_conn = mock_conn hook.user_agent_header = {"User-Agent": "test-agent"} hook.token_timeout_seconds = 10