apache-airflow-providers-microsoft-azure 10.0.0rc1__py3-none-any.whl → 10.1.0__py3-none-any.whl

airflow/providers/microsoft/azure/operators/msgraph.py

@@ -0,0 +1,281 @@
+#
+# 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 copy import deepcopy
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Sequence,
+)
+
+from airflow.exceptions import AirflowException, TaskDeferred
+from airflow.models import BaseOperator
+from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
+from airflow.providers.microsoft.azure.triggers.msgraph import (
+    MSGraphTrigger,
+    ResponseSerializer,
+)
+from airflow.utils.xcom import XCOM_RETURN_KEY
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from kiota_abstractions.request_adapter import ResponseType
+    from kiota_abstractions.request_information import QueryParams
+    from msgraph_core import APIVersion
+
+    from airflow.utils.context import Context
+
+
+class MSGraphAsyncOperator(BaseOperator):
+    """
+    A Microsoft Graph API operator which allows you to execute REST call to the Microsoft Graph API.
+
+    https://learn.microsoft.com/en-us/graph/use-the-api
+
+        .. seealso::
+            For more information on how to use this operator, take a look at the guide:
+            :ref:`howto/operator:MSGraphAsyncOperator`
+
+    :param url: The url being executed on the Microsoft Graph API (templated).
+    :param response_type: The expected return type of the response as a string. Possible value are: `bytes`,
+        `str`, `int`, `float`, `bool` and `datetime` (default is None).
+    :param method: The HTTP method being used to do the REST call (default is GET).
+    :param conn_id: The HTTP Connection ID to run the operator against (templated).
+    :param key: The key that will be used to store `XCom's` ("return_value" is default).
+    :param timeout: The HTTP timeout being used by the `KiotaRequestAdapter` (default is None).
+        When no timeout is specified or set to None then there is no HTTP timeout on each request.
+    :param proxies: A dict defining the HTTP proxies to be used (default is None).
+    :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
+        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
+        or you can pass a string as `v1.0` or `beta`.
+    :param result_processor: Function to further process the response from MS Graph API
+        (default is lambda: context, response: response).  When the response returned by the
+        `KiotaRequestAdapterHook` are bytes, then those will be base64 encoded into a string.
+    :param serializer: Class which handles response serialization (default is ResponseSerializer).
+        Bytes will be base64 encoded into a string, so it can be stored as an XCom.
+    """
+
+    template_fields: Sequence[str] = (
+        "url",
+        "response_type",
+        "path_parameters",
+        "url_template",
+        "query_parameters",
+        "headers",
+        "data",
+        "conn_id",
+    )
+
+    def __init__(
+        self,
+        *,
+        url: str,
+        response_type: ResponseType | None = None,
+        path_parameters: dict[str, Any] | None = None,
+        url_template: str | None = None,
+        method: str = "GET",
+        query_parameters: dict[str, QueryParams] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, Any] | str | BytesIO | None = None,
+        conn_id: str = KiotaRequestAdapterHook.default_conn_name,
+        key: str = XCOM_RETURN_KEY,
+        timeout: float | None = None,
+        proxies: dict | None = None,
+        api_version: APIVersion | None = None,
+        pagination_function: Callable[[MSGraphAsyncOperator, dict], tuple[str, dict]] | None = None,
+        result_processor: Callable[[Context, Any], Any] = lambda context, result: result,
+        serializer: type[ResponseSerializer] = ResponseSerializer,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.url = url
+        self.response_type = response_type
+        self.path_parameters = path_parameters
+        self.url_template = url_template
+        self.method = method
+        self.query_parameters = query_parameters
+        self.headers = headers
+        self.data = data
+        self.conn_id = conn_id
+        self.key = key
+        self.timeout = timeout
+        self.proxies = proxies
+        self.api_version = api_version
+        self.pagination_function = pagination_function or self.paginate
+        self.result_processor = result_processor
+        self.serializer: ResponseSerializer = serializer()
+        self.results: list[Any] | None = None
+
+    def execute(self, context: Context) -> None:
+        self.defer(
+            trigger=MSGraphTrigger(
+                url=self.url,
+                response_type=self.response_type,
+                path_parameters=self.path_parameters,
+                url_template=self.url_template,
+                method=self.method,
+                query_parameters=self.query_parameters,
+                headers=self.headers,
+                data=self.data,
+                conn_id=self.conn_id,
+                timeout=self.timeout,
+                proxies=self.proxies,
+                api_version=self.api_version,
+                serializer=type(self.serializer),
+            ),
+            method_name=self.execute_complete.__name__,
+        )
+
+    def execute_complete(
+        self,
+        context: Context,
+        event: dict[Any, Any] | None = None,
+    ) -> Any:
+        """
+        Execute callback when MSGraphTrigger finishes execution.
+
+        This method gets executed automatically when MSGraphTrigger completes its execution.
+        """
+        self.log.debug("context: %s", context)
+
+        if event:
+            self.log.debug("%s completed with %s: %s", self.task_id, event.get("status"), event)
+
+            if event.get("status") == "failure":
+                raise AirflowException(event.get("message"))
+
+            response = event.get("response")
+
+            self.log.debug("response: %s", response)
+
+            if response:
+                response = self.serializer.deserialize(response)
+
+                self.log.debug("deserialize response: %s", response)
+
+                result = self.result_processor(context, response)
+
+                self.log.debug("processed response: %s", result)
+
+                event["response"] = result
+
+                try:
+                    self.trigger_next_link(response, method_name=self.pull_execute_complete.__name__)
+                except TaskDeferred as exception:
+                    self.append_result(
+                        result=result,
+                        append_result_as_list_if_absent=True,
+                    )
+                    self.push_xcom(context=context, value=self.results)
+                    raise exception
+
+                self.append_result(result=result)
+                self.log.debug("results: %s", self.results)
+
+                return self.results
+        return None
+
+    def append_result(
+        self,
+        result: Any,
+        append_result_as_list_if_absent: bool = False,
+    ):
+        self.log.debug("value: %s", result)
+
+        if isinstance(self.results, list):
+            if isinstance(result, list):
+                self.results.extend(result)
+            else:
+                self.results.append(result)
+        else:
+            if append_result_as_list_if_absent:
+                if isinstance(result, list):
+                    self.results = result
+                else:
+                    self.results = [result]
+            else:
+                self.results = result
+
+    def push_xcom(self, context: Context, value) -> None:
+        self.log.debug("do_xcom_push: %s", self.do_xcom_push)
+        if self.do_xcom_push:
+            self.log.info("Pushing XCom with key '%s': %s", self.key, value)
+            self.xcom_push(context=context, key=self.key, value=value)
+
+    def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = None) -> Any:
+        self.results = list(
+            self.xcom_pull(
+                context=context,
+                task_ids=self.task_id,
+                dag_id=self.dag_id,
+                key=self.key,
+            )
+            or []
+        )
+        self.log.info(
+            "Pulled XCom with task_id '%s' and dag_id '%s' and key '%s': %s",
+            self.task_id,
+            self.dag_id,
+            self.key,
+            self.results,
+        )
+        return self.execute_complete(context, event)
+
+    @staticmethod
+    def paginate(operator: MSGraphAsyncOperator, response: dict) -> tuple[Any, dict[str, Any] | None]:
+        odata_count = response.get("@odata.count")
+        if odata_count and operator.query_parameters:
+            query_parameters = deepcopy(operator.query_parameters)
+            top = query_parameters.get("$top")
+            odata_count = response.get("@odata.count")
+
+            if top and odata_count:
+                if len(response.get("value", [])) == top:
+                    skip = (
+                        sum(map(lambda result: len(result["value"]), operator.results)) + top
+                        if operator.results
+                        else top
+                    )
+                    query_parameters["$skip"] = skip
+                    return operator.url, query_parameters
+        return response.get("@odata.nextLink"), operator.query_parameters
+
+    def trigger_next_link(self, response, method_name="execute_complete") -> None:
+        if isinstance(response, dict):
+            url, query_parameters = self.pagination_function(self, response)
+
+            self.log.debug("url: %s", url)
+            self.log.debug("query_parameters: %s", query_parameters)
+
+            if url:
+                self.defer(
+                    trigger=MSGraphTrigger(
+                        url=url,
+                        query_parameters=query_parameters,
+                        response_type=self.response_type,
+                        conn_id=self.conn_id,
+                        timeout=self.timeout,
+                        proxies=self.proxies,
+                        api_version=self.api_version,
+                        serializer=type(self.serializer),
+                    ),
+                    method_name=method_name,
+                )

