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 0000000000000..cb750ba444af0 --- /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 `__. + +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:@``. See `Add a service principal to the server administrator role +`__. + +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 0000000000000..c809698a2044c --- /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 ` +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 +`__. diff --git a/providers/microsoft/azure/provider.yaml b/providers/microsoft/azure/provider.yaml index 88af32c63a30e..3f49fdb1de37f 100644 --- a/providers/microsoft/azure/provider.yaml +++ b/providers/microsoft/azure/provider.yaml @@ -122,6 +122,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: @@ -213,6 +219,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 @@ -254,6 +263,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 @@ -275,6 +287,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 @@ -327,6 +342,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 @@ -387,6 +405,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 34509081d0eb4..d8f3f636dc23f 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 `__\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 0000000000000..6605ef948629d --- /dev/null +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/analysis_services.py @@ -0,0 +1,280 @@ +# 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 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 urllib.parse import SplitResult + + 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", + }, + } + + @property + def client(self) -> httpx.AsyncClient: + """Return the shared HTTP client, creating it on first use.""" + if self._client is None: + self._client = httpx.AsyncClient(timeout=self.request_timeout) + return self._client + + def get_conn(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 = self._client, None + credential, self._credential = 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.client.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 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.client.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]) + + def _assert_host(self, host: str, parsed_host: SplitResult) -> None: + # 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, urlsplit(f"//{host}")) + return f"https://{host}" + + async def _get_headers(self) -> dict[str, str]: + try: + token = await self.get_conn().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 0000000000000..d78738bd8ff41 --- /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 0000000000000..2e3ce44ed3563 --- /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 0000000000000..aa71cf85eef04 --- /dev/null +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/analysis_services.py @@ -0,0 +1,194 @@ +# 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 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 poll it without blocking the triggerer.""" + 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 + + while True: + status = await hook.get_refresh_status( + server_name=self.server_name, + database=self.database, + refresh_id=refresh_id, + ) + self.log.info("Refresh %s status: %s", refresh_id, status) + + if status == AzureAnalysisServicesRefreshStatus.SUCCEEDED: + yield TriggerEvent( + { + "status": "success", + "refresh_status": status, + "message": f"Refresh {refresh_id} completed successfully", + "refresh_id": refresh_id, + } + ) + return + if status in AzureAnalysisServicesRefreshStatus.FAILURE_STATUSES: + yield TriggerEvent( + { + "status": "error", + "refresh_status": status, + "message": f"Refresh {refresh_id} finished with status {status}", + "refresh_id": refresh_id, + } + ) + return + + await asyncio.sleep(self.poke_interval) + 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 0000000000000..15ebda3bc6dc5 --- /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 0000000000000..475b2202f2f87 --- /dev/null +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_analysis_services.py @@ -0,0 +1,444 @@ +# 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" +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}.ClientSecretCredential", autospec=True) + def test_get_conn_creates_and_caches_credential(self, credential_class): + hook = AzureAnalysisServicesHook(CONN_ID) + + first_credential = hook.get_conn() + second_credential = hook.get_conn() + + 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_conn_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_conn() + + @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", + "attacker@evil.example", + "@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 + @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.client + credential = hook.get_conn() + + 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.client.aclose.side_effect = RuntimeError("transport already gone") + credential = hook.get_conn() + + 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.client + credential = hook.get_conn() + + 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 0000000000000..8d459c915cbe3 --- /dev/null +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_analysis_services.py @@ -0,0 +1,179 @@ +# 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 ( + AzureAnalysisServicesRefreshException, + AzureAnalysisServicesRefreshStatus, + RefreshType, +) +from airflow.providers.microsoft.azure.operators.analysis_services import ( + AzureAnalysisServicesRefreshOperator, +) +from airflow.providers.microsoft.azure.triggers.analysis_services import ( + AzureAnalysisServicesRefreshTrigger, +) + +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 + + 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 0000000000000..073adba1978a8 --- /dev/null +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_analysis_services.py @@ -0,0 +1,107 @@ +# 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 + +import pytest + +from airflow.providers.common.compat.sdk import TaskDeferred +from airflow.providers.microsoft.azure.hooks.analysis_services import ( + AzureAnalysisServicesRefreshException, + AzureAnalysisServicesRefreshStatus, +) +from airflow.providers.microsoft.azure.sensors.analysis_services import AzureAnalysisServicesSensor +from airflow.providers.microsoft.azure.triggers.analysis_services import ( + AzureAnalysisServicesRefreshTrigger, +) + +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) + + def test_execute_complete_validates_event(self): + build_sensor().execute_complete( + context={}, + event={ + "status": "success", + "refresh_status": AzureAnalysisServicesRefreshStatus.SUCCEEDED, + "message": "ok", + "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 0000000000000..b2d18989965d9 --- /dev/null +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_analysis_services.py @@ -0,0 +1,351 @@ +# 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.get_refresh_status.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.get_refresh_status.assert_awaited_once_with( + server_name=SERVER_NAME, + database=DATABASE, + refresh_id=REFRESH_ID, + ) + 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.get_refresh_status.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) + @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=mock.AsyncMock) + @pytest.mark.asyncio + async def test_polls_again_after_non_terminal_status(self, sleep, hook_class): + hook = hook_class.return_value + hook.get_refresh_status.side_effect = [ + AzureAnalysisServicesRefreshStatus.IN_PROGRESS, + AzureAnalysisServicesRefreshStatus.SUCCEEDED, + ] + + events = [event async for event in build_trigger().run()] + + sleep.assert_awaited_once_with(POKE_INTERVAL) + assert hook.get_refresh_status.await_count == 2 + assert events[0].payload["status"] == "success" + + @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.get_refresh_status.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.get_refresh_status.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.get_refresh_status.side_effect = AzureAnalysisServicesRefreshException("boom") + else: + hook.get_refresh_status.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.get_refresh_status.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.get_refresh_status.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, + } + )