Skip to content

Fix async connection issue in BaseDatabricksHook #72014

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Paxx2303 wants to merge 1 commit into
base: main
Choose a base branch
from
+357 −50
Open

Fix async connection issue in BaseDatabricksHook #72014

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 92 additions & 50 deletions providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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``).
Expand Down Expand Up @@ -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:
Expand All @@ -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``.
Expand Down Expand Up @@ -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:
Expand All @@ -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],
Expand Down Expand Up @@ -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}
Expand All @@ -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":
Expand Down
Loading
Loading