airflow/providers/microsoft/azure/sensors/msgraph.py

@@ -0,0 +1,177 @@
+#
+# 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 typing import TYPE_CHECKING, Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
+from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger, ResponseSerializer
+from airflow.sensors.base import BaseSensorOperator
+from airflow.triggers.temporal import TimeDeltaTrigger
+
+if TYPE_CHECKING:
+    from datetime import timedelta
+    from io import BytesIO
+
+    from kiota_abstractions.request_information import QueryParams
+    from kiota_http.httpx_request_adapter import ResponseType
+    from msgraph_core import APIVersion
+
+    from airflow.utils.context import Context
+
+
+class MSGraphSensor(BaseSensorOperator):
+    """
+    A Microsoft Graph API sensor which allows you to poll an async REST call to the Microsoft Graph API.
+
+    :param url: The url being executed on the Microsoft Graph API (templated).
+    :param response_type: The expected return type of the response as a string. Possible value are: `bytes`,
+        `str`, `int`, `float`, `bool` and `datetime` (default is None).
+    :param method: The HTTP method being used to do the REST call (default is GET).
+    :param conn_id: The HTTP Connection ID to run the operator against (templated).
+    :param proxies: A dict defining the HTTP proxies to be used (default is None).
+    :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
+        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
+        or you can pass a string as `v1.0` or `beta`.
+    :param event_processor: Function which checks the response from MS Graph API (default is the
+        `default_event_processor` method) and returns a boolean.  When the result is True, the sensor
+        will stop poking, otherwise it will continue until it's True or times out.
+    :param result_processor: Function to further process the response from MS Graph API
+        (default is lambda: context, response: response).  When the response returned by the
+        `KiotaRequestAdapterHook` are bytes, then those will be base64 encoded into a string.
+    :param serializer: Class which handles response serialization (default is ResponseSerializer).
+        Bytes will be base64 encoded into a string, so it can be stored as an XCom.
+    """
+
+    template_fields: Sequence[str] = (
+        "url",
+        "response_type",
+        "path_parameters",
+        "url_template",
+        "query_parameters",
+        "headers",
+        "data",
+        "conn_id",
+    )
+
+    def __init__(
+        self,
+        url: str,
+        response_type: ResponseType | None = None,
+        path_parameters: dict[str, Any] | None = None,
+        url_template: str | None = None,
+        method: str = "GET",
+        query_parameters: dict[str, QueryParams] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, Any] | str | BytesIO | None = None,
+        conn_id: str = KiotaRequestAdapterHook.default_conn_name,
+        proxies: dict | None = None,
+        api_version: APIVersion | None = None,
+        event_processor: Callable[[Context, Any], bool] = lambda context, e: e.get("status") == "Succeeded",
+        result_processor: Callable[[Context, Any], Any] = lambda context, result: result,
+        serializer: type[ResponseSerializer] = ResponseSerializer,
+        retry_delay: timedelta | float = 60,
+        **kwargs,
+    ):
+        super().__init__(retry_delay=retry_delay, **kwargs)
+        self.url = url
+        self.response_type = response_type
+        self.path_parameters = path_parameters
+        self.url_template = url_template
+        self.method = method
+        self.query_parameters = query_parameters
+        self.headers = headers
+        self.data = data
+        self.conn_id = conn_id
+        self.proxies = proxies
+        self.api_version = api_version
+        self.event_processor = event_processor
+        self.result_processor = result_processor
+        self.serializer = serializer()
+
+    def execute(self, context: Context):
+        self.defer(
+            trigger=MSGraphTrigger(
+                url=self.url,
+                response_type=self.response_type,
+                path_parameters=self.path_parameters,
+                url_template=self.url_template,
+                method=self.method,
+                query_parameters=self.query_parameters,
+                headers=self.headers,
+                data=self.data,
+                conn_id=self.conn_id,
+                timeout=self.timeout,
+                proxies=self.proxies,
+                api_version=self.api_version,
+                serializer=type(self.serializer),
+            ),
+            method_name=self.execute_complete.__name__,
+        )
+
+    def retry_execute(
+        self,
+        context: Context,
+    ) -> Any:
+        self.execute(context=context)
+
+    def execute_complete(
+        self,
+        context: Context,
+        event: dict[Any, Any] | None = None,
+    ) -> Any:
+        """
+        Execute callback when MSGraphSensor finishes execution.
+
+        This method gets executed automatically when MSGraphTrigger completes its execution.
+        """
+        self.log.debug("context: %s", context)
+
+        if event:
+            self.log.debug("%s completed with %s: %s", self.task_id, event.get("status"), event)
+
+            if event.get("status") == "failure":
+                raise AirflowException(event.get("message"))
+
+            response = event.get("response")
+
+            self.log.debug("response: %s", response)
+
+            if response:
+                response = self.serializer.deserialize(response)
+
+                self.log.debug("deserialize response: %s", response)
+
+                is_done = self.event_processor(context, response)
+
+                self.log.debug("is_done: %s", is_done)
+
+                if is_done:
+                    result = self.result_processor(context, response)
+
+                    self.log.debug("processed response: %s", result)
+
+                    return result
+
+                self.defer(
+                    trigger=TimeDeltaTrigger(self.retry_delay),
+                    method_name=self.retry_execute.__name__,
+                )
+
+        return None

