This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 328730cbcd9 Add Azure Analysis Services model refresh support (#71350)
328730cbcd9 is described below
commit 328730cbcd95e68d416a8d9f0d4335058d2b8802
Author: Aaron Chen <[email protected]>
AuthorDate: Mon Sep 7 16:13:27 2026 +0800
Add Azure Analysis Services model refresh support (#71350)
* Add Azure Analysis Services model refresh support
---
.../azure/docs/connections/analysis_services.rst | 63 +++
.../azure/docs/operators/analysis_services.rst | 85 ++++
providers/microsoft/azure/provider.yaml | 34 ++
.../providers/microsoft/azure/get_provider_info.py | 40 ++
.../microsoft/azure/hooks/analysis_services.py | 300 +++++++++++++
.../microsoft/azure/operators/analysis_services.py | 136 ++++++
.../microsoft/azure/sensors/analysis_services.py | 98 ++++
.../microsoft/azure/triggers/analysis_services.py | 182 ++++++++
.../azure/example_azure_analysis_services.py | 80 ++++
.../azure/hooks/test_analysis_services.py | 500 +++++++++++++++++++++
.../azure/operators/test_analysis_services.py | 218 +++++++++
.../azure/sensors/test_analysis_services.py | 122 +++++
.../azure/triggers/test_analysis_services.py | 336 ++++++++++++++
13 files changed, 2194 insertions(+)
diff --git a/providers/microsoft/azure/docs/connections/analysis_services.rst
b/providers/microsoft/azure/docs/connections/analysis_services.rst
new file mode 100644
index 00000000000..cb750ba444a
--- /dev/null
+++ b/providers/microsoft/azure/docs/connections/analysis_services.rst
@@ -0,0 +1,63 @@
+ .. 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.
+
+.. _howto/connection:azure_analysis_services:
+
+Microsoft Azure Analysis Services Connection
+============================================
+
+The Microsoft Azure Analysis Services connection type enables model refresh
operations through the
+`asynchronous refresh REST API
<https://learn.microsoft.com/en-us/analysis-services/azure-analysis-services/analysis-services-async-refresh>`__.
+
+The
:class:`~airflow.providers.microsoft.azure.hooks.analysis_services.AzureAnalysisServicesHook`,
+:class:`~airflow.providers.microsoft.azure.operators.analysis_services.AzureAnalysisServicesRefreshOperator`,
+and
:class:`~airflow.providers.microsoft.azure.sensors.analysis_services.AzureAnalysisServicesSensor`
+use this connection.
+
+Default Connection ID
+---------------------
+
+The default connection ID is ``azure_analysis_services_default``.
+
+Configuring the Connection
+--------------------------
+
+Region Endpoint
+ Specify the rollout endpoint in ``host``, for example
``westus.asazure.windows.net``. Do not include
+ ``https://``, a port, the Analysis Services server name, or a path. You
can find this endpoint in the
+ server's full name, such as
``asazure://westus.asazure.windows.net/example-server``.
+
+Client ID
+ Specify the Microsoft Entra service principal application (client) ID in
``login``.
+
+Client Secret
+ Specify the service principal client secret in ``password``.
+
+Tenant ID
+ Specify the Microsoft Entra tenant ID in the ``tenantId`` extra field.
+
+Azure Analysis Services currently requires the service principal to be a
server administrator for
+asynchronous refresh REST API calls. Add it to the server administrator role
using the format
+``app:<client-id>@<tenant-id>``. See `Add a service principal to the server
administrator role
+<https://learn.microsoft.com/en-us/analysis-services/azure-analysis-services/analysis-services-addservprinc-admins>`__.
+
+This connection supports service principal client-secret authentication.
Managed identity authentication is
+not supported because Azure Analysis Services does not support managed
identities for these operations.
+
+.. spelling:word-list::
+
+ rollout
diff --git a/providers/microsoft/azure/docs/operators/analysis_services.rst
b/providers/microsoft/azure/docs/operators/analysis_services.rst
new file mode 100644
index 00000000000..c809698a204
--- /dev/null
+++ b/providers/microsoft/azure/docs/operators/analysis_services.rst
@@ -0,0 +1,85 @@
+ .. 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.
+
+Microsoft Azure Analysis Services Operators
+===========================================
+
+Use these components to start and monitor asynchronous model refresh
operations in Microsoft Azure Analysis
+Services. Configure an :ref:`Azure Analysis Services connection
<howto/connection:azure_analysis_services>`
+before using them.
+
+Prerequisite Tasks
+------------------
+
+.. include:: /operators/_partials/prerequisite_tasks.rst
+
+.. _howto/operator:AzureAnalysisServicesRefreshOperator:
+
+Start a model refresh
+---------------------
+
+Use
:class:`~airflow.providers.microsoft.azure.operators.analysis_services.AzureAnalysisServicesRefreshOperator`
+with ``wait_for_termination=False`` to start a refresh and immediately return
its refresh ID.
+
+.. exampleinclude::
/../tests/system/microsoft/azure/example_azure_analysis_services.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_azure_analysis_services_refresh]
+ :end-before: [END howto_operator_azure_analysis_services_refresh]
+
+Wait for the refresh to finish
+------------------------------
+
+With the default ``wait_for_termination=True`` the operator also waits for the
refresh to reach a
+terminal status. The ``timeout`` parameter limits that wait and starts once
the refresh has been
+submitted, while ``request_timeout`` limits each individual REST request.
+
+.. exampleinclude::
/../tests/system/microsoft/azure/example_azure_analysis_services.py
+ :language: python
+ :dedent: 4
+ :start-after: [START
howto_operator_azure_analysis_services_refresh_and_wait]
+ :end-before: [END howto_operator_azure_analysis_services_refresh_and_wait]
+
+.. note::
+
+ The operator and the sensor always run deferred. Both the request that
starts the refresh and
+ the status polling are performed by the triggerer, so no worker slot is
held while the model is
+ refreshing. A triggerer must be running in your deployment.
+
+.. _howto/sensor:AzureAnalysisServicesSensor:
+
+Wait with a sensor
+------------------
+
+The operator return value can be passed directly to
+:class:`~airflow.providers.microsoft.azure.sensors.analysis_services.AzureAnalysisServicesSensor`.
+This separates starting the refresh from waiting for it.
+
+.. exampleinclude::
/../tests/system/microsoft/azure/example_azure_analysis_services.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_azure_analysis_services_refresh]
+ :end-before: [END howto_sensor_azure_analysis_services_refresh]
+
+Azure Analysis Services accepts only one active refresh for a model. A second
request returns HTTP 409, so
+serialize refresh tasks that target the same model, for example with task
dependencies or an Airflow pool.
+
+Reference
+---------
+
+For more information, see `Asynchronous refresh with the REST API
+<https://learn.microsoft.com/en-us/analysis-services/azure-analysis-services/analysis-services-async-refresh>`__.
diff --git a/providers/microsoft/azure/provider.yaml
b/providers/microsoft/azure/provider.yaml
index 4c2f8112330..bb5c770251f 100644
--- a/providers/microsoft/azure/provider.yaml
+++ b/providers/microsoft/azure/provider.yaml
@@ -124,6 +124,12 @@ versions:
- 1.0.0
integrations:
+ - integration-name: Microsoft Azure Analysis Services
+ external-doc-url:
https://learn.microsoft.com/en-us/analysis-services/azure-analysis-services/
+ how-to-guide:
+ -
/docs/apache-airflow-providers-microsoft-azure/operators/analysis_services.rst
+ logo: /docs/integration-logos/Microsoft-Azure.png
+ tags: [azure]
- integration-name: Microsoft Azure Batch
external-doc-url: https://azure.microsoft.com/en-us/services/batch/
how-to-guide:
@@ -215,6 +221,9 @@ integrations:
tags: [azure]
operators:
+ - integration-name: Microsoft Azure Analysis Services
+ python-modules:
+ - airflow.providers.microsoft.azure.operators.analysis_services
- integration-name: Microsoft Azure Compute
python-modules:
- airflow.providers.microsoft.azure.operators.compute
@@ -256,6 +265,9 @@ operators:
- airflow.providers.microsoft.azure.operators.powerbi
sensors:
+ - integration-name: Microsoft Azure Analysis Services
+ python-modules:
+ - airflow.providers.microsoft.azure.sensors.analysis_services
- integration-name: Microsoft Azure Compute
python-modules:
- airflow.providers.microsoft.azure.sensors.compute
@@ -277,6 +289,9 @@ filesystems:
- airflow.providers.microsoft.azure.fs.msgraph
hooks:
+ - integration-name: Microsoft Azure Analysis Services
+ python-modules:
+ - airflow.providers.microsoft.azure.hooks.analysis_services
- integration-name: Microsoft Azure Compute
python-modules:
- airflow.providers.microsoft.azure.hooks.compute
@@ -329,6 +344,9 @@ hooks:
- airflow.providers.microsoft.azure.hooks.powerbi
triggers:
+ - integration-name: Microsoft Azure Analysis Services
+ python-modules:
+ - airflow.providers.microsoft.azure.triggers.analysis_services
- integration-name: Microsoft Azure Batch
python-modules:
- airflow.providers.microsoft.azure.triggers.batch
@@ -389,6 +407,22 @@ transfers:
python-module: airflow.providers.microsoft.azure.transfers.gcs_to_wasb
connection-types:
+ - hook-class-name:
airflow.providers.microsoft.azure.hooks.analysis_services.AzureAnalysisServicesHook
+ hook-name: "Azure Analysis Services"
+ connection-type: azure_analysis_services
+ ui-field-behaviour:
+ hidden-fields: ["schema", "port", "extra"]
+ relabeling:
+ host: Region Endpoint
+ login: Client ID
+ password: Client Secret
+ placeholders:
+ host: westus.asazure.windows.net
+ conn-fields:
+ tenantId:
+ label: Tenant ID
+ schema:
+ type: ["string", "null"]
- hook-class-name:
airflow.providers.microsoft.azure.hooks.base_azure.AzureBaseHook
hook-name: "Azure"
connection-type: azure
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
index 4e3ab2e9d30..f45dd5af9a6 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
@@ -27,6 +27,15 @@ def get_provider_info():
"name": "Microsoft Azure",
"description": "`Microsoft Azure <https://azure.microsoft.com/>`__\n",
"integrations": [
+ {
+ "integration-name": "Microsoft Azure Analysis Services",
+ "external-doc-url":
"https://learn.microsoft.com/en-us/analysis-services/azure-analysis-services/";,
+ "how-to-guide": [
+
"/docs/apache-airflow-providers-microsoft-azure/operators/analysis_services.rst"
+ ],
+ "logo": "/docs/integration-logos/Microsoft-Azure.png",
+ "tags": ["azure"],
+ },
{
"integration-name": "Microsoft Azure Batch",
"external-doc-url":
"https://azure.microsoft.com/en-us/services/batch/";,
@@ -149,6 +158,10 @@ def get_provider_info():
},
],
"operators": [
+ {
+ "integration-name": "Microsoft Azure Analysis Services",
+ "python-modules":
["airflow.providers.microsoft.azure.operators.analysis_services"],
+ },
{
"integration-name": "Microsoft Azure Compute",
"python-modules":
["airflow.providers.microsoft.azure.operators.compute"],
@@ -203,6 +216,10 @@ def get_provider_info():
},
],
"sensors": [
+ {
+ "integration-name": "Microsoft Azure Analysis Services",
+ "python-modules":
["airflow.providers.microsoft.azure.sensors.analysis_services"],
+ },
{
"integration-name": "Microsoft Azure Compute",
"python-modules":
["airflow.providers.microsoft.azure.sensors.compute"],
@@ -229,6 +246,10 @@ def get_provider_info():
"airflow.providers.microsoft.azure.fs.msgraph",
],
"hooks": [
+ {
+ "integration-name": "Microsoft Azure Analysis Services",
+ "python-modules":
["airflow.providers.microsoft.azure.hooks.analysis_services"],
+ },
{
"integration-name": "Microsoft Azure Compute",
"python-modules":
["airflow.providers.microsoft.azure.hooks.compute"],
@@ -299,6 +320,10 @@ def get_provider_info():
},
],
"triggers": [
+ {
+ "integration-name": "Microsoft Azure Analysis Services",
+ "python-modules":
["airflow.providers.microsoft.azure.triggers.analysis_services"],
+ },
{
"integration-name": "Microsoft Azure Batch",
"python-modules":
["airflow.providers.microsoft.azure.triggers.batch"],
@@ -379,6 +404,21 @@ def get_provider_info():
},
],
"connection-types": [
+ {
+ "hook-class-name":
"airflow.providers.microsoft.azure.hooks.analysis_services.AzureAnalysisServicesHook",
+ "hook-name": "Azure Analysis Services",
+ "connection-type": "azure_analysis_services",
+ "ui-field-behaviour": {
+ "hidden-fields": ["schema", "port", "extra"],
+ "relabeling": {
+ "host": "Region Endpoint",
+ "login": "Client ID",
+ "password": "Client Secret",
+ },
+ "placeholders": {"host": "westus.asazure.windows.net"},
+ },
+ "conn-fields": {"tenantId": {"label": "Tenant ID", "schema":
{"type": ["string", "null"]}}},
+ },
{
"hook-class-name":
"airflow.providers.microsoft.azure.hooks.base_azure.AzureBaseHook",
"hook-name": "Azure",
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/analysis_services.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/analysis_services.py
new file mode 100644
index 00000000000..cdce36a4eb8
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/analysis_services.py
@@ -0,0 +1,300 @@
+# 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
+
+TOKEN_SCOPE = "https://*.asazure.windows.net/.default";
+
+RefreshType = Literal["full", "clearValues", "calculate", "dataOnly",
"automatic", "defragment"]
+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}"
+
+
+class AzureAnalysisServicesRefreshStatus:
+ """Azure Analysis Services model refresh statuses."""
+
+ SUCCEEDED = "succeeded"
+ FAILED = "failed"
+ CANCELLED = "cancelled"
+ TIMED_OUT = "timedOut"
+ NOT_STARTED = "notStarted"
+ IN_PROGRESS = "inProgress"
+
+ FAILURE_STATUSES = frozenset({FAILED, CANCELLED, TIMED_OUT})
+ VALID_STATUSES = frozenset({SUCCEEDED, FAILED, CANCELLED, TIMED_OUT,
NOT_STARTED, IN_PROGRESS})
+
+
+class AzureAnalysisServicesRefreshException(AirflowException):
+ """Indicate that an Azure Analysis Services model refresh operation
failed."""
+
+
+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.
+ """
+
+ conn_type: str = "azure_analysis_services"
+ conn_name_attr: str = "azure_analysis_services_conn_id"
+ default_conn_name: str = "azure_analysis_services_default"
+ 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")
+ self.azure_analysis_services_conn_id = azure_analysis_services_conn_id
+ self.request_timeout = request_timeout
+ self._credential: AsyncTokenCredential | None = None
+ self._client: httpx.AsyncClient | None = None
+
+ @cached_property
+ def connection(self) -> Connection:
+ """Return the Azure Analysis Services connection."""
+ return self.get_connection(self.azure_analysis_services_conn_id)
+
+ @classmethod
+ def get_connection_form_widgets(cls) -> dict[str, Any]:
+ """Return connection widgets to add to the connection form."""
+ from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
+ from flask_babel import lazy_gettext
+ from wtforms import StringField
+
+ return {
+ "tenantId": StringField(lazy_gettext("Tenant ID"),
widget=BS3TextFieldWidget()),
+ }
+
+ @classmethod
+ 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",
+ },
+ }
+
+ 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
+
+ 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()
+
+ 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
+
+ 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)
+
+ 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"
+ )
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/analysis_services.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/analysis_services.py
new file mode 100644
index 00000000000..d78738bd8ff
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/analysis_services.py
@@ -0,0 +1,136 @@
+# 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
+
+from collections.abc import Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import BaseOperator
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ AzureAnalysisServicesHook,
+ RefreshType,
+)
+from airflow.providers.microsoft.azure.triggers.analysis_services import (
+ AzureAnalysisServicesRefreshTrigger,
+ validate_completed_refresh_event,
+ validate_refresh_event,
+)
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class AzureAnalysisServicesRefreshOperator(BaseOperator):
+ """
+ Trigger an Azure Analysis Services model refresh and optionally wait for
completion.
+
+ The operator always runs deferred: both the request that starts the
refresh and the status
+ polling happen in the triggerer, so no worker slot is held while the model
is refreshing.
+ A triggerer must therefore be running in the deployment.
+
+ .. seealso::
+ For more information, see
+ :ref:`howto/operator:AzureAnalysisServicesRefreshOperator`.
+
+ :param server_name: The Analysis Services server name.
+ :param database: The model database name.
+ :param azure_analysis_services_conn_id: The Azure Analysis Services
connection ID.
+ :param refresh_type: The processing type to request.
+ :param wait_for_termination: Wait for the refresh to reach a terminal
status.
+ :param check_interval: Time in seconds between status requests.
+ :param timeout: Maximum time in seconds to wait for the refresh to
complete. The clock starts
+ once the refresh has been submitted.
+ :param request_timeout: Timeout in seconds for each HTTP request.
+ """
+
+ template_fields: Sequence[str] = (
+ "azure_analysis_services_conn_id",
+ "server_name",
+ "database",
+ "refresh_type",
+ )
+ ui_color = "#0078d4"
+ ui_fgcolor = "#ffffff"
+
+ def __init__(
+ self,
+ *,
+ server_name: str,
+ database: str,
+ azure_analysis_services_conn_id: str =
AzureAnalysisServicesHook.default_conn_name,
+ refresh_type: RefreshType = "full",
+ wait_for_termination: bool = True,
+ check_interval: float = 60,
+ timeout: float = 60 * 60 * 24 * 7,
+ request_timeout: float = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if check_interval <= 0:
+ raise ValueError("check_interval must be greater than zero")
+ if timeout <= 0:
+ raise ValueError("timeout must be greater than zero")
+ if request_timeout <= 0:
+ raise ValueError("request_timeout must be greater than zero")
+ self.server_name = server_name
+ self.database = database
+ self.azure_analysis_services_conn_id = azure_analysis_services_conn_id
+ self.refresh_type = refresh_type
+ self.wait_for_termination = wait_for_termination
+ self.check_interval = check_interval
+ self.timeout = timeout
+ self.request_timeout = request_timeout
+
+ def execute(self, context: Context) -> None:
+ """Defer to the trigger so the refresh is submitted off the worker."""
+ self.defer(
+ trigger=self._build_trigger(refresh_id=None),
+ method_name=self.handle_refresh.__name__,
+ )
+
+ def handle_refresh(self, context: Context, event: dict[str, Any] | None)
-> str | None:
+ """Record the new refresh ID and defer again when the refresh has to
be awaited."""
+ refresh_id = validate_refresh_event(event)
+ self.log.info("Triggered Azure Analysis Services refresh %s",
refresh_id)
+ context["ti"].xcom_push(key=f"{self.task_id}.refresh_id",
value=refresh_id)
+ if not self.wait_for_termination:
+ return refresh_id
+
+ # The timeout covers waiting for the refresh, so it starts once it has
been submitted.
+ self.defer(
+ timeout=timedelta(seconds=self.timeout),
+ trigger=self._build_trigger(refresh_id=refresh_id),
+ method_name=self.execute_complete.__name__,
+ )
+
+ def _build_trigger(self, *, refresh_id: str | None) ->
AzureAnalysisServicesRefreshTrigger:
+ return AzureAnalysisServicesRefreshTrigger(
+ conn_id=self.azure_analysis_services_conn_id,
+ server_name=self.server_name,
+ database=self.database,
+ refresh_id=refresh_id,
+ refresh_type=self.refresh_type,
+ poke_interval=self.check_interval,
+ request_timeout=self.request_timeout,
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None)
-> str:
+ """Validate the terminal trigger event and return the refresh ID."""
+ refresh_id = validate_completed_refresh_event(event)
+ self.log.info("Azure Analysis Services refresh %s completed
successfully", refresh_id)
+ return refresh_id
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/analysis_services.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/analysis_services.py
new file mode 100644
index 00000000000..2e3ce44ed35
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/analysis_services.py
@@ -0,0 +1,98 @@
+# 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
+
+from collections.abc import Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.microsoft.azure.hooks.analysis_services import
AzureAnalysisServicesHook
+from airflow.providers.microsoft.azure.triggers.analysis_services import (
+ AzureAnalysisServicesRefreshTrigger,
+ validate_completed_refresh_event,
+)
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class AzureAnalysisServicesSensor(BaseSensorOperator):
+ """
+ Wait for an Azure Analysis Services model refresh to finish.
+
+ The sensor always runs deferred: status polling happens in the triggerer,
so no worker slot is
+ held while the model is refreshing. A triggerer must therefore be running
in the deployment.
+
+ .. seealso::
+ For more information, see
+ :ref:`howto/sensor:AzureAnalysisServicesSensor`.
+
+ :param server_name: The Analysis Services server name.
+ :param database: The model database name.
+ :param refresh_id: The refresh operation ID to monitor.
+ :param azure_analysis_services_conn_id: The Azure Analysis Services
connection ID.
+ :param request_timeout: Timeout in seconds for each HTTP request.
+ """
+
+ template_fields: Sequence[str] = (
+ "azure_analysis_services_conn_id",
+ "server_name",
+ "database",
+ "refresh_id",
+ )
+ ui_color = "#0078d4"
+ ui_fgcolor = "#ffffff"
+
+ def __init__(
+ self,
+ *,
+ server_name: str,
+ database: str,
+ refresh_id: str,
+ azure_analysis_services_conn_id: str =
AzureAnalysisServicesHook.default_conn_name,
+ request_timeout: float = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if request_timeout <= 0:
+ raise ValueError("request_timeout must be greater than zero")
+ self.server_name = server_name
+ self.database = database
+ self.refresh_id = refresh_id
+ self.azure_analysis_services_conn_id = azure_analysis_services_conn_id
+ self.request_timeout = request_timeout
+
+ def execute(self, context: Context) -> None:
+ """Defer status polling to the triggerer."""
+ self.defer(
+ timeout=timedelta(seconds=self.timeout),
+ trigger=AzureAnalysisServicesRefreshTrigger(
+ conn_id=self.azure_analysis_services_conn_id,
+ server_name=self.server_name,
+ database=self.database,
+ refresh_id=self.refresh_id,
+ poke_interval=self.poke_interval,
+ request_timeout=self.request_timeout,
+ ),
+ method_name=self.execute_complete.__name__,
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None)
-> None:
+ """Validate the terminal trigger event."""
+ refresh_id = validate_completed_refresh_event(event)
+ self.log.info("Azure Analysis Services refresh %s completed
successfully", refresh_id)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/analysis_services.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/analysis_services.py
new file mode 100644
index 00000000000..666859656fd
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/analysis_services.py
@@ -0,0 +1,182 @@
+# 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
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ AzureAnalysisServicesHook,
+ AzureAnalysisServicesRefreshException,
+ AzureAnalysisServicesRefreshStatus,
+ RefreshType,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+def validate_refresh_event(event: dict[str, Any] | None) -> str:
+ """Validate a trigger event and return its refresh ID."""
+ if not isinstance(event, dict):
+ raise AzureAnalysisServicesRefreshException(
+ "Did not receive a valid event from the Azure Analysis Services
trigger"
+ )
+
+ # Errors are reported before the refresh ID is validated: a failed POST
yields an event
+ # without one, and its message is the only useful diagnostic.
+ event_status = event.get("status")
+ if event_status == "error":
+ message = event.get("message")
+ if not isinstance(message, str) or not message:
+ message = "Azure Analysis Services refresh failed"
+ raise AzureAnalysisServicesRefreshException(message)
+ if event_status != "success":
+ raise AzureAnalysisServicesRefreshException(
+ f"Azure Analysis Services trigger returned unknown event status
{event_status!r}"
+ )
+
+ refresh_id = event.get("refresh_id")
+ if not isinstance(refresh_id, str) or not refresh_id:
+ raise AzureAnalysisServicesRefreshException(
+ "Azure Analysis Services trigger event did not contain a valid
refresh ID"
+ )
+ return refresh_id
+
+
+def validate_completed_refresh_event(event: dict[str, Any] | None) -> str:
+ """Validate a terminal trigger event and return the completed refresh
ID."""
+ refresh_id = validate_refresh_event(event)
+
+ refresh_status = (event or {}).get("refresh_status")
+ if refresh_status != AzureAnalysisServicesRefreshStatus.SUCCEEDED:
+ raise AzureAnalysisServicesRefreshException(
+ f"Azure Analysis Services trigger returned unexpected refresh
status {refresh_status!r}"
+ )
+ return refresh_id
+
+
+class AzureAnalysisServicesRefreshTrigger(BaseTrigger):
+ """
+ Poll an Azure Analysis Services model refresh until it reaches a terminal
status.
+
+ When ``refresh_id`` is ``None`` the trigger starts a new refresh and
yields its ID without
+ polling; the caller defers again with that ID to wait for completion.
Serializing the actual
+ refresh ID is what makes the polling stage survive a triggerer restart.
+
+ :param conn_id: The Azure Analysis Services connection ID.
+ :param server_name: The Analysis Services server name.
+ :param database: The model database name.
+ :param refresh_id: The refresh operation ID to poll, or ``None`` to start
a new refresh.
+ :param refresh_type: The refresh type used when starting a new refresh.
+ :param poke_interval: Time in seconds between status requests.
+ :param request_timeout: Timeout in seconds for each HTTP request.
+ """
+
+ def __init__(
+ self,
+ *,
+ conn_id: str,
+ server_name: str,
+ database: str,
+ refresh_id: str | None = None,
+ refresh_type: RefreshType = "full",
+ poke_interval: float = 60,
+ request_timeout: float = 60,
+ ) -> None:
+ super().__init__()
+ if poke_interval <= 0:
+ raise ValueError("poke_interval must be greater than zero")
+ if request_timeout <= 0:
+ raise ValueError("request_timeout must be greater than zero")
+ self.conn_id = conn_id
+ self.server_name = server_name
+ self.database = database
+ self.refresh_id = refresh_id
+ self.refresh_type = refresh_type
+ self.poke_interval = poke_interval
+ self.request_timeout = request_timeout
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize the trigger arguments and classpath."""
+ return (
+ f"{self.__class__.__module__}.{self.__class__.__name__}",
+ {
+ "conn_id": self.conn_id,
+ "server_name": self.server_name,
+ "database": self.database,
+ "refresh_id": self.refresh_id,
+ "refresh_type": self.refresh_type,
+ "poke_interval": self.poke_interval,
+ "request_timeout": self.request_timeout,
+ },
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Start the refresh when needed, then delegate polling to the hook."""
+ hook = AzureAnalysisServicesHook(
+ azure_analysis_services_conn_id=self.conn_id,
+ request_timeout=self.request_timeout,
+ )
+ refresh_id = self.refresh_id
+ try:
+ if refresh_id is None:
+ refresh_id = await hook.trigger_refresh(
+ server_name=self.server_name,
+ database=self.database,
+ refresh_type=self.refresh_type,
+ )
+ self.log.info("Triggered Azure Analysis Services refresh %s",
refresh_id)
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "refresh_status": None,
+ "message": f"Refresh {refresh_id} has been triggered",
+ "refresh_id": refresh_id,
+ }
+ )
+ return
+
+ status = await hook.wait_for_refresh(
+ server_name=self.server_name,
+ database=self.database,
+ refresh_id=refresh_id,
+ poke_interval=self.poke_interval,
+ )
+ is_success = status == AzureAnalysisServicesRefreshStatus.SUCCEEDED
+ yield TriggerEvent(
+ {
+ "status": "success" if is_success else "error",
+ "refresh_status": status,
+ "message": (
+ f"Refresh {refresh_id} completed successfully"
+ if is_success
+ else f"Refresh {refresh_id} finished with status
{status}"
+ ),
+ "refresh_id": refresh_id,
+ }
+ )
+ except Exception as error:
+ message = str(error) or type(error).__name__
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "refresh_status": None,
+ "message": message,
+ "refresh_id": refresh_id,
+ }
+ )
+ finally:
+ await hook.aclose()
diff --git
a/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_analysis_services.py
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_analysis_services.py
new file mode 100644
index 00000000000..15ebda3bc6d
--- /dev/null
+++
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_analysis_services.py
@@ -0,0 +1,80 @@
+# 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 os
+from datetime import datetime
+from typing import TYPE_CHECKING, cast
+
+from airflow.providers.microsoft.azure.operators.analysis_services import (
+ AzureAnalysisServicesRefreshOperator,
+)
+from airflow.providers.microsoft.azure.sensors.analysis_services import
AzureAnalysisServicesSensor
+from airflow.sdk import DAG
+
+if TYPE_CHECKING:
+ from airflow.providers.microsoft.azure.hooks.analysis_services import
RefreshType
+
+DAG_ID = "example_azure_analysis_services"
+SERVER_NAME = os.environ.get("AZURE_ANALYSIS_SERVICES_SERVER_NAME",
"testserver")
+DATABASE = os.environ.get("AZURE_ANALYSIS_SERVICES_DATABASE", "adventureworks")
+REFRESH_TYPE = cast("RefreshType",
os.environ.get("AZURE_ANALYSIS_SERVICES_REFRESH_TYPE", "calculate"))
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule=None,
+ start_date=datetime(2026, 1, 1),
+ catchup=False,
+ tags=["example", "azure", "analysis-services"],
+) as dag:
+ # [START howto_operator_azure_analysis_services_refresh]
+ start_refresh_without_wait = AzureAnalysisServicesRefreshOperator(
+ task_id="start_refresh_without_wait",
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_type=REFRESH_TYPE,
+ wait_for_termination=False,
+ )
+ # [END howto_operator_azure_analysis_services_refresh]
+
+ # [START howto_sensor_azure_analysis_services_refresh]
+ wait_for_refresh = AzureAnalysisServicesSensor(
+ task_id="wait_for_refresh",
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=start_refresh_without_wait.output,
+ poke_interval=10,
+ timeout=600,
+ )
+ # [END howto_sensor_azure_analysis_services_refresh]
+
+ # [START howto_operator_azure_analysis_services_refresh_and_wait]
+ refresh_and_wait = AzureAnalysisServicesRefreshOperator(
+ task_id="refresh_and_wait",
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_type=REFRESH_TYPE,
+ check_interval=10,
+ timeout=600,
+ )
+ # [END howto_operator_azure_analysis_services_refresh_and_wait]
+
+ start_refresh_without_wait >> wait_for_refresh >> refresh_and_wait
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_analysis_services.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_analysis_services.py
new file mode 100644
index 00000000000..c1b0d57d5ee
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_analysis_services.py
@@ -0,0 +1,500 @@
+# 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 json
+from unittest import mock
+
+import httpx
+import pytest
+from azure.core.credentials import AccessToken
+from azure.core.exceptions import ClientAuthenticationError
+
+from airflow.models import Connection
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ TOKEN_SCOPE,
+ VALID_REFRESH_TYPES,
+ AzureAnalysisServicesHook,
+ AzureAnalysisServicesRefreshException,
+ AzureAnalysisServicesRefreshStatus,
+)
+
+CONN_ID = "azure_analysis_services_test"
+HOST = "westus.asazure.windows.net"
+SERVER_NAME = "testserver"
+DATABASE = "Adventure Works"
+REFRESH_ID = "refresh-id"
+POKE_INTERVAL = 5
+REQUEST_TIMEOUT = 30
+HEADERS = {"Authorization": "Bearer token", "Content-Type": "application/json"}
+MODULE = "airflow.providers.microsoft.azure.hooks.analysis_services"
+
+
+def build_response(
+ *, headers: dict[str, str] | None = None, body: object | None = None,
text: str = ""
+) -> mock.Mock:
+ """Build a response mock with the requested headers and JSON body."""
+ response = mock.Mock(spec=httpx.Response)
+ response.headers = headers or {}
+ response.json.return_value = body
+ response.text = text
+ return response
+
+
+def build_client(client_class: mock.MagicMock) -> mock.MagicMock:
+ """Return the ``httpx.AsyncClient`` instance that the patched class hands
to the hook."""
+ return client_class.return_value
+
+
+class TestAzureAnalysisServicesHook:
+ @pytest.fixture(autouse=True)
+ def setup_connection(self, create_mock_connection):
+ create_mock_connection(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_analysis_services",
+ host=HOST,
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id"},
+ )
+ )
+
+ @pytest.mark.parametrize("request_timeout", [0, -1])
+ def test_rejects_invalid_request_timeout(self, request_timeout):
+ with pytest.raises(ValueError, match="request_timeout must be greater
than zero"):
+ AzureAnalysisServicesHook(CONN_ID, request_timeout=request_timeout)
+
+ def test_defines_connection_form_widget(self):
+ pytest.importorskip("flask_appbuilder")
+ assert set(AzureAnalysisServicesHook.get_connection_form_widgets()) ==
{"tenantId"}
+
+ def test_defines_connection_ui_field_behaviour(self):
+ assert AzureAnalysisServicesHook.get_ui_field_behaviour() == {
+ "hidden_fields": ["schema", "port", "extra"],
+ "relabeling": {
+ "host": "Region Endpoint",
+ "login": "Client ID",
+ "password": "Client Secret",
+ },
+ "placeholders": {"host": "westus.asazure.windows.net"},
+ }
+
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ def test_get_conn_creates_and_caches_client(self, client_class):
+ hook = AzureAnalysisServicesHook(CONN_ID,
request_timeout=REQUEST_TIMEOUT)
+
+ first_client = hook.get_conn()
+ second_client = hook.get_conn()
+
+ client_class.assert_called_once_with(timeout=REQUEST_TIMEOUT)
+ assert first_client is client_class.return_value
+ assert second_client is first_client
+
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ def test_get_credential_creates_and_caches_credential(self,
credential_class):
+ hook = AzureAnalysisServicesHook(CONN_ID)
+
+ first_credential = hook._get_credential()
+ second_credential = hook._get_credential()
+
+ credential_class.assert_called_once_with(
+ tenant_id="tenant-id",
+ client_id="client-id",
+ client_secret="client-secret",
+ )
+ assert first_credential is credential_class.return_value
+ assert second_credential is first_credential
+
+ @mock.patch(f"{MODULE}.BaseHook.get_connection", autospec=True)
+ def test_caches_connection(self, get_connection):
+ get_connection.return_value = Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_analysis_services",
+ host=HOST,
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id"},
+ )
+ hook = AzureAnalysisServicesHook(CONN_ID)
+
+ first_connection = hook.connection
+ second_connection = hook.connection
+
+ get_connection.assert_called_once_with(CONN_ID)
+ assert second_connection is first_connection
+
+ @pytest.mark.parametrize(
+ ("login", "password", "extra", "message"),
+ [
+ (None, "secret", {"tenantId": "tenant"}, "Client ID is required"),
+ ("client", None, {"tenantId": "tenant"}, "Client secret is
required"),
+ ("client", "secret", {}, "Tenant ID is required"),
+ ("client", "secret", {"tenantId": 123}, "Tenant ID is required"),
+ ],
+ )
+ def test_get_credential_requires_service_principal_fields(
+ self, create_mock_connection, login, password, extra, message
+ ):
+ create_mock_connection(
+ Connection(
+ conn_id="invalid-auth",
+ conn_type="azure_analysis_services",
+ host=HOST,
+ login=login,
+ password=password,
+ extra=extra,
+ )
+ )
+
+ with pytest.raises(ValueError, match=message):
+ AzureAnalysisServicesHook("invalid-auth")._get_credential()
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ async def test_get_headers_uses_literal_token_scope(self,
credential_class):
+ credential_class.return_value.get_token.return_value =
AccessToken("token", 0)
+
+ headers = await AzureAnalysisServicesHook(CONN_ID)._get_headers()
+
+
credential_class.return_value.get_token.assert_awaited_once_with(TOKEN_SCOPE)
+ assert TOKEN_SCOPE == "https://*.asazure.windows.net/.default";
+ assert headers == HEADERS
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ async def test_get_headers_wraps_authentication_errors(self,
credential_class):
+ credential_class.return_value.get_token.side_effect =
ClientAuthenticationError("bad credential")
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Failed to authenticate"):
+ await AzureAnalysisServicesHook(CONN_ID)._get_headers()
+
+ @pytest.mark.parametrize(
+ "host",
+ [
+ None,
+ "",
+ "https://westus.asazure.windows.net";,
+ "host/path",
+ "[email protected]",
+ "@evil.example",
+ "evil.example:9999",
+ "host:abc",
+ "evil.example:0",
+ ],
+ )
+ def test_rejects_invalid_region_endpoint(self, host,
create_mock_connection):
+ create_mock_connection(
+ Connection(
+ conn_id="invalid-host",
+ conn_type="azure_analysis_services",
+ host=host,
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id"},
+ )
+ )
+
+ with pytest.raises(ValueError, match="valid region endpoint"):
+ AzureAnalysisServicesHook("invalid-host")._get_base_url()
+
+ def test_strips_trailing_slash_from_region_endpoint(self,
create_mock_connection):
+ create_mock_connection(
+ Connection(
+ conn_id="trailing-slash",
+ conn_type="azure_analysis_services",
+ host=f"{HOST}/",
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id"},
+ )
+ )
+
+ assert AzureAnalysisServicesHook("trailing-slash")._get_base_url() ==
f"https://{HOST}";
+
+ @pytest.mark.parametrize(
+ ("server_name", "database", "message"), [("", "db", "server_name"),
("s", "", "database")]
+ )
+ def test_rejects_empty_resource_names(self, server_name, database,
message):
+ with pytest.raises(ValueError, match=message):
+ AzureAnalysisServicesHook(CONN_ID)._get_refreshes_url(server_name,
database)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("refresh_type", sorted(VALID_REFRESH_TYPES))
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_trigger_refresh_posts_official_request_body(self,
client_class, get_headers, refresh_type):
+ client = build_client(client_class)
+ client.post.return_value = build_response(
+ headers={
+ "Location":
f"https://{HOST}/servers/{SERVER_NAME}/models/Adventure%20Works/refreshes/{REFRESH_ID}";
+ }
+ )
+ hook = AzureAnalysisServicesHook(CONN_ID,
request_timeout=REQUEST_TIMEOUT)
+
+ result = await hook.trigger_refresh(SERVER_NAME, DATABASE,
refresh_type)
+
+ assert result == REFRESH_ID
+ client_class.assert_called_once_with(timeout=REQUEST_TIMEOUT)
+ client.post.assert_awaited_once_with(
+
f"https://{HOST}/servers/{SERVER_NAME}/models/Adventure%20Works/refreshes";,
+ json={"Type": refresh_type},
+ headers=HEADERS,
+ )
+
+ @pytest.mark.asyncio
+ async def test_trigger_refresh_rejects_unknown_type(self):
+ with pytest.raises(ValueError, match="Invalid refresh_type"):
+ await
AzureAnalysisServicesHook(CONN_ID).trigger_refresh(SERVER_NAME, DATABASE,
"invalid")
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "location",
+ [None, "",
f"https://{HOST}/servers/{SERVER_NAME}/models/{DATABASE}/refreshes/";,
"https://[";],
+ )
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_trigger_refresh_rejects_invalid_location(self,
client_class, get_headers, location):
+ build_client(client_class).post.return_value = build_response(
+ headers={"Location": location} if location is not None else {}
+ )
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Location header"):
+ await
AzureAnalysisServicesHook(CONN_ID).trigger_refresh(SERVER_NAME, DATABASE)
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_trigger_refresh_decodes_refresh_id(self, client_class,
get_headers):
+ build_client(client_class).post.return_value = build_response(
+ headers={"Location":
f"https://{HOST}/models/model/refreshes/refresh%20id"}
+ )
+
+ result = await
AzureAnalysisServicesHook(CONN_ID).trigger_refresh(SERVER_NAME, DATABASE)
+
+ assert result == "refresh id"
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("error", [httpx.ConnectError("offline"),
httpx.ReadTimeout("slow")])
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_trigger_refresh_wraps_request_errors(self, client_class,
get_headers, error):
+ build_client(client_class).post.side_effect = error
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Failed to trigger"):
+ await
AzureAnalysisServicesHook(CONN_ID).trigger_refresh(SERVER_NAME, DATABASE)
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def
test_trigger_refresh_includes_truncated_http_error_response(self, client_class,
get_headers):
+ response_text = "A refresh is already running. " + "x" * 1000
+ response = build_response(text=response_text)
+ response.raise_for_status.side_effect = httpx.HTTPStatusError(
+ "409 Client Error: Conflict",
request=mock.Mock(spec=httpx.Request), response=response
+ )
+ build_client(client_class).post.return_value = response
+
+ with pytest.raises(AzureAnalysisServicesRefreshException) as error:
+ await
AzureAnalysisServicesHook(CONN_ID).trigger_refresh(SERVER_NAME, DATABASE)
+
+ assert "409 Client Error: Conflict" in str(error.value)
+ assert f"response body: {response_text[:1000]}" in str(error.value)
+ assert response_text[:1001] not in str(error.value)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("status",
sorted(AzureAnalysisServicesRefreshStatus.VALID_STATUSES))
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_get_refresh_status_returns_valid_status(self, client_class,
get_headers, status):
+ client = build_client(client_class)
+ client.get.return_value = build_response(body={"status": status})
+ hook = AzureAnalysisServicesHook(CONN_ID,
request_timeout=REQUEST_TIMEOUT)
+
+ result = await hook.get_refresh_status(SERVER_NAME, DATABASE, "refresh
id")
+
+ assert result == status
+ client_class.assert_called_once_with(timeout=REQUEST_TIMEOUT)
+ client.get.assert_awaited_once_with(
+
f"https://{HOST}/servers/{SERVER_NAME}/models/Adventure%20Works/refreshes/refresh%20id";,
+ headers=HEADERS,
+ )
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("body", "message"),
+ [
+ ([], "invalid status response"),
+ ({}, "unknown status None"),
+ ({"status": "unexpected"}, "unknown status 'unexpected'"),
+ ({"status": 1}, "unknown status 1"),
+ ],
+ )
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_get_refresh_status_rejects_invalid_response(
+ self, client_class, get_headers, body, message
+ ):
+ build_client(client_class).get.return_value = build_response(body=body)
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match=message):
+ await
AzureAnalysisServicesHook(CONN_ID).get_refresh_status(SERVER_NAME, DATABASE,
REFRESH_ID)
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_get_refresh_status_rejects_non_json_response(self,
client_class, get_headers):
+ response = build_response()
+ response.json.side_effect = json.JSONDecodeError("bad json", "", 0)
+ build_client(client_class).get.return_value = response
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="non-JSON"):
+ await
AzureAnalysisServicesHook(CONN_ID).get_refresh_status(SERVER_NAME, DATABASE,
REFRESH_ID)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("error", [httpx.ConnectError("offline"),
httpx.ReadTimeout("slow")])
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_get_refresh_status_wraps_request_errors(self, client_class,
get_headers, error):
+ build_client(client_class).get.side_effect = error
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Failed to get status"):
+ await
AzureAnalysisServicesHook(CONN_ID).get_refresh_status(SERVER_NAME, DATABASE,
REFRESH_ID)
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_get_refresh_status_includes_http_error_response(self,
client_class, get_headers):
+ response = build_response(text='{"error":"model not found"}')
+ response.raise_for_status.side_effect = httpx.HTTPStatusError(
+ "500 Server Error", request=mock.Mock(spec=httpx.Request),
response=response
+ )
+ build_client(client_class).get.return_value = response
+
+ with pytest.raises(AzureAnalysisServicesRefreshException) as error:
+ await
AzureAnalysisServicesHook(CONN_ID).get_refresh_status(SERVER_NAME, DATABASE,
REFRESH_ID)
+
+ assert "500 Server Error" in str(error.value)
+ assert 'response body: {"error":"model not found"}' in str(error.value)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "status",
+ [
+ AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ *sorted(AzureAnalysisServicesRefreshStatus.FAILURE_STATUSES),
+ ],
+ )
+ @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=mock.AsyncMock)
+ @mock.patch.object(AzureAnalysisServicesHook, "get_refresh_status",
autospec=True)
+ async def test_wait_for_refresh_returns_terminal_status(self,
get_refresh_status, sleep, status):
+ get_refresh_status.return_value = status
+
+ result = await AzureAnalysisServicesHook(CONN_ID).wait_for_refresh(
+ SERVER_NAME, DATABASE, REFRESH_ID, POKE_INTERVAL
+ )
+
+ assert result == status
+ get_refresh_status.assert_awaited_once_with(
+ mock.ANY,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ )
+ sleep.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=mock.AsyncMock)
+ @mock.patch.object(AzureAnalysisServicesHook, "get_refresh_status",
autospec=True)
+ async def test_wait_for_refresh_polls_until_terminal_status(self,
get_refresh_status, sleep):
+ get_refresh_status.side_effect = [
+ AzureAnalysisServicesRefreshStatus.NOT_STARTED,
+ AzureAnalysisServicesRefreshStatus.IN_PROGRESS,
+ AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ ]
+
+ result = await AzureAnalysisServicesHook(CONN_ID).wait_for_refresh(
+ SERVER_NAME, DATABASE, REFRESH_ID, POKE_INTERVAL
+ )
+
+ assert result == AzureAnalysisServicesRefreshStatus.SUCCEEDED
+ assert get_refresh_status.await_count == 3
+ assert sleep.await_args_list == [mock.call(POKE_INTERVAL),
mock.call(POKE_INTERVAL)]
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAnalysisServicesHook, "_get_headers",
autospec=True, return_value=HEADERS)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_reuses_single_async_client_across_calls(self, client_class,
get_headers):
+ client = build_client(client_class)
+ client.get.return_value = build_response(
+ body={"status": AzureAnalysisServicesRefreshStatus.IN_PROGRESS}
+ )
+ hook = AzureAnalysisServicesHook(CONN_ID,
request_timeout=REQUEST_TIMEOUT)
+
+ await hook.get_refresh_status(SERVER_NAME, DATABASE, REFRESH_ID)
+ await hook.get_refresh_status(SERVER_NAME, DATABASE, REFRESH_ID)
+
+ client_class.assert_called_once_with(timeout=REQUEST_TIMEOUT)
+ assert client.get.await_count == 2
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_aclose_closes_client_and_credential(self, client_class,
credential_class):
+ hook = AzureAnalysisServicesHook(CONN_ID)
+ client = hook.get_conn()
+ credential = hook._get_credential()
+
+ await hook.aclose()
+
+ client.aclose.assert_awaited_once_with()
+ credential.close.assert_awaited_once_with()
+ assert hook._client is None
+ assert hook._credential is None
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_aclose_closes_credential_even_if_client_fails(self,
client_class, credential_class):
+ hook = AzureAnalysisServicesHook(CONN_ID)
+ hook.get_conn().aclose.side_effect = RuntimeError("transport already
gone")
+ credential = hook._get_credential()
+
+ with pytest.raises(RuntimeError, match="transport already gone"):
+ await hook.aclose()
+
+ credential.close.assert_awaited_once_with()
+ assert hook._client is None
+ assert hook._credential is None
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ @mock.patch(f"{MODULE}.httpx.AsyncClient", autospec=True)
+ async def test_aclose_is_idempotent(self, client_class, credential_class):
+ hook = AzureAnalysisServicesHook(CONN_ID)
+ client = hook.get_conn()
+ credential = hook._get_credential()
+
+ await hook.aclose()
+ await hook.aclose()
+
+ client.aclose.assert_awaited_once_with()
+ credential.close.assert_awaited_once_with()
+
+ @pytest.mark.asyncio
+ async def test_aclose_is_noop_when_nothing_created(self):
+ await AzureAnalysisServicesHook(CONN_ID).aclose()
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_analysis_services.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_analysis_services.py
new file mode 100644
index 00000000000..7bf216c50a3
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_analysis_services.py
@@ -0,0 +1,218 @@
+# 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
+
+from datetime import timedelta
+from unittest import mock
+
+import pytest
+
+from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ AzureAnalysisServicesHook,
+ AzureAnalysisServicesRefreshException,
+ AzureAnalysisServicesRefreshStatus,
+ RefreshType,
+)
+from airflow.providers.microsoft.azure.operators.analysis_services import (
+ AzureAnalysisServicesRefreshOperator,
+)
+from airflow.providers.microsoft.azure.triggers.analysis_services import (
+ AzureAnalysisServicesRefreshTrigger,
+)
+
+from tests_common.test_utils.operators.run_deferrable import execute_operator
+
+CONN_ID = "azure_analysis_services_test"
+SERVER_NAME = "testserver"
+DATABASE = "adventureworks"
+REFRESH_ID = "refresh-id"
+REFRESH_TYPE: RefreshType = "calculate"
+CHECK_INTERVAL = 5
+TIMEOUT = 120
+REQUEST_TIMEOUT = 30
+
+
+def build_operator(
+ *,
+ wait_for_termination: bool = True,
+ check_interval: float = CHECK_INTERVAL,
+ timeout: float = TIMEOUT,
+ request_timeout: float = REQUEST_TIMEOUT,
+) -> AzureAnalysisServicesRefreshOperator:
+ """Build an operator with standard test arguments."""
+ return AzureAnalysisServicesRefreshOperator(
+ task_id="refresh_model",
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ azure_analysis_services_conn_id=CONN_ID,
+ refresh_type=REFRESH_TYPE,
+ wait_for_termination=wait_for_termination,
+ check_interval=check_interval,
+ timeout=timeout,
+ request_timeout=request_timeout,
+ )
+
+
+def build_context() -> dict:
+ """Build a context whose task instance records XCom pushes."""
+ return {"ti": mock.MagicMock()}
+
+
+def success_event(refresh_status: str | None) -> dict:
+ return {
+ "status": "success",
+ "refresh_status": refresh_status,
+ "message": "ok",
+ "refresh_id": REFRESH_ID,
+ }
+
+
+class TestAzureAnalysisServicesRefreshOperator:
+ @pytest.mark.parametrize(
+ ("check_interval", "timeout", "request_timeout", "message"),
+ [
+ (0, TIMEOUT, REQUEST_TIMEOUT, "check_interval"),
+ (CHECK_INTERVAL, 0, REQUEST_TIMEOUT, "timeout"),
+ (CHECK_INTERVAL, TIMEOUT, 0, "request_timeout"),
+ ],
+ )
+ def test_rejects_invalid_polling_arguments(self, check_interval, timeout,
request_timeout, message):
+ with pytest.raises(ValueError, match=message):
+ build_operator(
+ check_interval=check_interval,
+ timeout=timeout,
+ request_timeout=request_timeout,
+ )
+
+ def test_defines_template_fields(self):
+ assert AzureAnalysisServicesRefreshOperator.template_fields == (
+ "azure_analysis_services_conn_id",
+ "server_name",
+ "database",
+ "refresh_type",
+ )
+
+ def test_execute_defers_without_refresh_id(self):
+ with pytest.raises(TaskDeferred) as deferred:
+ build_operator().execute(context=build_context())
+
+ trigger = deferred.value.trigger
+ assert isinstance(trigger, AzureAnalysisServicesRefreshTrigger)
+ assert trigger.refresh_id is None
+ assert trigger.refresh_type == REFRESH_TYPE
+ assert trigger.poke_interval == CHECK_INTERVAL
+ assert trigger.request_timeout == REQUEST_TIMEOUT
+ assert deferred.value.method_name == "handle_refresh"
+ # The timeout covers waiting for the refresh, so it only starts on the
second deferral.
+ assert deferred.value.timeout is None
+
+ @mock.patch.object(AzureAnalysisServicesHook, "get_refresh_status",
autospec=True)
+ @mock.patch.object(AzureAnalysisServicesHook, "trigger_refresh",
autospec=True)
+ def test_execute_operator_full_lifecycle(self, trigger_refresh,
get_refresh_status):
+ trigger_refresh.return_value = REFRESH_ID
+ get_refresh_status.return_value =
AzureAnalysisServicesRefreshStatus.SUCCEEDED
+
+ result, events = execute_operator(build_operator())
+
+ assert result == REFRESH_ID
+ trigger_refresh.assert_awaited_once_with(
+ mock.ANY,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_type=REFRESH_TYPE,
+ )
+ get_refresh_status.assert_awaited_once_with(
+ mock.ANY,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ )
+ assert [event.payload for event in events] == [
+ {
+ "status": "success",
+ "refresh_status": None,
+ "message": f"Refresh {REFRESH_ID} has been triggered",
+ "refresh_id": REFRESH_ID,
+ },
+ {
+ "status": "success",
+ "refresh_status": AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ "message": f"Refresh {REFRESH_ID} completed successfully",
+ "refresh_id": REFRESH_ID,
+ },
+ ]
+
+ def test_handle_refresh_defers_again_with_refresh_id(self):
+ operator = build_operator()
+
+ with pytest.raises(TaskDeferred) as deferred:
+ operator.handle_refresh(context=build_context(),
event=success_event(None))
+
+ trigger = deferred.value.trigger
+ # The serialised refresh ID is what lets polling survive a triggerer
restart.
+ assert trigger.refresh_id == REFRESH_ID
+ assert trigger.serialize()[1]["refresh_id"] == REFRESH_ID
+ assert deferred.value.method_name == "execute_complete"
+
+ def test_handle_refresh_passes_timeout_to_polling_defer(self):
+ with pytest.raises(TaskDeferred) as deferred:
+ build_operator().handle_refresh(context=build_context(),
event=success_event(None))
+
+ assert deferred.value.timeout == timedelta(seconds=TIMEOUT)
+
+ def test_handle_refresh_pushes_refresh_id_to_xcom(self):
+ operator = build_operator(wait_for_termination=False)
+ context = build_context()
+
+ operator.handle_refresh(context=context, event=success_event(None))
+
+
context["ti"].xcom_push.assert_called_once_with(key="refresh_model.refresh_id",
value=REFRESH_ID)
+
+ def test_handle_refresh_returns_without_waiting(self):
+ result = build_operator(wait_for_termination=False).handle_refresh(
+ context=build_context(), event=success_event(None)
+ )
+
+ assert result == REFRESH_ID
+
+ def test_handle_refresh_raises_on_error_event(self):
+ event = {
+ "status": "error",
+ "refresh_status": None,
+ "message": "Failed to trigger an Azure Analysis Services model
refresh",
+ "refresh_id": None,
+ }
+
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Failed to trigger"):
+ build_operator().handle_refresh(context=build_context(),
event=event)
+
+ def test_execute_complete_returns_refresh_id(self):
+ result = build_operator().execute_complete(
+ context=build_context(),
+ event=success_event(AzureAnalysisServicesRefreshStatus.SUCCEEDED),
+ )
+
+ assert result == REFRESH_ID
+
+ @pytest.mark.parametrize(
+ "event",
+ [None, {"status": "success", "refresh_status":
AzureAnalysisServicesRefreshStatus.FAILED}],
+ )
+ def test_execute_complete_rejects_malformed_event(self, event):
+ with pytest.raises(AzureAnalysisServicesRefreshException):
+ build_operator().execute_complete(context=build_context(),
event=event)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_analysis_services.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_analysis_services.py
new file mode 100644
index 00000000000..a86854aab70
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_analysis_services.py
@@ -0,0 +1,122 @@
+# 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
+
+from datetime import timedelta
+from unittest import mock
+
+import pytest
+
+from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ AzureAnalysisServicesHook,
+ AzureAnalysisServicesRefreshException,
+ AzureAnalysisServicesRefreshStatus,
+)
+from airflow.providers.microsoft.azure.sensors.analysis_services import
AzureAnalysisServicesSensor
+from airflow.providers.microsoft.azure.triggers.analysis_services import (
+ AzureAnalysisServicesRefreshTrigger,
+)
+
+from tests_common.test_utils.operators.run_deferrable import execute_operator
+
+CONN_ID = "azure_analysis_services_test"
+SERVER_NAME = "testserver"
+DATABASE = "adventureworks"
+REFRESH_ID = "refresh-id"
+POKE_INTERVAL = 5
+TIMEOUT = 120
+REQUEST_TIMEOUT = 30
+
+
+def build_sensor(*, request_timeout: float = REQUEST_TIMEOUT) ->
AzureAnalysisServicesSensor:
+ """Build a sensor with standard test arguments."""
+ return AzureAnalysisServicesSensor(
+ task_id="wait_for_refresh",
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ azure_analysis_services_conn_id=CONN_ID,
+ poke_interval=POKE_INTERVAL,
+ timeout=TIMEOUT,
+ request_timeout=request_timeout,
+ )
+
+
+class TestAzureAnalysisServicesSensor:
+ def test_rejects_invalid_request_timeout(self):
+ with pytest.raises(ValueError, match="request_timeout must be greater
than zero"):
+ build_sensor(request_timeout=0)
+
+ def test_defines_template_fields(self):
+ assert AzureAnalysisServicesSensor.template_fields == (
+ "azure_analysis_services_conn_id",
+ "server_name",
+ "database",
+ "refresh_id",
+ )
+
+ def test_execute_defers_with_provided_refresh_id(self):
+ with pytest.raises(TaskDeferred) as deferred:
+ build_sensor().execute(context={})
+
+ trigger = deferred.value.trigger
+ assert isinstance(trigger, AzureAnalysisServicesRefreshTrigger)
+ assert trigger.refresh_id == REFRESH_ID
+ assert trigger.poke_interval == POKE_INTERVAL
+ assert trigger.request_timeout == REQUEST_TIMEOUT
+ assert deferred.value.method_name == "execute_complete"
+
+ def test_execute_passes_sensor_timeout_to_defer(self):
+ with pytest.raises(TaskDeferred) as deferred:
+ build_sensor().execute(context={})
+
+ assert deferred.value.timeout == timedelta(seconds=TIMEOUT)
+
+ @mock.patch.object(AzureAnalysisServicesHook, "get_refresh_status",
autospec=True)
+ def test_execute_sensor_full_lifecycle(self, get_refresh_status):
+ get_refresh_status.return_value =
AzureAnalysisServicesRefreshStatus.SUCCEEDED
+
+ result, events = execute_operator(build_sensor())
+
+ assert result is None
+ get_refresh_status.assert_awaited_once_with(
+ mock.ANY,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ )
+ assert [event.payload for event in events] == [
+ {
+ "status": "success",
+ "refresh_status": AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ "message": f"Refresh {REFRESH_ID} completed successfully",
+ "refresh_id": REFRESH_ID,
+ }
+ ]
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ None,
+ {"status": "success", "refresh_status":
AzureAnalysisServicesRefreshStatus.FAILED},
+ {"status": "error", "message": "refresh failed", "refresh_id":
REFRESH_ID},
+ ],
+ )
+ def test_execute_complete_rejects_malformed_event(self, event):
+ with pytest.raises(AzureAnalysisServicesRefreshException):
+ build_sensor().execute_complete(context={}, event=event)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_analysis_services.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_analysis_services.py
new file mode 100644
index 00000000000..d4a64744eb1
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_analysis_services.py
@@ -0,0 +1,336 @@
+# 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
+
+from unittest import mock
+
+import pytest
+
+from airflow.providers.microsoft.azure.hooks.analysis_services import (
+ AzureAnalysisServicesRefreshException,
+ AzureAnalysisServicesRefreshStatus,
+ RefreshType,
+)
+from airflow.providers.microsoft.azure.triggers.analysis_services import (
+ AzureAnalysisServicesRefreshTrigger,
+ validate_completed_refresh_event,
+ validate_refresh_event,
+)
+from airflow.triggers.base import TriggerEvent
+
+CONN_ID = "azure_analysis_services_test"
+SERVER_NAME = "testserver"
+DATABASE = "adventureworks"
+REFRESH_ID = "refresh-id"
+REFRESH_TYPE: RefreshType = "full"
+POKE_INTERVAL = 5
+REQUEST_TIMEOUT = 30
+MODULE = "airflow.providers.microsoft.azure.triggers.analysis_services"
+
+
+def build_trigger(refresh_id: str | None = REFRESH_ID) ->
AzureAnalysisServicesRefreshTrigger:
+ """Build a trigger with standard test arguments."""
+ return AzureAnalysisServicesRefreshTrigger(
+ conn_id=CONN_ID,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=refresh_id,
+ refresh_type=REFRESH_TYPE,
+ poke_interval=POKE_INTERVAL,
+ request_timeout=REQUEST_TIMEOUT,
+ )
+
+
+class TestAzureAnalysisServicesRefreshTrigger:
+ def test_serializes_all_arguments(self):
+ trigger = build_trigger()
+
+ classpath, arguments = trigger.serialize()
+
+ assert (
+ classpath
+ ==
f"{AzureAnalysisServicesRefreshTrigger.__module__}.AzureAnalysisServicesRefreshTrigger"
+ )
+ assert arguments == {
+ "conn_id": CONN_ID,
+ "server_name": SERVER_NAME,
+ "database": DATABASE,
+ "refresh_id": REFRESH_ID,
+ "refresh_type": REFRESH_TYPE,
+ "poke_interval": POKE_INTERVAL,
+ "request_timeout": REQUEST_TIMEOUT,
+ }
+
+ def test_rejects_invalid_poke_interval(self):
+ with pytest.raises(ValueError, match="poke_interval must be greater
than zero"):
+ AzureAnalysisServicesRefreshTrigger(
+ conn_id=CONN_ID,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ poke_interval=0,
+ )
+
+ def test_rejects_invalid_request_timeout(self):
+ with pytest.raises(ValueError, match="request_timeout must be greater
than zero"):
+ AzureAnalysisServicesRefreshTrigger(
+ conn_id=CONN_ID,
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ request_timeout=0,
+ )
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_emits_success_event(self, hook_class):
+ hook = hook_class.return_value
+ hook.wait_for_refresh.return_value =
AzureAnalysisServicesRefreshStatus.SUCCEEDED
+
+ events = [event async for event in build_trigger().run()]
+
+ hook_class.assert_called_once_with(
+ azure_analysis_services_conn_id=CONN_ID,
+ request_timeout=REQUEST_TIMEOUT,
+ )
+ hook.wait_for_refresh.assert_awaited_once_with(
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_id=REFRESH_ID,
+ poke_interval=POKE_INTERVAL,
+ )
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "success",
+ "refresh_status":
AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ "message": f"Refresh {REFRESH_ID} completed successfully",
+ "refresh_id": REFRESH_ID,
+ }
+ )
+ ]
+
+ @pytest.mark.parametrize("status",
sorted(AzureAnalysisServicesRefreshStatus.FAILURE_STATUSES))
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_emits_error_event_for_failure_status(self, hook_class,
status):
+ hook_class.return_value.wait_for_refresh.return_value = status
+
+ events = [event async for event in build_trigger().run()]
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "refresh_status": status,
+ "message": f"Refresh {REFRESH_ID} finished with status
{status}",
+ "refresh_id": REFRESH_ID,
+ }
+ )
+ ]
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_triggers_refresh_and_yields_without_polling(self,
hook_class):
+ hook = hook_class.return_value
+ hook.trigger_refresh.return_value = REFRESH_ID
+
+ events = [event async for event in
build_trigger(refresh_id=None).run()]
+
+ hook.trigger_refresh.assert_awaited_once_with(
+ server_name=SERVER_NAME,
+ database=DATABASE,
+ refresh_type=REFRESH_TYPE,
+ )
+ hook.wait_for_refresh.assert_not_awaited()
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "success",
+ "refresh_status": None,
+ "message": f"Refresh {REFRESH_ID} has been triggered",
+ "refresh_id": REFRESH_ID,
+ }
+ )
+ ]
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_skips_trigger_when_refresh_id_provided(self, hook_class):
+ hook = hook_class.return_value
+ hook.wait_for_refresh.return_value =
AzureAnalysisServicesRefreshStatus.SUCCEEDED
+
+ [event async for event in build_trigger().run()]
+
+ hook.trigger_refresh.assert_not_awaited()
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_emits_error_event_when_trigger_refresh_fails(self,
hook_class):
+ hook_class.return_value.trigger_refresh.side_effect =
AzureAnalysisServicesRefreshException(
+ "Failed to authenticate with Azure Analysis Services"
+ )
+
+ events = [event async for event in
build_trigger(refresh_id=None).run()]
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "refresh_status": None,
+ "message": "Failed to authenticate with Azure Analysis
Services",
+ "refresh_id": None,
+ }
+ )
+ ]
+
+ @pytest.mark.parametrize("fails", [False, True])
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_closes_hook_on_exit(self, hook_class, fails):
+ hook = hook_class.return_value
+ if fails:
+ hook.wait_for_refresh.side_effect =
AzureAnalysisServicesRefreshException("boom")
+ else:
+ hook.wait_for_refresh.return_value =
AzureAnalysisServicesRefreshStatus.SUCCEEDED
+
+ [event async for event in build_trigger().run()]
+
+ hook.aclose.assert_awaited_once_with()
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_converts_hook_exception_to_error_event(self, hook_class):
+ hook_class.return_value.wait_for_refresh.side_effect =
AzureAnalysisServicesRefreshException(
+ "API unavailable"
+ )
+
+ events = [event async for event in build_trigger().run()]
+
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "error",
+ "refresh_status": None,
+ "message": "API unavailable",
+ "refresh_id": REFRESH_ID,
+ }
+ )
+ ]
+
+ @mock.patch(f"{MODULE}.AzureAnalysisServicesHook", autospec=True)
+ @pytest.mark.asyncio
+ async def test_uses_exception_class_name_when_message_is_empty(self,
hook_class):
+ hook_class.return_value.wait_for_refresh.side_effect = RuntimeError()
+
+ events = [event async for event in build_trigger().run()]
+
+ assert events[0].payload["message"] == "RuntimeError"
+
+
+class TestValidateRefreshEvent:
+ def test_returns_refresh_id_for_success(self):
+ event = {
+ "status": "success",
+ "refresh_status": AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ "message": "completed",
+ "refresh_id": REFRESH_ID,
+ }
+
+ assert validate_refresh_event(event) == REFRESH_ID
+
+ @pytest.mark.parametrize("event", [None, [], "event"])
+ def test_rejects_non_dictionary_event(self, event):
+ with pytest.raises(AzureAnalysisServicesRefreshException, match="valid
event"):
+ validate_refresh_event(event)
+
+ @pytest.mark.parametrize("refresh_id", [None, "", 1])
+ def test_rejects_invalid_refresh_id(self, refresh_id):
+ with pytest.raises(AzureAnalysisServicesRefreshException, match="valid
refresh ID"):
+ validate_refresh_event({"status": "success", "refresh_id":
refresh_id})
+
+ def test_raises_error_event_message(self):
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="refresh failed"):
+ validate_refresh_event(
+ {
+ "status": "error",
+ "refresh_status":
AzureAnalysisServicesRefreshStatus.FAILED,
+ "message": "refresh failed",
+ "refresh_id": REFRESH_ID,
+ }
+ )
+
+ @pytest.mark.parametrize("message", [None, "", 1])
+ def test_raises_fallback_for_error_without_message(self, message):
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="refresh failed"):
+ validate_refresh_event(
+ {
+ "status": "error",
+ "refresh_status":
AzureAnalysisServicesRefreshStatus.FAILED,
+ "message": message,
+ "refresh_id": REFRESH_ID,
+ }
+ )
+
+ def test_raises_error_message_when_refresh_id_absent(self):
+ # A POST that fails before a refresh exists yields no ID; the message
is the only diagnostic.
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="Failed to authenticate"):
+ validate_refresh_event(
+ {
+ "status": "error",
+ "refresh_status": None,
+ "message": "Failed to authenticate with Azure Analysis
Services",
+ "refresh_id": None,
+ }
+ )
+
+ def test_rejects_unknown_event_status(self):
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="unknown event status"):
+ validate_refresh_event({"status": "unknown", "refresh_id":
REFRESH_ID})
+
+ def test_base_validator_accepts_null_refresh_status(self):
+ event = {"status": "success", "refresh_status": None, "refresh_id":
REFRESH_ID}
+
+ assert validate_refresh_event(event) == REFRESH_ID
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="unexpected refresh status"):
+ validate_completed_refresh_event(event)
+
+ def test_completed_validator_returns_refresh_id(self):
+ assert (
+ validate_completed_refresh_event(
+ {
+ "status": "success",
+ "refresh_status":
AzureAnalysisServicesRefreshStatus.SUCCEEDED,
+ "refresh_id": REFRESH_ID,
+ }
+ )
+ == REFRESH_ID
+ )
+
+ @pytest.mark.parametrize(
+ "refresh_status",
+ [None, AzureAnalysisServicesRefreshStatus.IN_PROGRESS,
AzureAnalysisServicesRefreshStatus.FAILED],
+ )
+ def test_completed_validator_rejects_non_success_refresh_status(self,
refresh_status):
+ with pytest.raises(AzureAnalysisServicesRefreshException,
match="unexpected refresh status"):
+ validate_completed_refresh_event(
+ {
+ "status": "success",
+ "refresh_status": refresh_status,
+ "refresh_id": REFRESH_ID,
+ }
+ )
