# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import asyncio
from functools import cached_property
from typing import TYPE_CHECKING, Any, Literal, get_args
from urllib.parse import quote, unquote, urlsplit
import httpx
from azure.core.exceptions import AzureError
from azure.identity.aio import ClientSecretCredential
from airflow.providers.common.compat.sdk import AirflowException, BaseHook
if TYPE_CHECKING:
from azure.core.credentials_async import AsyncTokenCredential
from airflow.sdk import Connection
[docs]
TOKEN_SCOPE = "https://*.asazure.windows.net/.default"
[docs]
RefreshType = Literal["full", "clearValues", "calculate", "dataOnly", "automatic", "defragment"]
[docs]
VALID_REFRESH_TYPES: frozenset[str] = frozenset(get_args(RefreshType))
def _format_request_error(error: httpx.HTTPError) -> str:
# Only HTTPStatusError carries a response; transport errors have no such attribute.
response = getattr(error, "response", None)
response_body = getattr(response, "text", "").strip()[:1000]
response_detail = f"; response body: {response_body}" if response_body else ""
return f"{error}{response_detail}"
[docs]
class AzureAnalysisServicesRefreshStatus:
"""Azure Analysis Services model refresh statuses."""
[docs]
SUCCEEDED = "succeeded"
[docs]
CANCELLED = "cancelled"
[docs]
NOT_STARTED = "notStarted"
[docs]
IN_PROGRESS = "inProgress"
[docs]
FAILURE_STATUSES = frozenset({FAILED, CANCELLED, TIMED_OUT})
[docs]
VALID_STATUSES = frozenset({SUCCEEDED, FAILED, CANCELLED, TIMED_OUT, NOT_STARTED, IN_PROGRESS})
[docs]
class AzureAnalysisServicesRefreshException(AirflowException):
"""Indicate that an Azure Analysis Services model refresh operation failed."""
[docs]
class AzureAnalysisServicesHook(BaseHook):
"""
Interact with the Azure Analysis Services asynchronous refresh REST API.
All request methods are asynchronous and are meant to be awaited from a trigger. Call
:meth:`aclose` when done so the HTTP client and the credential release their resources.
:param azure_analysis_services_conn_id: The Azure Analysis Services connection ID.
:param request_timeout: Timeout in seconds for each HTTP request.
The connection must define the region endpoint in ``host``, the service principal client ID in
``login``, the client secret in ``password``, and the Microsoft Entra tenant ID in the
``tenantId`` extra field.
"""
[docs]
conn_type: str = "azure_analysis_services"
[docs]
conn_name_attr: str = "azure_analysis_services_conn_id"
[docs]
default_conn_name: str = "azure_analysis_services_default"
[docs]
hook_name: str = "Azure Analysis Services"
def __init__(
self,
azure_analysis_services_conn_id: str = default_conn_name,
request_timeout: float = 60,
) -> None:
super().__init__()
if request_timeout <= 0:
raise ValueError("request_timeout must be greater than zero")
[docs]
self.azure_analysis_services_conn_id = azure_analysis_services_conn_id
[docs]
self.request_timeout = request_timeout
self._credential: AsyncTokenCredential | None = None
self._client: httpx.AsyncClient | None = None
@cached_property
[docs]
def connection(self) -> Connection:
"""Return the Azure Analysis Services connection."""
return self.get_connection(self.azure_analysis_services_conn_id)
@classmethod
@classmethod
[docs]
def get_ui_field_behaviour(cls) -> dict[str, Any]:
"""Return custom field behaviour for the connection form."""
return {
"hidden_fields": ["schema", "port", "extra"],
"relabeling": {
"host": "Region Endpoint",
"login": "Client ID",
"password": "Client Secret",
},
"placeholders": {
"host": "westus.asazure.windows.net",
},
}
[docs]
def get_conn(self) -> httpx.AsyncClient:
"""Return and cache the HTTP client used to communicate with Analysis Services."""
if self._client is None:
self._client = httpx.AsyncClient(timeout=self.request_timeout)
return self._client
def _get_credential(self) -> AsyncTokenCredential:
"""Return and cache the service principal credential."""
if self._credential is not None:
return self._credential
connection = self.connection
tenant_id = connection.extra_dejson.get("tenantId")
if not connection.login:
raise ValueError("Client ID is required for Azure Analysis Services authentication")
if not connection.password:
raise ValueError("Client secret is required for Azure Analysis Services authentication")
if not isinstance(tenant_id, str) or not tenant_id:
raise ValueError("Tenant ID is required for Azure Analysis Services authentication")
self._credential = ClientSecretCredential(
tenant_id=tenant_id,
client_id=connection.login,
client_secret=connection.password,
)
return self._credential
[docs]
async def aclose(self) -> None:
"""Release the HTTP client and the credential."""
client = self._client
credential = self._credential
self._client = None
self._credential = None
try:
if client is not None:
await client.aclose()
finally:
if credential is not None:
await credential.close()
[docs]
async def get_refresh_status(self, server_name: str, database: str, refresh_id: str) -> str:
"""Return the validated status of an Azure Analysis Services model refresh."""
refresh_url = f"{self._get_refreshes_url(server_name, database)}/{quote(refresh_id, safe='')}"
try:
response = await self.get_conn().get(refresh_url, headers=await self._get_headers())
response.raise_for_status()
except httpx.HTTPError as error:
raise AzureAnalysisServicesRefreshException(
f"Failed to get status for Azure Analysis Services refresh {refresh_id}: "
f"{_format_request_error(error)}"
) from error
try:
response_body = response.json()
except ValueError as error:
raise AzureAnalysisServicesRefreshException(
f"Azure Analysis Services returned a non-JSON status response for refresh {refresh_id}"
) from error
if not isinstance(response_body, dict):
raise AzureAnalysisServicesRefreshException(
f"Azure Analysis Services returned an invalid status response for refresh {refresh_id}"
)
status = response_body.get("status")
if not isinstance(status, str) or status not in AzureAnalysisServicesRefreshStatus.VALID_STATUSES:
raise AzureAnalysisServicesRefreshException(
f"Azure Analysis Services returned unknown status {status!r} for refresh {refresh_id}"
)
return status
[docs]
async def wait_for_refresh(
self, server_name: str, database: str, refresh_id: str, poke_interval: float
) -> str:
"""Poll until the refresh reaches a terminal status and return it."""
while True:
status = await self.get_refresh_status(
server_name=server_name,
database=database,
refresh_id=refresh_id,
)
self.log.info("Refresh %s status: %s", refresh_id, status)
if (
status == AzureAnalysisServicesRefreshStatus.SUCCEEDED
or status in AzureAnalysisServicesRefreshStatus.FAILURE_STATUSES
):
return status
await asyncio.sleep(poke_interval)
[docs]
async def trigger_refresh(
self, server_name: str, database: str, refresh_type: RefreshType = "full"
) -> str:
"""Trigger a model refresh and return its refresh ID."""
if refresh_type not in VALID_REFRESH_TYPES:
raise ValueError(
f"Invalid refresh_type {refresh_type!r}. Valid values are: {sorted(VALID_REFRESH_TYPES)}"
)
try:
response = await self.get_conn().post(
self._get_refreshes_url(server_name, database),
json={"Type": refresh_type},
headers=await self._get_headers(),
)
response.raise_for_status()
except httpx.HTTPError as error:
raise AzureAnalysisServicesRefreshException(
f"Failed to trigger an Azure Analysis Services model refresh: {_format_request_error(error)}"
) from error
location = response.headers.get("Location")
if not location:
raise AzureAnalysisServicesRefreshException(
"Azure Analysis Services did not return a refresh ID in the Location header"
)
try:
location_parts = [part for part in urlsplit(location).path.split("/") if part]
except ValueError as error:
raise AzureAnalysisServicesRefreshException(
"Azure Analysis Services returned an invalid refresh Location header"
) from error
if len(location_parts) < 2 or location_parts[-2] != "refreshes":
raise AzureAnalysisServicesRefreshException(
"Azure Analysis Services returned an invalid refresh Location header"
)
return unquote(location_parts[-1])
@staticmethod
def _assert_host(host: str) -> None:
parsed_host = urlsplit(f"//{host}")
# netloc, not .username/.port: those miss "@host" and ":0", and raise on ":abc".
if (
not host
or not parsed_host.hostname
or "@" in parsed_host.netloc
or ":" in parsed_host.netloc
or parsed_host.path
or parsed_host.query
or parsed_host.fragment
):
raise ValueError(
"A valid region endpoint without a URL scheme, credentials, port, or path is "
"required in the Azure Analysis Services connection host"
)
def _get_base_url(self) -> str:
host = (self.connection.host or "").strip().rstrip("/")
self._assert_host(host)
return f"https://{host}"
async def _get_headers(self) -> dict[str, str]:
try:
token = await self._get_credential().get_token(TOKEN_SCOPE)
except AzureError as error:
raise AzureAnalysisServicesRefreshException(
"Failed to authenticate with Azure Analysis Services"
) from error
return {
"Authorization": f"Bearer {token.token}",
"Content-Type": "application/json",
}
def _get_refreshes_url(self, server_name: str, database: str) -> str:
if not server_name:
raise ValueError("server_name must not be empty")
if not database:
raise ValueError("database must not be empty")
return (
f"{self._get_base_url()}/servers/{quote(server_name, safe='')}"
f"/models/{quote(database, safe='')}/refreshes"
)