airflow/providers/microsoft/azure/triggers/msgraph.py

@@ -0,0 +1,231 @@
+#
+# 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
+import locale
+from base64 import b64encode
+from contextlib import suppress
+from datetime import datetime
+from json import JSONDecodeError
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Sequence,
+)
+from uuid import UUID
+
+import pendulum
+
+from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.utils.module_loading import import_string
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from kiota_abstractions.request_adapter import RequestAdapter
+    from kiota_abstractions.request_information import QueryParams
+    from kiota_http.httpx_request_adapter import ResponseType
+    from msgraph_core import APIVersion
+
+
+class ResponseSerializer:
+    """ResponseSerializer serializes the response as a string."""
+
+    def __init__(self, encoding: str | None = None):
+        self.encoding = encoding or locale.getpreferredencoding()
+
+    def serialize(self, response) -> str | None:
+        def convert(value) -> str | None:
+            if value is not None:
+                if isinstance(value, UUID):
+                    return str(value)
+                if isinstance(value, datetime):
+                    return value.isoformat()
+                if isinstance(value, pendulum.DateTime):
+                    return value.to_iso8601_string()  # Adjust the format as needed
+                raise TypeError(f"Object of type {type(value)} is not JSON serializable!")
+            return None
+
+        if response is not None:
+            if isinstance(response, bytes):
+                return b64encode(response).decode(self.encoding)
+            with suppress(JSONDecodeError):
+                return json.dumps(response, default=convert)
+            return response
+        return None
+
+    def deserialize(self, response) -> Any:
+        if isinstance(response, str):
+            with suppress(JSONDecodeError):
+                response = json.loads(response)
+        return response
+
+
+class MSGraphTrigger(BaseTrigger):
+    """
+    A Microsoft Graph API trigger which allows you to execute an async REST call to the Microsoft Graph API.
+
+    :param url: The url being executed on the Microsoft Graph API (templated).
+    :param response_type: The expected return type of the response as a string. Possible value are: `bytes`,
+        `str`, `int`, `float`, `bool` and `datetime` (default is None).
+    :param method: The HTTP method being used to do the REST call (default is GET).
+    :param conn_id: The HTTP Connection ID to run the operator against (templated).
+    :param timeout: The HTTP timeout being used by the `KiotaRequestAdapter` (default is None).
+        When no timeout is specified or set to None then there is no HTTP timeout on each request.
+    :param proxies: A dict defining the HTTP proxies to be used (default is None).
+    :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
+        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
+        or you can pass a string as `v1.0` or `beta`.
+    :param serializer: Class which handles response serialization (default is ResponseSerializer).
+        Bytes will be base64 encoded into a string, so it can be stored as an XCom.
+    """
+
+    template_fields: Sequence[str] = (
+        "url",
+        "response_type",
+        "path_parameters",
+        "url_template",
+        "query_parameters",
+        "headers",
+        "data",
+        "conn_id",
+    )
+
+    def __init__(
+        self,
+        url: str,
+        response_type: ResponseType | None = None,
+        path_parameters: dict[str, Any] | None = None,
+        url_template: str | None = None,
+        method: str = "GET",
+        query_parameters: dict[str, QueryParams] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, Any] | str | BytesIO | None = None,
+        conn_id: str = KiotaRequestAdapterHook.default_conn_name,
+        timeout: float | None = None,
+        proxies: dict | None = None,
+        api_version: APIVersion | None = None,
+        serializer: type[ResponseSerializer] = ResponseSerializer,
+    ):
+        super().__init__()
+        self.hook = KiotaRequestAdapterHook(
+            conn_id=conn_id,
+            timeout=timeout,
+            proxies=proxies,
+            api_version=api_version,
+        )
+        self.url = url
+        self.response_type = response_type
+        self.path_parameters = path_parameters
+        self.url_template = url_template
+        self.method = method
+        self.query_parameters = query_parameters
+        self.headers = headers
+        self.data = data
+        self.serializer: ResponseSerializer = self.resolve_type(serializer, default=ResponseSerializer)()
+
+    @classmethod
+    def resolve_type(cls, value: str | type, default) -> type:
+        if isinstance(value, str):
+            with suppress(ImportError):
+                return import_string(value)
+            return default
+        return value or default
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize the HttpTrigger arguments and classpath."""
+        api_version = self.api_version.value if self.api_version else None
+        return (
+            f"{self.__class__.__module__}.{self.__class__.__name__}",
+            {
+                "conn_id": self.conn_id,
+                "timeout": self.timeout,
+                "proxies": self.proxies,
+                "api_version": api_version,
+                "serializer": f"{self.serializer.__class__.__module__}.{self.serializer.__class__.__name__}",
+                "url": self.url,
+                "path_parameters": self.path_parameters,
+                "url_template": self.url_template,
+                "method": self.method,
+                "query_parameters": self.query_parameters,
+                "headers": self.headers,
+                "data": self.data,
+                "response_type": self.response_type,
+            },
+        )
+
+    def get_conn(self) -> RequestAdapter:
+        return self.hook.get_conn()
+
+    @property
+    def conn_id(self) -> str:
+        return self.hook.conn_id
+
+    @property
+    def timeout(self) -> float | None:
+        return self.hook.timeout
+
+    @property
+    def proxies(self) -> dict | None:
+        return self.hook.proxies
+
+    @property
+    def api_version(self) -> APIVersion:
+        return self.hook.api_version
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Make a series of asynchronous HTTP calls via a KiotaRequestAdapterHook."""
+        try:
+            response = await self.hook.run(
+                url=self.url,
+                response_type=self.response_type,
+                path_parameters=self.path_parameters,
+                method=self.method,
+                query_parameters=self.query_parameters,
+                headers=self.headers,
+                data=self.data,
+            )
+
+            self.log.debug("response: %s", response)
+
+            if response:
+                response_type = type(response)
+
+                self.log.debug("response type: %s", response_type)
+
+                yield TriggerEvent(
+                    {
+                        "status": "success",
+                        "type": f"{response_type.__module__}.{response_type.__name__}",
+                        "response": self.serializer.serialize(response),
+                    }
+                )
+            else:
+                yield TriggerEvent(
+                    {
+                        "status": "success",
+                        "type": None,
+                        "response": None,
+                    }
+                )
+        except Exception as e:
+            self.log.exception("An error occurred: %s", e)
+            yield TriggerEvent({"status": "failure", "message": str(e)})