From f82717c9587d12c027e061380410d0005385efdf Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 13 Mar 2024 14:27:34 +0100 Subject: [PATCH 001/105] refactor: Initial commit contains the new MSGraphOperator --- .../microsoft/azure/hooks/msgraph.py | 174 +++++++++++ .../microsoft/azure/operators/msgraph.py | 271 ++++++++++++++++++ .../microsoft/azure/serialization/__init__.py | 0 .../azure/serialization/response_handler.py | 21 ++ .../azure/serialization/serializer.py | 43 +++ .../microsoft/azure/triggers/msgraph.py | 256 +++++++++++++++++ pyproject.toml | 1 + tests/providers/microsoft/azure/base.py | 102 +++++++ .../microsoft/azure/hooks/test_msgraph.py | 77 +++++ .../microsoft/azure/operators/test_msgraph.py | 129 +++++++++ .../microsoft/azure/resources/dummy.pdf | Bin 0 -> 13264 bytes .../microsoft/azure/resources/next_users.json | 1 + .../microsoft/azure/resources/users.json | 1 + .../microsoft/azure/serializer/__init__.py | 0 .../azure/serializer/test_serializer.py | 46 +++ .../microsoft/azure/triggers/test_trigger.py | 123 ++++++++ tests/providers/microsoft/conftest.py | 63 +++- 17 files changed, 1307 insertions(+), 1 deletion(-) create mode 100644 airflow/providers/microsoft/azure/hooks/msgraph.py create mode 100644 airflow/providers/microsoft/azure/operators/msgraph.py create mode 100644 airflow/providers/microsoft/azure/serialization/__init__.py create mode 100644 airflow/providers/microsoft/azure/serialization/response_handler.py create mode 100644 airflow/providers/microsoft/azure/serialization/serializer.py create mode 100644 airflow/providers/microsoft/azure/triggers/msgraph.py create mode 100644 tests/providers/microsoft/azure/base.py create mode 100644 tests/providers/microsoft/azure/hooks/test_msgraph.py create mode 100644 tests/providers/microsoft/azure/operators/test_msgraph.py create mode 100644 tests/providers/microsoft/azure/resources/dummy.pdf create mode 100644 tests/providers/microsoft/azure/resources/next_users.json create mode 100644 tests/providers/microsoft/azure/resources/users.json create mode 100644 tests/providers/microsoft/azure/serializer/__init__.py create mode 100644 tests/providers/microsoft/azure/serializer/test_serializer.py create mode 100644 tests/providers/microsoft/azure/triggers/test_trigger.py diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py new file mode 100644 index 0000000000000..67d89a05606fd --- /dev/null +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -0,0 +1,174 @@ +# +# 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 typing import Dict, Optional, Union, TYPE_CHECKING, Tuple +from urllib.parse import urljoin + +import httpx +from azure import identity +from httpx import Timeout +from kiota_authentication_azure import azure_identity_authentication_provider +from kiota_http import httpx_request_adapter +from msgraph_core import GraphClientFactory +from msgraph_core._enums import APIVersion, NationalClouds + +from airflow.exceptions import AirflowException +from airflow.hooks.base import BaseHook + +if TYPE_CHECKING: + from airflow.models import Connection + from kiota_abstractions.request_adapter import RequestAdapter + + +class KiotaRequestAdapterHook(BaseHook): + """ + A Microsoft Graph API interaction hook, a Wrapper around KiotaRequestAdapter. + + https://github.com/microsoftgraph/msgraph-sdk-python-core + + :param conn_id: The HTTP Connection ID to run the trigger against. + :param timeout: The HTTP timeout being used by the KiotaRequestAdapter (default is None). + When no timeout is specified or set to None then no HTTP timeout is applied 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". + """ + + cached_request_adapters: Dict[str, Tuple[APIVersion, RequestAdapter]] = {} + default_conn_name: str = "msgraph_default" + + def __init__( + self, + conn_id: str = default_conn_name, + timeout: Optional[float] = None, + proxies: Optional[Dict] = None, + api_version: Optional[Union[APIVersion, str]] = None, + ) -> None: + self.conn_id = conn_id + self.timeout = timeout + self.proxies = proxies + self._api_version = self.resolve_api_version_from_value(api_version) + + @property + def api_version(self) -> APIVersion: + self.get_conn() # Make sure config has been loaded through get_conn to have correct api version! + return self._api_version + + @staticmethod + def resolve_api_version_from_value( + api_version: Union[APIVersion, str], default: Optional[APIVersion] = None + ) -> APIVersion: + if isinstance(api_version, APIVersion): + return api_version + return next( + filter(lambda version: version.value == api_version, APIVersion), + default, + ) + + def get_api_version(self, config: Dict) -> APIVersion: + if self._api_version is None: + return self.resolve_api_version_from_value( + api_version=config.get("api_version"), default=APIVersion.v1 + ) + return self._api_version + + @staticmethod + def get_host(connection: Connection) -> str: + if connection.schema and connection.host: + return f"{connection.schema}://{connection.host}" + return NationalClouds.Global.value + + @staticmethod + def to_httpx_proxies(proxies: Dict) -> Dict: + proxies = proxies.copy() + if proxies.get("http"): + proxies["http://"] = proxies.pop("http") + if proxies.get("https"): + proxies["https://"] = proxies.pop("https") + return proxies + + def get_conn(self) -> RequestAdapter: + if not self.conn_id: + raise AirflowException( + "Failed to create the KiotaRequestAdapterHook. No conn_id provided!" + ) + + api_version, request_adapter = self.cached_request_adapters.get( + self.conn_id, (None, None) + ) + + if not request_adapter: + connection = self.get_connection(conn_id=self.conn_id) + client_id = connection.login + client_secret = connection.password + config = connection.extra_dejson if connection.extra else {} + tenant_id = config.get("tenant_id") + api_version = self.get_api_version(config) + host = self.get_host(connection) + base_url = config.get("base_url", urljoin(host, api_version.value)) + proxies = self.proxies or config.get("proxies", {}) + scopes = config.get("scopes", ["https://graph.microsoft.com/.default"]) + verify = config.get("verify", True) + trust_env = config.get("trust_env", False) + + self.log.info( + "Creating Microsoft Graph SDK client %s for conn_id: %s", + api_version.value, + self.conn_id, + ) + self.log.info("Host: %s", host) + self.log.info("Base URL: %s", base_url) + self.log.info("Tenant id: %s", tenant_id) + self.log.info("Client id: %s", client_id) + self.log.info("Client secret: %s", client_secret) + self.log.info("API version: %s", api_version.value) + self.log.info("Scope: %s", scopes) + self.log.info("Verify: %s", verify) + self.log.info("Timeout: %s", self.timeout) + self.log.info("Trust env: %s", trust_env) + self.log.info("Proxies: %s", json.dumps(proxies)) + credentials = identity.ClientSecretCredential( + tenant_id=tenant_id, + client_id=connection.login, + client_secret=connection.password, + proxies=proxies, + ) + http_client = GraphClientFactory.create_with_default_middleware( + api_version=api_version, + client=httpx.AsyncClient( + proxies=self.to_httpx_proxies(proxies), + timeout=Timeout(timeout=self.timeout), + verify=verify, + trust_env=trust_env, + ), + host=host, + ) + auth_provider = azure_identity_authentication_provider.AzureIdentityAuthenticationProvider( + credentials=credentials, scopes=scopes + ) + request_adapter = httpx_request_adapter.HttpxRequestAdapter( + authentication_provider=auth_provider, + http_client=http_client, + base_url=base_url, + ) + self.cached_request_adapters[self.conn_id] = (api_version, request_adapter) + self._api_version = api_version + return request_adapter diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py new file mode 100644 index 0000000000000..f7ffa910579f7 --- /dev/null +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -0,0 +1,271 @@ +# +# 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 ( + Dict, + Optional, + Any, + TYPE_CHECKING, + Sequence, + Callable, + Type, +) + +from airflow import AirflowException +from airflow.exceptions import TaskDeferred +from airflow.models import BaseOperator +from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook +from airflow.providers.microsoft.azure.serialization.serializer import ( + ResponseSerializer, +) +from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger +from airflow.utils.xcom import XCOM_RETURN_KEY + +if TYPE_CHECKING: + from msgraph_core import APIVersion + from io import BytesIO + from airflow.utils.context import Context + from kiota_abstractions.request_adapter import ResponseType + from kiota_abstractions.request_information import QueryParams + from kiota_abstractions.response_handler import NativeResponseType + from kiota_abstractions.serialization import ParsableFactory + + +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 + + :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 url: The url being executed on the Microsoft Graph API (templated). + :param timeout: The HTTP timeout being used by the KiotaRequestAdapter (default is None). + When no timeout is specified or set to None then no HTTP timeout is applied 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 + GraphServiceClientHook are bytes, then those will be base64 encoded into a string. + :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 response_handler: Function to convert the native HTTPX response returned by the hook (default is + lambda response, error_map: response.json()). The default expression will convert the native response + to JSON. If response_type parameter is specified, then the response_handler will be ignored. + :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", "conn_id") + + def __init__( + self, + *, + url: Optional[str] = None, + response_type: Optional[ResponseType] = None, + response_handler: Callable[ + [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + ] = lambda response, error_map: response.json(), + path_parameters: Optional[Dict[str, Any]] = None, + url_template: Optional[str] = None, + method: str = "GET", + query_parameters: Optional[Dict[str, QueryParams]] = None, + headers: Optional[Dict[str, str]] = None, + content: Optional[BytesIO] = None, + conn_id: str = KiotaRequestAdapterHook.default_conn_name, + key: str = XCOM_RETURN_KEY, + timeout: Optional[float] = None, + proxies: Optional[Dict] = None, + api_version: Optional[APIVersion] = None, + result_processor: Callable[ + [Context, Any], Any + ] = lambda context, result: result, + serializer: Type[ResponseSerializer] = ResponseSerializer, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.url = url + self.response_type = response_type + self.response_handler = response_handler + self.path_parameters = path_parameters + self.url_template = url_template + self.method = method + self.query_parameters = query_parameters + self.headers = headers + self.content = content + self.conn_id = conn_id + self.key = key + self.timeout = timeout + self.proxies = proxies + self.api_version = api_version + self.result_processor = result_processor + self.serializer: ResponseSerializer = serializer() + self.results = None + + def execute(self, context: Context) -> None: + self.log.info("Executing url '%s' as '%s'", self.url, self.method) + 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, + content=self.content, + conn_id=self.conn_id, + timeout=self.timeout, + proxies=self.proxies, + api_version=self.api_version, + serializer=type(self.serializer), + ), + method_name="execute_complete", + ) + + def execute_complete( + self, + context: Context, + event: Optional[Dict[Any, Any]] = None, + ) -> Any: + """ + Callback for when the trigger fires - returns immediately. + Relies on trigger to throw an exception, otherwise it assumes execution was + successful. + """ + self.log.debug("context: %s", context) + + if event: + self.log.info( + "%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.info("response: %s", response) + + if response: + self.log.debug("response type: %s", type(response)) + + response = self.serializer.deserialize(response) + + self.log.debug("deserialized response type: %s", type(response)) + + result = self.result_processor(context, response) + + self.log.debug("processed response: %s", result) + + event["response"] = result + + self.log.debug("parsed response type: %s", type(response)) + + try: + self.trigger_next_link( + response, method_name="pull_execute_complete" + ) + 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: Optional[Dict[Any, Any]] = None + ) -> Any: + self.results = list( + self.xcom_pull( + context=context, + task_ids=self.task_id, + dag_id=self.dag_id, + key=self.key, + ) + or [] # noqa: W503 + ) + 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) + + def trigger_next_link( + self, response, method_name="execute_complete" + ) -> None: + if isinstance(response, dict): + odata_next_link = response.get("@odata.nextLink") + + self.log.debug("odata_next_link: %s", odata_next_link) + + if odata_next_link: + self.defer( + trigger=MSGraphTrigger( + url=odata_next_link, + response_type=self.response_type, + response_handler=self.response_handler, + conn_id=self.conn_id, + timeout=self.timeout, + proxies=self.proxies, + api_version=self.api_version, + serializer=type(self.serializer), + ), + method_name=method_name, + ) diff --git a/airflow/providers/microsoft/azure/serialization/__init__.py b/airflow/providers/microsoft/azure/serialization/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py new file mode 100644 index 0000000000000..71dfbcdec2c8a --- /dev/null +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -0,0 +1,21 @@ +from typing import Optional, Dict, Any, Callable + +from kiota_abstractions.response_handler import ResponseHandler, NativeResponseType +from kiota_abstractions.serialization import ParsableFactory # noqa: TC002 + + +class CallableResponseHandler(ResponseHandler): + def __init__( + self, + callable_function: Callable[ + [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + ], + ): + self.callable_function = callable_function + + async def handle_response_async( + self, + response: NativeResponseType, + error_map: Optional[Dict[str, Optional[ParsableFactory]]], + ) -> Any: + return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py new file mode 100644 index 0000000000000..95b51937bf6da --- /dev/null +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -0,0 +1,43 @@ +import json +import locale +from base64 import b64encode +from contextlib import suppress +from datetime import datetime +from json import JSONDecodeError +from typing import Optional, Any +from uuid import UUID + +import pendulum + + +class ResponseSerializer: + def __init__(self, encoding: Optional[str] = None): + self.encoding = encoding or locale.getpreferredencoding() + + def serialize(self, response) -> Optional[str]: + def convert(value) -> Optional[str]: + 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 diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py new file mode 100644 index 0000000000000..07cc038a5150e --- /dev/null +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -0,0 +1,256 @@ +# +# 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 ( + Dict, + Optional, + Any, + AsyncIterator, + Sequence, + Union, + Type, + TYPE_CHECKING, + Callable, +) + +from kiota_abstractions.api_error import APIError +from kiota_abstractions.method import Method +from kiota_abstractions.request_information import RequestInformation +from kiota_http.middleware.options import ResponseHandlerOption +from msgraph_core import APIVersion + +from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook +from airflow.providers.microsoft.azure.serialization.response_handler import ( + CallableResponseHandler, +) +from airflow.providers.microsoft.azure.serialization.serializer import ( + ResponseSerializer, +) +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_abstractions.response_handler import NativeResponseType + from kiota_abstractions.serialization import ParsableFactory + from kiota_http.httpx_request_adapter import ResponseType + + +class MSGraphTrigger(BaseTrigger): + """ + A Microsoft Graph API trigger which allows you to execute an async REST call to the Microsoft Graph API. + + https://github.com/microsoftgraph/msgraph-sdk-python + + :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 trigger 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 no HTTP timeout is applied 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". + """ + + DEFAULT_HEADERS = {"Accept": "application/json;q=1"} + template_fields: Sequence[str] = ( + "url", + "response_type", + "path_parameters", + "url_template", + "query_parameters", + "headers", + "content", + "conn_id", + ) + + def __init__( + self, + url: Optional[str] = None, + response_type: Optional[ResponseType] = None, + response_handler: Callable[ + [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + ] = lambda response, error_map: response.json(), + path_parameters: Optional[Dict[str, Any]] = None, + url_template: Optional[str] = None, + method: str = "GET", + query_parameters: Optional[Dict[str, QueryParams]] = None, + headers: Optional[Dict[str, str]] = None, + content: Optional[BytesIO] = None, + conn_id: str = KiotaRequestAdapterHook.default_conn_name, + timeout: Optional[float] = None, + proxies: Optional[Dict] = None, + api_version: Union[APIVersion, str] = APIVersion.v1, + serializer: Union[str, 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.response_handler = response_handler + self.path_parameters = path_parameters + self.url_template = url_template + self.method = method + self.query_parameters = query_parameters + self.headers = headers + self.content = content + self.serializer: ResponseSerializer = self.resolve_type( + serializer, default=ResponseSerializer + )() + + @classmethod + def resolve_type(cls, value: Union[str, Type], default) -> Type: + if isinstance(value, str): + try: + return import_string(value) + except ImportError: + return default + return value or default + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serializes 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, + "content": self.content, + "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) -> Optional[float]: + return self.hook.timeout + + @property + def proxies(self) -> Optional[Dict]: + return self.hook.proxies + + @property + def api_version(self) -> APIVersion: + return self.hook.api_version + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Makes a series of asynchronous http calls via a KiotaRequestAdapterHook.""" + try: + response = await self.execute() + + self.log.debug("response: %s", response) + + if response: + response_type = type(response) + + self.log.debug("response type: %s", type(response)) + + response = self.serializer.serialize(response) + + self.log.debug("serialized response type: %s", type(response)) + + yield TriggerEvent( + { + "status": "success", + "type": f"{response_type.__module__}.{response_type.__name__}", + "response": 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)}) + + def normalize_url(self) -> str: + if self.url.startswith("/"): + return self.url.replace("/", "", 1) + return self.url + + def request_information(self) -> RequestInformation: + request_information = RequestInformation() + if self.url.startswith("http"): + request_information.url = self.url + else: + request_information.url_template = f"{{+baseurl}}/{self.normalize_url()}" + request_information.path_parameters = self.path_parameters or {} + request_information.http_method = Method(self.method.strip().upper()) + request_information.query_parameters = self.query_parameters or {} + if not self.response_type: + request_information.request_options[ + ResponseHandlerOption.get_key() + ] = ResponseHandlerOption( + response_handler=CallableResponseHandler(self.response_handler) + ) + request_information.content = self.content + headers = ( + {**self.DEFAULT_HEADERS, **self.headers} + if self.headers + else self.DEFAULT_HEADERS + ) + for header_name, header_value in headers.items(): + request_information.headers.try_add( + header_name=header_name, header_value=header_value + ) + return request_information + + @staticmethod + def error_mapping() -> Dict[str, Optional[ParsableFactory]]: + return { + "4XX": APIError, + "5XX": APIError, + } + + async def execute(self) -> AsyncIterator[TriggerEvent]: + return await self.get_conn().send_primitive_async( + request_info=self.request_information(), + response_type=self.response_type, + error_map=self.error_mapping(), + ) diff --git a/pyproject.toml b/pyproject.toml index 63df5c74fb45a..ab962240e76da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -819,6 +819,7 @@ microsoft-azure = [ # source: airflow/providers/microsoft/azure/provider.yaml "azure-storage-file-share", "azure-synapse-artifacts>=0.17.0", "azure-synapse-spark", + "msgraph-core>=1.0.0", # Devel dependencies for the microsoft.azure provider "pywinrm", ] diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py new file mode 100644 index 0000000000000..8486eae988c5b --- /dev/null +++ b/tests/providers/microsoft/azure/base.py @@ -0,0 +1,102 @@ +# 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. +import asyncio +from copy import deepcopy +from datetime import datetime +from typing import List, Tuple, Any, Iterable, Union, Optional + +import pytest +from sqlalchemy.orm import Session + +from airflow.exceptions import TaskDeferred +from airflow.models import Operator, TaskInstance +from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook +from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.utils.session import NEW_SESSION +from airflow.utils.state import TaskInstanceState +from airflow.utils.xcom import XCOM_RETURN_KEY + + +class MockedTaskInstance(TaskInstance): + values = {} + + def xcom_pull( + self, + task_ids: Optional[Union[Iterable[str], str]] = None, + dag_id: Optional[str] = None, + key: str = XCOM_RETURN_KEY, + include_prior_dates: bool = False, + session: Session = NEW_SESSION, + *, + map_indexes: Optional[Union[Iterable[int], int]] = None, + default: Optional[Any] = None, + ) -> Any: + self.task_id = task_ids + self.dag_id = dag_id + return self.values.get(f"{task_ids}_{dag_id}_{key}") + + def xcom_push( + self, + key: str, + value: Any, + execution_date: Optional[datetime] = None, + session: Session = NEW_SESSION, + ) -> None: + self.values[f"{self.task_id}_{self.dag_id}_{key}"] = value + + +class Base: + _loop = asyncio.get_event_loop() + + def teardown_method(self, method): + KiotaRequestAdapterHook.cached_request_adapters.clear() + MockedTaskInstance.values.clear() + + @staticmethod + async def run_tigger(trigger: BaseTrigger) -> List[TriggerEvent]: + events = [] + async for event in trigger.run(): + events.append(event) + return events + + def execute_operator(self, operator: Operator) -> Tuple[Any, Any]: + task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) + context = {"ti": task_instance} + result = None + triggered_events = [] + + with pytest.raises(TaskDeferred) as deferred: + operator.execute(context=context) + + task = deferred.value + + while task: + events = self._loop.run_until_complete(self.run_tigger(deferred.value.trigger)) + + if not events: + break + + triggered_events.extend(deepcopy(events)) + + try: + method = getattr(operator, deferred.value.method_name) + result = method(context=context, event=next(iter(events)).payload) + task = None + except TaskDeferred as exception: + task = exception + + return result, triggered_events diff --git a/tests/providers/microsoft/azure/hooks/test_msgraph.py b/tests/providers/microsoft/azure/hooks/test_msgraph.py new file mode 100644 index 0000000000000..a3c8c7ac9dba7 --- /dev/null +++ b/tests/providers/microsoft/azure/hooks/test_msgraph.py @@ -0,0 +1,77 @@ +# 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 unittest.mock import patch + +from kiota_http.httpx_request_adapter import HttpxRequestAdapter +from msgraph_core import APIVersion, NationalClouds + +from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook +from tests.providers.microsoft.conftest import get_airflow_connection, mock_connection + + +class TestKiotaRequestAdapterHook: + def test_get_conn(self): + with (patch( + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, + )): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + actual = hook.get_conn() + + assert isinstance(actual, HttpxRequestAdapter) + assert actual.base_url == "https://graph.microsoft.com/v1.0" + + def test_api_version(self): + with (patch( + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, + )): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + + assert hook.api_version == APIVersion.v1 + + def test_get_api_version_when_empty_config_dict(self): + with (patch( + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, + )): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + actual = hook.get_api_version({}) + + assert actual == APIVersion.v1 + + def test_get_api_version_when_api_version_in_config_dict(self): + with (patch( + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, + )): + hook = KiotaRequestAdapterHook(conn_id="msgraph_api") + actual = hook.get_api_version({"api_version": "beta"}) + + assert actual == APIVersion.beta + + def test_get_host_when_connection_has_scheme_and_host(self): + connection = mock_connection(schema="https", host="graph.microsoft.de") + actual = KiotaRequestAdapterHook.get_host(connection) + + assert actual == NationalClouds.Germany.value + + def test_get_host_when_connection_has_no_scheme_or_host(self): + connection = mock_connection() + actual = KiotaRequestAdapterHook.get_host(connection) + + assert actual == NationalClouds.Global.value diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py new file mode 100644 index 0000000000000..7c9ce073a1f08 --- /dev/null +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -0,0 +1,129 @@ +# 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. +import json +import locale +from base64 import b64encode +from unittest.mock import patch + +import pytest +from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator +from kiota_http.httpx_request_adapter import HttpxRequestAdapter + +from airflow.exceptions import AirflowException +from airflow.triggers.base import TriggerEvent +from tests.providers.microsoft.azure.base import Base +from tests.providers.microsoft.conftest import load_json, mock_json_response, get_airflow_connection, \ + load_file, mock_response + + +class TestMSGraphAsyncOperator(Base): + def test_run_when_expression_is_valid(self): + users = load_json("resources", "users.json") + next_users = load_json("resources", "next_users.json") + response = mock_json_response(200, users, next_users) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + + operator = MSGraphAsyncOperator( + task_id="users_delta", + conn_id="msgraph_api", + url="users", + result_processor=lambda context, result: result.get("value") + ) + + results, events = self.execute_operator(operator) + + assert len(results) == 30 + assert results == users.get("value") + next_users.get("value") + assert len(events) == 2 + assert isinstance(events[0], TriggerEvent) + assert events[0].payload["status"] == "success" + assert events[0].payload["type"] == "builtins.dict" + assert events[0].payload["response"] == json.dumps(users) + assert isinstance(events[1], TriggerEvent) + assert events[1].payload["status"] == "success" + assert events[1].payload["type"] == "builtins.dict" + assert events[1].payload["response"] == json.dumps(next_users) + + def test_run_when_expression_is_valid_and_do_xcom_push_is_false(self): + users = load_json("resources", "users.json") + users.pop("@odata.nextLink") + response = mock_json_response(200, users) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + operator = MSGraphAsyncOperator( + task_id="users_delta", + conn_id="msgraph_api", + url="users/delta", + do_xcom_push=False, + ) + + results, events = self.execute_operator(operator) + + assert isinstance(results, dict) + assert len(events) == 1 + assert isinstance(events[0], TriggerEvent) + assert events[0].payload["status"] == "success" + assert events[0].payload["type"] == "builtins.dict" + assert events[0].payload["response"] == json.dumps(users) + + def test_run_when_an_exception_occurs(self): + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter,"get_http_response_message", side_effect=AirflowException()), + ): + operator = MSGraphAsyncOperator( + task_id="users_delta", + conn_id="msgraph_api", + url="users/delta", + do_xcom_push=False, + ) + + with pytest.raises(AirflowException): + self.execute_operator(operator) + + def test_run_when_url_which_returns_bytes(self): + content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) + drive_id = "82f9d24d-6891-4790-8b6d-f1b2a1d0ca22" + response = mock_response(200, content) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + operator = MSGraphAsyncOperator( + task_id="drive_item_content", + conn_id="msgraph_api", + response_type="bytes", + url=f"/drives/{drive_id}/root/content", + ) + + results, events = self.execute_operator(operator) + + assert results == base64_encoded_content + assert len(events) == 1 + assert isinstance(events[0], TriggerEvent) + assert events[0].payload["status"] == "success" + assert events[0].payload["type"] == "builtins.bytes" + assert events[0].payload["response"] == base64_encoded_content diff --git a/tests/providers/microsoft/azure/resources/dummy.pdf b/tests/providers/microsoft/azure/resources/dummy.pdf new file mode 100644 index 0000000000000000000000000000000000000000..774c2ea70c55104973794121eae56bcad918da97 GIT binary patch literal 13264 zcmaibWmsIxvUW%|5FkJZ7A&~y%m9Oj;I6>~WPrgfxD$eVfZ*=#?hsspJHa(bATYRn zGueBev(G*EKHr+BrK+pDs^6;aH9u<6Dv3$30@ygwX}fZ|TDt1G($Rqw927PN=I8~c_R69-cY5S*jJE@5Wr0JUS6u!J~3#h`{ZMo=LkbbALoD8vfgB}Fh|2>mhOnfS$3 zNV5}8Ox=$fj;C0=UKy*{myZZPRVS|0mqr-HxZAy;()@wxQ}MN`QWAZTXb3Z&Om9W2 zbnA^OWoQbAW|3W^fw#J;YzDato8*`rHQs+@W70D&SyT{wb`SN*3nI z5G%$wJlq932=n{60Eii*9H8dFih2ks?QY=>nAFL=5g^P@#b{YUEHt0S$D7WbX zx%TzvzIK%zpvzLEd9LNr0ch#LFf_(9 zEGt0C9v~%b54vynAc{~;v&2?S(-sTTft@9CABMNFZHtY1W0-99CEbUNfp_yu{LDBz z@8z^$LPN$wX4Hi+dZQs6K3QiKKF0}Nme@EII;;F}IplC(YvT*C3-Oh#(A}e5pIz01 zyR}D2|ftBF0T=1moHZy}$wS*PSCmSzHQ%x z2tCQQCx4jt7w1cuhY69~eH`31KC4)ZZJ^)f=IabocAkBPa zEeg25yPX&9-i_N(Qiq!I3RDrfx&0t^i)&MSQ1D(w%|%#LTNr>1cPiltAYO;6kBn(B?r11c^Bz~#)z5~~V+*`U)lDFtKbZ|;? z&4wTUtK=KE&uQIWUQv1mDE;LIhXXgx44PMa@%Z<7a& zx45^oYSnei^~%}`?!O-+cgfSmn_c?`=Gmm*Z^I(96ve&$zDs|)r84)IEEiE1kfQ$q zm3km*m1)PjdU9nkk9BTlidI1~M|O~WfP7AUu2T}d>5is9l$<%;7r2&Re06w>W$KM~ zqITBTd=Ln>^crw`_N?{ z;2d_=E0n!*NisQ|XYuX9q3+UcqdA(MC45|>2tz^c6HdZOmXTB?X2Elx@_0f)1z&-gS;UxN`>Ll-kWb0X0 zTrQis=w9sJ(q7k|@|k3SA~DJ@uMXP@4(Mgn+LJC+3F~3NHW71pIzY(aHg~{O+squi zWO_|F>78)L5*gcRXXRD9IzQ(ddSxh}E7(8sC~EYrOz$9BkSMBCkGGO9FuZ{#*mW+h zvwE7d)6Ag=a*R5URs>}qdqb_E6g)kN2Wel;pWe9=hZ)XvRZR!RQg&gxAPGj8J0!gR zrdV<2@MZQ?_Ocbd5@0zI?t>$z3eD80_h^{DI)H5lk`T4lbn8kteH3%fOBH^g26#lLN2&P^s zr&d05GDs)u_8OKzCgNxllk5pLC<2wKmghL{zW%}5^}%S$?d=3OzjaSzT3>uWYikZN z2ZcR7*L|%UMs|u)wMi7#vkN?cxlBcyAM80Tyzzv&zHMF1TH9?Mx5&E57P^)^zE5N| z^foq}!--if$Uj=U6Tc>EM!Pv)e^_SZSdvtQ=@>)(ONejQ!XW8u6>ESl<*s^6cH;Q1 z#n}nL{#|{l}}@td^zNSA;R{`3A&Jjr8L9(3^2FSyZ1W9$%;!XP#N2 z-SAzyRfxtgq^py7_3*GJFO%x_v<`xJ46`~S*IukgQDKfLxzFnS&GYL!1LA{I z!c#{A90{k(b*tUfbgjOH>}{#V;%^O+LUU<*#QkLtWzjho*Kb?Cr&wC38%wxpn}^Wy zG6EpV9x3xioCWA6H6=aE3)%jmZePu#Ji7wy0CmkDZNG`a{J1i-2`Bt&UrFb&<~V$^ zy9i`R1<35M&{mtCz144%v#7LKBTPPApjoV}#W-gDc5cn;A@Mbt#zXUK@J9^vj*ME( zo8(%K{c-KDr8n1-I&Mjn)*i|pF|7l*`fXvo8-z&j{$NOfUPM-xILbX1D29IHp|__B zL*JQ8*7-VrZVY*&$!PiE%zv@osg`qx0M8+w9iy7Az7;HYezs;5NRvrdNM~t@o}5Gc zjagk3Y_>6!Ct;ITqhu3FojJO^(^SG-($M4|frkp?4y-QoSmFcw9Z%(z?eC0kGi9@? zm(vAgXU|%!6_)CrnqYL-Hj@B5hA?#8C3G^cjd?0dMSZ!wbe%O4bWvlIG=nwOEInVj zhjzd`Bry8sXBTfIUr+juZH5JyE#7~UQiwR!gmG@wm}aNyo`13xEo)tzP64MWWG|j8 z8u8a2_=C2FdRZ9(eG&Au`@$mY9vvWldP-@wj5@38H0W2V8wnaQO?!)qoS_J=(ieoI zOvH}mkBRh_p1oTW66+?3u-GH2Ex~c=BQiwpJ zJlF7O2PBaCojRRL_mp44*Iq}vcRFpBD>V9M7do5{w&b;4^<_V~Vr{+O_&hz9k5Sm` zq3|%Z(6B5~wz2k0iH-QlafAa>1%ZebdxkR;6SdA?@dK|4Jf8PIO%64Fpw$6RYG2R# zX>Iq(xf`5Xk)79-@;BAQjlWu|w@Ss3sJv3Ew&%lBu-H?vYsC8XPJD!lkv*A~z_-k= zLOaM?B5}$Sf-KF5BWHoB51WFA{GlweQna618{*tqVn)YKUVq?khU_=QER9uW?N17xgAponbjg0W`=>f;sulH3?st)Y_@k$We2-__a>^{E78lUiI13qq!3# zwxMEl75MK1q`~J>ST#?`mUx#vr%-jwpZ+DV;W!0KNkZmO#sK)zt)H@`EQl6RRWhwb z0&E7|fG~@z)wlK1-RsxN#8Gr)D5=xpv=b}=CWPbwz@(9bIhD0Crd-Q>qEo>~Gh{X7 z77AK5>TfF0wK!?7Nx!<5uDy?D{Qg$SEc_R3J9EuH!Z@qmEJ*QRRHd3BPirM6783nv zAnab$>rhdDJ6pO@%Ox(}BYw{Ba<3|=A%Fg5_Hfxj{%CfzZCFO{?%h&=?%CNBvi&p; z(otqN>+5giLLa^*G?xzN30=IgQrV+r7dW4bX;zKtuD)O$UnwAKC?CpkPt{77nUArH ze-jKcCfRrOlp(Q^b&W}mrgt4n%wikNxeSBBE_n>K-IOIzi6!<)xGRYA)wGgqp^s@d46N#krDHPc#9SOgXhI7Vbj?B z%c6@8dCOGPYBoNE#3N7HD^ihbC9*xGm6chu;?fcuv)s01keHHZ1vXl5D;29O7wZBr zyPzyLZHKMtUI%PK+*X2zTFtaDzU1qn(H=hRRj-SoJw7I5i%4b0u=&InEAKgoae-lp zXk0SkjlJ52HruS*1QykTZ&aCN`PbcKuw$1st{peJ@&aF^aR@~{XA@L&YvK%+VU}G4 ze5iuesu&i6=*#nvHbm_v-ZLr5^Ij#|YSAper4XpsH;0x(2h1-tIobIy;0~2a( z!G($SB!iu#P;;hGeI~C`O=-3|d~zoB0!`*JrU-)Ko_X5#kSpy5o^z49RG;{j#l~45 zF?X9Ih4IdviT(8@+q|`BveLTprbESZ6^2I&ew|V3pDXRe9gSyXT)zzqKQ;gCD;p+( zM)2(;YJ%P5)X(N3ZSn>dn6UIcEcvQOXZBn}uD!7V0yXr$f+d@eTSYoquPit2S8cPW zA8t3dX)Cv{0cKF`@e|PP(xS0|z2_R0(P6)#+kC$0^5- z$7Hs|bOQanE z1oJ;uh(dYiDt}mVmtC3&HaGT6-dY429v#ySHJ7V)C8ow=PSmnEI)=b3_RJsU(S*+J zV$p3>RkK?DFvTc;(-T=h!1u~CP!pE=0eSSu#c@N7S0Z57CPg}!5z{QL#`2v?DJDt^ zCGN{0p-&&=)Sb28Xlo;ZXc^CGdwL9prf30uu$y5aPeWD6WIk4%%~DEhTiwOvy!rS% z&3z#DWo2qBA*=M2xIu=_R0sbrmP;Y?_rRa^k}3WYU6n9H^(})Zi-woMKKXfgbab@J zWx3DUr0MLpdDYk_LO8As}d*Z=x^K+uIv#T&SnY6&C$9 zBn1u`G#TBt+n5b%a;Cr0h^sm5Fl^OdxJ^8IebW);DWATq#Ba=#rggj*wNKy5NMzz& zBm`bk9bcSVPJbC`dHrI>o^=LSvTFpT`VAK`x_naOpvS~*l2$1vIk$avBA!|aeZ+7c z$_9Zzh>fc4$uX&w@-$VORCscG(B)OA@SPj>BNY3gxkkcPgNi9bE=?&3A4`3ekrdsb zn~`M;p8I>4?@@ZI{9Afv(tC@pp@Oe5BYUw-%&J_WaTBGls)&d8q?t$i<<@=_CNfH! z4H!ww7#gkp_^`bxZaJI9@C+A9x7@E1ZRoG5PL?w3GDi>`8Qq%I+0ygfT78%{Zt#mP zqX0CzaHKn@hAOQsv=^8UbfpuyFnT8Ht++Vmmx$~09!e{5t8fMkEjr~tfIxMlIpr4zGwvEIWKC2`Q#C)c7QF9wet?hE zLKoU?t@nqm=iBc` z8_((*(i(g}7z)3{%SJ!uya{?Ir-2^Fiap*VC4pF@N zpL5F*DG+(taLhdu4DbyAP(0&60n@%?G~hHugBI^-X6@_YOu}8UqwbQ8V`2vwDRLMz z)aRFo+r1f?5idT9xRF`cjgx$a-IpH3AH|bs$emw}d23*3aU0hYNh4(D0o-Z+wIX{d zeann?lzjgsAt62`er@<$`G755?i7tl%CHNgXp}#j>j&S1n5wZ;ofNbI>B2*4L1}@3 zq(LzPqn()w{KBsX!5*a&=dv<}t=R%II;TcQatbnKM7S4Q1PQIoT=^$#=>Y(m{mBYtl5W z6}|l4kxikOcJ`C3o{TSxIi?8|N6sH7Lkhq5qttl@uBTA|-cBluU$hU0&xYKvNidrL z4q>|j76}G1Db23Fa|XlFm%W&jW0h#7B$_FD-ZhqJ5#7i!0ZmCrereX z|Jlf`<1zR2akFe|boWv-r=}kM03o|%$mZA7Of2T99u~e56~6sh$P=yk9f!H6msn)n zvFOLF?W?iqi6fK9C)a42Sgt0kz4#M6 z-UY6451Er~=V;ITs1O-q*>}{;bs74MMZ(Z&=Z{5#q+i@cw^vI#0|Dh~-Dh-tn2I(S zTXXp-bLEG{p0#BbIqIcTM|DWZmr`&br8u)jQ`CR*^+g_fIX%=K+)x}F%Oak-Uh$6nIHUavnNV5M7YffU80QPRD%y>T{bIzn<6Rsy zb6cW6`?0EwSn;uJddPn@`?^Cry2s(6ccP1ykKr!kmDg2~zbTJq@+e(z5N>ZNr|8$j zPi-~ofp7E|Xx1#H+f@UR@AS}iLP!}}dRwf{u!avAq-_hNw#uaoOD{2jo*eRn8$~bDK`h1&ssOC6ekGV38+hU!KR z+kpnSzT;y#o|V2h|F?SY4-z1MFxz0;)@Lk`H>Cj zSl@fR%*@F79;HJcsX%L8_d!%TwmQyi$|n&C{oBMJ9~Xm!@@#lZdz(WB9SgJ#NIC%@ zy+~ZnI|4E`7f@W0Y9I@N7UTs1fTPD-ZiU%Lr2MnP+2h8AGh?(WGVf>h@W-_M>jRkD z(KNxvo(UJ7)o+*t%fCcM10;2XM$1NAFKwhp(c917^io_ynn-yv58IFIF*UJUw*2Ma zm?a-a1yp9B?WxpLzap-c^$HKkX_IfT_W8Lqaltl*A%vZSZWAe`Kv}vjz}>Tc;Hw9T zA+Nc49X&{WDmxY~ReV0YceXdL!$9mTL$Q@_vXIW6I{G=`$KR7jFcE&IsHwnKX;KldV#YL z(xwKAB5cFiz+r6m*5iJvo&E)XQqVWjmA}BfyVS&dm9&Y%$Sp^sW!JE3iI0v(kQHdo zmhWk|gC!e@CFKPv4BE*U;mYo0y}J0J-Fhu!c%v+paQf9+3Ed2EkfPt(D7|Ok#t)^PGr3Y)RGfvO=k;@Xry=Cf3fLCQ# zi`%oCt+vyB-t{iEgI&+2dczmnMXj>EOmSpMuuL8Ob`1$D;fc$wM6j2HH4Q$ zqaoj&M$2sLhpptdJMbs!krJId=iOd}HdP4Lt@yf42OZ{pOoQ4_gShz_sMoWYX}yQd zDQ8(tc7UvTt%`0#?9K!C^J>GpucEnBhnsWg102Z=uzOlwez^q^j7nV$krID#wC}A$ zcRfc2)T5Y~({6@1`{yL-Lzs;miT@C9|1SIFBMK7cz*E;v2H|EStZphjfb5mGMpw{q z!pl;Vw772tuvDH4o$;j4u8)@=m+&BIf4Ix(u75P?Q{4Y8^uvpq)mCW(enuQc)hx$B zOY{`_*%~bm%k*x6y;)D8_-yYbMsC8y#1H}89X;M=a#*HT>d*NFf}x$pQ&X?nFtvzA zKH|l8y;frsm|&}<%&*}Yu}Yn0M=Jy8qe%<1qXRR%Nut}Aqr+1pQS*D7Cp`+8Y`RO02p14DyVOmSYlEzZ;9&JzYhtybMZ%e4s zlks=V(+aJ!LK-()3ox`%9c)lx#3#y4{ulL6KpG|&>9`n?Uh#m3G-mZy-3h98Scyja zH^3Pb7?P z+2hAkyvg}g$#)n$Gs2fL19JNOZ|~>Nx(|}lmwesC!>?Y~72mpf4XZ8t^TIwbCk;i0 z+a2ymSZ^=OrtrSH!(y#Vn!8KWk#O7<1-!if+`dDDy18U7wS3k$lIeM}Z0fhYqI)+x zo*o4*S$S|hGf6vL>PaQ(OQ_%eskx-G-FV|dXHbTH<#w@RbeIx9I$d$xqHh`{*&d3y zevlYNk)}w@cuu4A$^DYJsOvO7VBaom@Rx@gb$V5IKJ{Xue16H-1H0j=U0brW-aVRG znWCQRkESBmD^4?a7mB@!jf2>(Hs=Bd-;XX1oEilevb9axB^NhIPLO>jl03S+Rw|fx z&oIsIk(~W!4$zzKF|uSR<@S#;{r;fKup)iDaxz_9JouroY>XHcrN(Mm@UHV?-8bCh zXGfY~7U`rCasv(h-R*ava)^ zF1`BMT*n3xQBTdM?`n&h2Ecf*XXuLo7Zyl_El(v~oh>}mK01$%0a@#uzyiX_g>Bav2XWwH%YekAxU%pBT!p*?%cS#zA zv;^eDC#KZP@7o=^GDc_V8<3w>`*L(+=A#(fcH)dGjqM}Vk_el+c>B`{9xm<>IZ-Zm zLL!-Yf*3nju_(8ZGUd9*K`iofWW+BYFnZF&+a|=yxqV?oUOcG#ulnSR$DMs|e5Tph%WW zVjzE3nMh7+rG!}av)+~;o$#+EHyPX zzOUO?^#)Jh*t^b7pTW+I%f;xy&JMPCO&5RR``BmHX-Mw{qoJp9BjKea$;A9%>-iEZ zvuUBm%0j5UWax~`ue!K6dDdip+zs3f{+qQKqH;9C(1Z@95()-Ew=`BdLh2VS3zI8qYGH&&7m9+vpUc+x8l!i-ATXKhw34XL2;ya_VIQz!OL^)8mtqnb?q=~&^h-$;Zn^HRZ2p(gH z39An;`AWT=i&VP0u&CUe7OYW51Icv=q%Vc7%Zm z_uAp9n}osEUdk2*pV)*i`WRSa-FWtCwGqS-75@K#V0)r;+0(0XVp9vnb7lWiMj!q= z>Zf(ioa@gSwA55Jil$lh)%4U<)$j@HTQU2KwuUUsZA*2O^QTKobak8g0Qb~ROMTW7 zfTF2yF*na6i(lQ*Nq^rPen^0>$$b`K!Kp{FVa-VF`kCiXZg0Vtr}i*rcpny_YOR!} z+?Jiv?dWlT`}o$s9Fxt%%684d7ek-q-Q~jS*I5+8HtvSw+Rp!D=+gVr!gqcYy9K74 z&eClx6f6{1Din;ynjz?XZlJ~W7^A@0wiHIt8$aou;f>MYpU%gUlDwAK*nX0#vHtyl z_C=B+ZkOffY|oR^2>(+IlZCTMFirZMhn>bqzR=38hvJpcM4-@gUYY7_k^G*FW9;5r zc9q4c>C?hd{uS3{MThN*(w!3e05e?bI#SNlo$U&%>((Dz0_JeqbG|}!wI$& z%q2JQ)Vas;i0RYqNXW!CC~QK%u$K$beGI zT2KuzMjus26(zmofK;m2gY%d*o~sHBKA#`RBNc9c*-GLmbgh?*9V;^TBSot2E%~Q5 zl+R!WA_h_JT;+irbJ#Z-tSy-;B^t&&dOSwPV(T!CB)no8Y4sP%k(MD^0P!NL1vK&7 z`3luW2$gkI#Zf>IZT2=m4R&e@d zeo#B=Q|9`w8}%|)f%GBjYO01&Dk5qjm$+#1yia#CE=Sh~88Vdp%|VU}0a6mF@JkhUY&~W3f#rHK-1Qdo z>0*z5?#-hQUY}k^X7~1bkI?($-~3#c3mF4Cl@2%|0@1=ARZ z^qlNaN63&>;O_~mmto}?tAhznb}p;GpyIq1Z^yf<_6Ui~cpbbP;uV7W!+ke>wYG-f zPPz2~%UgSs(>vsKFle%uo=WIDYz;BR!doAy)aQ0QCpE_Wz1XK+3Kpr=V_H8w zqzaizn9ALx#?fo-N)_CtENYH*1|ID|x=xa9d#;9~1Wgrcx^8=evrfky*Xj`269~A;kh^O|ewZnM}=SmM7NX=?h#jjLh&1kIT+A z)If4luYo@s+e_L&eRJ$gw1`)>u#efOq=M0iYIPS$GII0z`T56eNxK@~Y%*^~Q&w$1b)jM9Z~kuRc~YX`6r#ySCskW5cq|#a39s;ZiaL~OdEpgu z1k*sKkLZ&?6fAi=)77yKI1xii%)@DG8r}663xkJcwLTj?s`h{GP@_2}`A|;w7zrzk4QOQ*O$(e|M^<`vLD*1^i>Nr*= z+A`y@f{!zLi)ys9OrFM5`Qw0292Ciyq>zC>8(TkG1O;#UUh?#I08kuwpS_vhufJ0v&p^Yr`=^WG7!qVG(8n9u7=J64fr zQq7B|9rzl7s)I_|8UeVp?=cqGILQ}0O(n+^vJz=vFBU9JmG$=DWzi+qCHw@D0a7`M zA`%pmU8+8W{u0{2*^tg&3;I&i`4`{YJe_n8 z{viTJZL?$}#l9w${3mydrW>Z%nY!WXf$HJv5$Zw4F%7^mXWsZ-s&olv31;C*KlH)j z?j?Eika^cI`l>)WJ*ga?%>0HwJm{%<)OP8pdvwMG@fm;Ca`jfy7ixY-sic42*f&ld zJg3(O0~;=Zsp@cdUj@&Zj~#~LX=F5Ws@!Ik0-~(wlbJO6&)S~s6WrAW9lrQ%6+S03 z&P&xJ{;BC%2s%J#uxZy3=Fc}fkwE9(T}QAK9b{FT!L3^PQ~;#X$T|9v&JFq)ru$h|ls zvPxYyWT}V&Dol3#)t6pVE4nIClEq=r++eGcG-tkOW4{n$Ra~3z?`@_gXRUiR`SrhY4K z#>C+t>pNtm>!Zw*;p^qI0|g<)Ob`r0jaN6asw2ZGLT}bMbHnQ$OH8cR7{Rq?=4%&x z2Qe&O`w$~b%fuo>fkgT`PVx=uto@&SdDpIXL)<da|A*x(b?o zdUj^iN+B9%;2{1URo7=%m@r*RJi3fQNO_`AZY;b#tClm;A}NQF#!Y;pMMdh=^fO@9 z>J>Xv^joKJM>M7x=xh!oSLO3JlxVwTn$DPHdGsnkAvB)9d)IE6ZHgd1vd+Z;W1d682CBy4zti z&6;T6!rzSKIy&zKKfAx9J%7q-=Mac{u-_GIYEaZt*`h25Ne?ch`E_c2{pGA<;nVkx z102u6#||N$g5MhA{!rFwaI(;8$S{1DePGc^L~j6?Q$2QMIO09 zPdma#_kX(|;oOau(pX877ac9V4O8x3g{Mdbr6oS)7 zN0v#H_j!bhUNl;q>GrkeA~){;lCg@&Mg5(z%E1HV`d7{>_}@9JZ(VJn>=HKC4q{My zLpw8D2OD@&E}T?=SV7rE-XI?4H+E(aOI8sZOC$NW=!leE6MG6ycn2;fB4XpB!^#Z= zQ?P=-+!R0#4h{+c2LPbUF6{uZG&6i-ZDI+f;6P`8V{ZtxcA((p;6i6ds6r4x005m` z6k;m{H8U}FK+J;+syaZe)G2u2J;eI(G+`)^0+C~@0#BIzJLi_?-}e8NR15?I|34|k zx>2LneiYApj|7nW4k1sp9h-vz^G);Jq7ONB*clw!(IJ2QT3sYWS)>yb_Ual2Um3r5 zw706UJD48HLY73$&Gm=sl|EYND&Uk>VT!eN_p49f6HS<{TU>u{4&#WYh1dwy^E8il ziH`_=$2m8k)y$Q2yDZQluP+AZbND!Yi7Co@fwHnw2pV1bo*=wGx2n7Urt$y1@imz1&#&nK47Nw zT-dLY@^1NHY?5B#-Qf9?`lA_={@NnLpmwJGQG7&oU}0>) ziZ`GdjY(jIKi2Q?e+d=de}nq3pkP;ZG;lyf$Xh!{=x?qF#2$)p%>NM^W_I=tqNWf# zgv;e1fAtY=)-W@2FtyhKb8%3Bfj|mw00#vR4=)857d&XdU z(4fLD4>dA_AWjHkeJ)-u3LZ|NF1w_ijiW6*A6^xXD#Y5}7O{k(E4!#F{9rhl8A4Sg zMcAb&9N>rx39*a9v4(4~r$8jq|MLt0{*hTPYU2nu0sub&aQG~$!9>qU@%LGVw1{ZAdD5crj3WAdl2KV62-uIT7sX=aUZ*>8aV1F3(c z_P=p-FtxG!8!9*^U<3>RcoByeFaipAK|lhB5)AqaI)n^@hmeEwxOw0OKK@%C0pZ{C z5o^F{FbEE(DEt!$_$B<8DlYiaV7ME855ql#Py+_S#o(c8`L;d6lqRR~$cn(zq-4};(pf)4`xt=`PWS`7YO27?$MdgtpDP{`vCa4 z{2x3Z5bm@8-~oUj5Zv+q!Gl}N`CoDX0N4M*gTIpgb1nb?;)Y)s|FIqb0Ot6gw!m#h zTnhg~j+YZ2)c?r?0yzIm4hZ1=FTFrc;D6}=a`OJeW(PY6{AFi{I1;L6ZcsR+>?$@k z@FNVDLEL!K*2XpzfZwk|I3Y%%Lm?mm76XGtKw?0k2(JV$kO#;s#>p!o!6gRf5#f;l j@(7{-|3%=32kuUL2Z)`+Z(jm{U>-0!Ev>ks1p5C2Hj`#V literal 0 HcmV?d00001 diff --git a/tests/providers/microsoft/azure/resources/next_users.json b/tests/providers/microsoft/azure/resources/next_users.json new file mode 100644 index 0000000000000..f9376541b6f15 --- /dev/null +++ b/tests/providers/microsoft/azure/resources/next_users.json @@ -0,0 +1 @@ +{"@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Leonardo DiCaprio", "mailNickname": "LeoD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Meryl Streep", "mailNickname": "MerylS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Tom Hanks", "mailNickname": "TomH"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Jennifer Lawrence", "mailNickname": "JenniferL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Denzel Washington", "mailNickname": "DenzelW"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Angelina Jolie", "mailNickname": "AngelinaJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Johnny Depp", "mailNickname": "JohnnyD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Scarlett Johansson", "mailNickname": "ScarlettJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Daniel Craig", "mailNickname": "DanielC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Charlize Theron", "mailNickname": "CharlizeT"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "George Clooney", "mailNickname": "GeorgeC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/149136c3-62f0-4d27-94f1-8a27bfc4cf73/Microsoft.DirectoryServices.User", "displayName": "Julia Roberts", "mailNickname": "JuliaR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/08f9d41c-a653-4de3-85c2-324eb53bcbff/Microsoft.DirectoryServices.User", "displayName": "Ryan Reynolds", "mailNickname": "RyanR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/cf5ce2ea-132b-4699-bd8f-6bfcac619fc3/Microsoft.DirectoryServices.User", "displayName": "Nicole Kidman", "mailNickname": "NicoleK"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e255041c-28e5-4abc-bb03-09205096cd87/Microsoft.DirectoryServices.User", "displayName": "Sean Connery", "mailNickname": "SeanC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/52f8faad-6b39-4cbb-8614-e968f5af9e9e/Microsoft.DirectoryServices.User", "displayName": "Emma Stone", "mailNickname": "EmmaS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e9bf2db3-4149-4396-aafd-278a0212179a/Microsoft.DirectoryServices.User", "displayName": "Robert Downey Jr.", "mailNickname": "RobertDJr"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/f66aba7d-5284-418d-a025-6d6f0639350b/Microsoft.DirectoryServices.User", "displayName": "Cate Blanchett", "mailNickname": "CateB"}]} \ No newline at end of file diff --git a/tests/providers/microsoft/azure/resources/users.json b/tests/providers/microsoft/azure/resources/users.json new file mode 100644 index 0000000000000..abaceebcf9a54 --- /dev/null +++ b/tests/providers/microsoft/azure/resources/users.json @@ -0,0 +1 @@ +{"@odata.nextLink": "https://graph.microsoft.com/v1.0/users/delta()?$skiptoken=qLMhmnoTon81CQ1VVWyx9MNESqEIMKNpEWZWAfnn5F7tBNFuSgWh_pXZOweu67nEThGR0yQewi_a3Ixe75S6PoB8pdllphCEev0fMe5Uc1lWMtn3byOS8_OPTzPGZIZ17x-dVyxaE_4I55YyLJ0cgBxg8wsBrkYgaNE9vy5Su2HeCKxJODDQk4zRgP8QGo0pZatReTpqisVbrW5Gl1H_Xgy4lhenv1SmoRcBQtWBa5iAh-MURoaTo7i0kQjFhH6SCrkjBkfkRFVy9dafOOt2Owbxfn5hKGfEnfmG0RBmgdUsZPgX-ap0mjjf7PjExoxMek4CDnb8Yv737oGkh9C_G0XTJGeGxPBbkD-w4SaQookde4yxOzceAw1MuamBy63uJdbXt1ul61tDvfPwrJVHq99FxGU1n-i_RfHh65nRCHju3N3ApOFKrAi6933l3VupyaXsvmm5pCPh0T70dYK01CYktBce8Mc1HaVqB7SR-R9-X4PHYfozPWeqv4hng4YEvqA41XjRPUzOaS1VTH08k-HhR9ENpqw5UTFnimAu5RbMT5fTbUAMTQC2XcWF_5aDgjuw8D2VQvr5VsB7qmu4mgGb_dNrHM47QyJCKY2QcgLmmsTnr5Z-7Qe2AGGy1b5DREBVoLnSL-aJ_6m9TAlYD9oGityZbDJ1ssVxS0XsGYxAwSa5z_E_lgedr_ikHc0zzDAdj3TgLGD5gWIjwvxnkNEXa0onk-jqGfANEDFh0vGuDz1mlgwiHGIKN-QjmfTX-Equ_uY2lbuPcIXTVpdrgQob5BaXONZ78uVh-AIlm96PPqbwNj1QMf3EE0DfgGaMIloFdByjSD5FjwQ_6COxueUJw2iUtKh8l4fpja7yOs4QCP5_tPvUUZT26ylLHjbPZO7I1TOaJb9OTBM-kplBJitW3SAm7DJM9jocA89Iqbm_ap6mlvEET2cyw6Dn7j2AP_0NsdoP1MJn6H0JE6HOJddbUVuzhba58QBujgHsEHo.WPJMTkIZstDbZSnvkZm1eX2ASvooRj4BgY-GW_pQEZo", "@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Sylvester Stallone", "mailNickname": "SylvesterS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Jason Statham", "mailNickname": "JasonS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Jet Li", "mailNickname": "JetL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Dolph Lundgren", "mailNickname": "DolphL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Randy Couture", "mailNickname": "RandyC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Terry Crews", "mailNickname": "TerryC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Arnold Schwarzenegger", "mailNickname": "ArnoldS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Wesley Snipes", "mailNickname": "WesleyS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Mel Gibson", "mailNickname": "MelG"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Harrison Ford", "mailNickname": "HarrisonF"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "Antonio Banderas", "mailNickname": "AntonioB"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/34c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Chuck Norris", "mailNickname": "ChuckN"}]} \ No newline at end of file diff --git a/tests/providers/microsoft/azure/serializer/__init__.py b/tests/providers/microsoft/azure/serializer/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tests/providers/microsoft/azure/serializer/test_serializer.py b/tests/providers/microsoft/azure/serializer/test_serializer.py new file mode 100644 index 0000000000000..3318666cc968a --- /dev/null +++ b/tests/providers/microsoft/azure/serializer/test_serializer.py @@ -0,0 +1,46 @@ +import locale +from base64 import b64encode, b64decode +from datetime import datetime +from uuid import uuid4 + +import pendulum +from airflow.providers.microsoft.msgraph.serialization.serializer import ResponseSerializer + +from tests.unit.conftest import load_json, load_file + + +class TestResponseSerializer: + def test_serialize_when_bytes_then_base64_encoded(self): + response = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + content = b64encode(response).decode(locale.getpreferredencoding()) + + actual = ResponseSerializer().serialize(response) + + assert isinstance(actual, str) + assert actual == content + + def test_serialize_when_dict_with_uuid_datatime_and_pendulum_then_json(self): + id = uuid4() + response = {"id": id, "creationDate": datetime(2024, 2, 5), "modificationTime": pendulum.datetime(2024, 2, 5)} + + actual = ResponseSerializer().serialize(response) + + assert isinstance(actual, str) + assert actual == f'{{"id": "{id}", "creationDate": "2024-02-05T00:00:00", "modificationTime": "2024-02-05T00:00:00+00:00"}}' + + def test_deserialize_when_json(self): + response = load_file("resources", "users.json") + + actual = ResponseSerializer().deserialize(response) + + assert isinstance(actual, dict) + assert actual == load_json("resources", "users.json") + + def test_deserialize_when_base64_encoded_string(self): + content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + response = b64encode(content).decode(locale.getpreferredencoding()) + + actual = ResponseSerializer().deserialize(response) + + assert actual == response + assert b64decode(actual) == content diff --git a/tests/providers/microsoft/azure/triggers/test_trigger.py b/tests/providers/microsoft/azure/triggers/test_trigger.py new file mode 100644 index 0000000000000..59715f907442f --- /dev/null +++ b/tests/providers/microsoft/azure/triggers/test_trigger.py @@ -0,0 +1,123 @@ +# 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. +import json +import locale +from base64 import b64encode +from unittest.mock import patch + +from kiota_http.httpx_request_adapter import HttpxRequestAdapter + +from airflow import AirflowException +from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger +from airflow.triggers.base import TriggerEvent +from tests.providers.microsoft.azure.base import Base +from tests.providers.microsoft.conftest import load_json, mock_json_response, get_airflow_connection, \ + load_file, mock_response + + +class TestMSGraphTrigger(Base): + def test_run_when_valid_response(self): + users = load_json("resources", "users.json") + response = mock_json_response(200, users) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") + actual = self._loop.run_until_complete(self.run_tigger(trigger)) + + assert len(actual) == 1 + assert isinstance(actual[0], TriggerEvent) + assert actual[0].payload["status"] == "success" + assert actual[0].payload["type"] == "builtins.dict" + assert actual[0].payload["response"] == json.dumps(users) + + def test_run_when_response_is_none(self): + response = mock_json_response(200) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") + actual = self._loop.run_until_complete(self.run_tigger(trigger)) + + assert len(actual) == 1 + assert isinstance(actual[0], TriggerEvent) + assert actual[0].payload["status"] == "success" + assert actual[0].payload["type"] is None + assert actual[0].payload["response"] is None + + def test_run_when_response_cannot_be_converted_to_json(self): + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", side_effect=AirflowException()), + ): + trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") + actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) + + assert isinstance(actual, TriggerEvent) + assert actual.payload["status"] == "failure" + assert actual.payload["message"] == "" + + def test_run_when_url_with_response_type_bytes(self): + content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) + response = mock_response(200, content) + + with ( + patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), + patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), + ): + url = "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content" + trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") + actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) + + assert isinstance(actual, TriggerEvent) + assert actual.payload["status"] == "success" + assert actual.payload["type"] == "builtins.bytes" + assert isinstance(actual.payload["response"], str) + assert actual.payload["response"] == base64_encoded_content + + def test_serialize(self): + with (patch( + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, + )): + url = "https://graph.microsoft.com/v1.0/me/drive/items" + trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") + + actual = trigger.serialize() + + assert isinstance(actual, tuple) + assert actual[0] == "airflow.providers.microsoft.azure.triggers.msgraph.MSGraphTrigger" + assert actual[1] == { + "url": "https://graph.microsoft.com/v1.0/me/drive/items", + "path_parameters": None, + "url_template": None, + "method": 'GET', + "query_parameters": None, + "headers": None, + "content": None, + "response_type": 'bytes', + "conn_id": 'msgraph_api', + "timeout": None, + "proxies": None, + "api_version": "v1.0", + "serializer": "airflow.providers.microsoft.azure.serialization.serializer.ResponseSerializer" + } diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index bcf5aa65fe6eb..ab410e579f23a 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -17,11 +17,16 @@ from __future__ import annotations +import json import random import string -from typing import TypeVar +from os.path import join, dirname +from typing import TypeVar, Iterable, Any, Optional, Dict +from unittest.mock import MagicMock import pytest +from httpx import Response +from msgraph_core import APIVersion from airflow.models import Connection @@ -68,3 +73,59 @@ def wrapper(*conns: T): def mocked_connection(request, create_mock_connection): """Helper indirect fixture for create test connection.""" return create_mock_connection(request.param) + + +def mock_connection(schema: Optional[str] = None, host: Optional[str] = None) -> Connection: + connection = MagicMock(spec=Connection) + connection.schema = schema + connection.host = host + return connection + + +def mock_json_response(status_code, *contents) -> Response: + response = MagicMock(spec=Response) + response.status_code = status_code + if contents: + contents = list(contents) + response.json.side_effect = lambda: contents.pop(0) + else: + response.json.return_value = None + return response + + +def mock_response(status_code, content: Any = None) -> Response: + response = MagicMock(spec=Response) + response.status_code = status_code + response.content = content + return response + + +def load_json(*locations: Iterable[str]): + with open(join(dirname(__file__), "azure", join(*locations)), encoding="utf-8") as file: + return json.load(file) + + +def load_file(*locations: Iterable[str], mode="r", encoding="utf-8"): + with open(join(dirname(__file__), "azure", join(*locations)), mode=mode, encoding=encoding) as file: + return file.read() + + +def get_airflow_connection( + conn_id: str, + login: str = "client_id", + password: str = "client_secret", + tenant_id: str = "tenant-id", + proxies: Optional[Dict] = None, + api_version: APIVersion = APIVersion.v1): + from airflow.models import Connection + + return Connection( + schema="https", + conn_id=conn_id, + conn_type="http", + host="graph.microsoft.com", + port="80", + login=login, + password=password, + extra={"tenant_id": tenant_id, "api_version": api_version.value, "proxies": proxies or {}}, + ) From 86643d40135acd587b31c72ce38e1e0c8a9a6373 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 13 Mar 2024 17:19:23 +0100 Subject: [PATCH 002/105] refactor: Extracted common method into Base class for patching airflow connection and request adapter + make multiple patches into one context manager Python 3.8 compatible --- tests/providers/microsoft/azure/base.py | 14 ++++++++++ .../microsoft/azure/operators/test_msgraph.py | 28 ++++--------------- .../microsoft/azure/triggers/test_trigger.py | 26 ++++------------- 3 files changed, 26 insertions(+), 42 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 8486eae988c5b..ca96eb2ebea60 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -15,11 +15,14 @@ # specific language governing permissions and limitations # under the License. import asyncio +from contextlib import contextmanager from copy import deepcopy from datetime import datetime from typing import List, Tuple, Any, Iterable, Union, Optional +from unittest.mock import patch import pytest +from kiota_http.httpx_request_adapter import HttpxRequestAdapter from sqlalchemy.orm import Session from airflow.exceptions import TaskDeferred @@ -29,6 +32,7 @@ from airflow.utils.session import NEW_SESSION from airflow.utils.state import TaskInstanceState from airflow.utils.xcom import XCOM_RETURN_KEY +from tests.providers.microsoft.conftest import get_airflow_connection class MockedTaskInstance(TaskInstance): @@ -66,6 +70,16 @@ def teardown_method(self, method): KiotaRequestAdapterHook.cached_request_adapters.clear() MockedTaskInstance.values.clear() + @contextmanager + def patch_hook_and_request_adapter(self, response): + with patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), \ + patch.object(HttpxRequestAdapter, "get_http_response_message") as mock_get_http_response: + if isinstance(response, Exception): + mock_get_http_response.side_effect = response + else: + mock_get_http_response.return_value = response + yield + @staticmethod async def run_tigger(trigger: BaseTrigger) -> List[TriggerEvent]: events = [] diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 7c9ce073a1f08..1bd5602a46a09 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -17,17 +17,14 @@ import json import locale from base64 import b64encode -from unittest.mock import patch import pytest -from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator -from kiota_http.httpx_request_adapter import HttpxRequestAdapter from airflow.exceptions import AirflowException +from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base -from tests.providers.microsoft.conftest import load_json, mock_json_response, get_airflow_connection, \ - load_file, mock_response +from tests.providers.microsoft.conftest import load_json, mock_json_response, load_file, mock_response class TestMSGraphAsyncOperator(Base): @@ -36,11 +33,7 @@ def test_run_when_expression_is_valid(self): next_users = load_json("resources", "next_users.json") response = mock_json_response(200, users, next_users) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): - + with self.patch_hook_and_request_adapter(response): operator = MSGraphAsyncOperator( task_id="users_delta", conn_id="msgraph_api", @@ -67,10 +60,7 @@ def test_run_when_expression_is_valid_and_do_xcom_push_is_false(self): users.pop("@odata.nextLink") response = mock_json_response(200, users) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): + with self.patch_hook_and_request_adapter(response): operator = MSGraphAsyncOperator( task_id="users_delta", conn_id="msgraph_api", @@ -88,10 +78,7 @@ def test_run_when_expression_is_valid_and_do_xcom_push_is_false(self): assert events[0].payload["response"] == json.dumps(users) def test_run_when_an_exception_occurs(self): - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter,"get_http_response_message", side_effect=AirflowException()), - ): + with self.patch_hook_and_request_adapter(AirflowException()): operator = MSGraphAsyncOperator( task_id="users_delta", conn_id="msgraph_api", @@ -108,10 +95,7 @@ def test_run_when_url_which_returns_bytes(self): drive_id = "82f9d24d-6891-4790-8b6d-f1b2a1d0ca22" response = mock_response(200, content) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): + with self.patch_hook_and_request_adapter(response): operator = MSGraphAsyncOperator( task_id="drive_item_content", conn_id="msgraph_api", diff --git a/tests/providers/microsoft/azure/triggers/test_trigger.py b/tests/providers/microsoft/azure/triggers/test_trigger.py index 59715f907442f..ad6d12d267f06 100644 --- a/tests/providers/microsoft/azure/triggers/test_trigger.py +++ b/tests/providers/microsoft/azure/triggers/test_trigger.py @@ -19,8 +19,6 @@ from base64 import b64encode from unittest.mock import patch -from kiota_http.httpx_request_adapter import HttpxRequestAdapter - from airflow import AirflowException from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger from airflow.triggers.base import TriggerEvent @@ -34,10 +32,7 @@ def test_run_when_valid_response(self): users = load_json("resources", "users.json") response = mock_json_response(200, users) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): + with self.patch_hook_and_request_adapter(response): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") actual = self._loop.run_until_complete(self.run_tigger(trigger)) @@ -50,10 +45,7 @@ def test_run_when_valid_response(self): def test_run_when_response_is_none(self): response = mock_json_response(200) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): + with self.patch_hook_and_request_adapter(response): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") actual = self._loop.run_until_complete(self.run_tigger(trigger)) @@ -64,10 +56,7 @@ def test_run_when_response_is_none(self): assert actual[0].payload["response"] is None def test_run_when_response_cannot_be_converted_to_json(self): - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", side_effect=AirflowException()), - ): + with self.patch_hook_and_request_adapter(AirflowException()): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) @@ -80,10 +69,7 @@ def test_run_when_url_with_response_type_bytes(self): base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) response = mock_response(200, content) - with ( - patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), - patch.object(HttpxRequestAdapter, "get_http_response_message", return_value=response), - ): + with self.patch_hook_and_request_adapter(response): url = "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content" trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) @@ -95,10 +81,10 @@ def test_run_when_url_with_response_type_bytes(self): assert actual.payload["response"] == base64_encoded_content def test_serialize(self): - with (patch( + with patch( "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection, - )): + ): url = "https://graph.microsoft.com/v1.0/me/drive/items" trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") From 6b9219502360a9a2d1d9111b424db71ac3383fc7 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 13 Mar 2024 19:15:56 +0100 Subject: [PATCH 003/105] refactor: Refactored some typing issues related to msgraph --- .../microsoft/azure/hooks/msgraph.py | 17 ++++---- .../microsoft/azure/operators/msgraph.py | 39 +++++++++---------- .../azure/serialization/response_handler.py | 8 ++-- .../azure/serialization/serializer.py | 8 ++-- .../microsoft/azure/triggers/msgraph.py | 35 ++++++++--------- tests/providers/microsoft/azure/base.py | 10 ++--- tests/providers/microsoft/conftest.py | 6 +-- 7 files changed, 60 insertions(+), 63 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index 67d89a05606fd..a70f4f1143bdf 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -18,7 +18,7 @@ from __future__ import annotations import json -from typing import Dict, Optional, Union, TYPE_CHECKING, Tuple +from typing import Optional, Union, TYPE_CHECKING, Tuple from urllib.parse import urljoin import httpx @@ -52,16 +52,17 @@ class KiotaRequestAdapterHook(BaseHook): or you can pass a string as "v1.0" or "beta". """ - cached_request_adapters: Dict[str, Tuple[APIVersion, RequestAdapter]] = {} + cached_request_adapters: dict[str, tuple[APIVersion, RequestAdapter]] = {} default_conn_name: str = "msgraph_default" def __init__( self, conn_id: str = default_conn_name, - timeout: Optional[float] = None, - proxies: Optional[Dict] = None, - api_version: Optional[Union[APIVersion, str]] = None, - ) -> None: + timeout: float | None = None, + proxies: dict | None = None, + api_version: APIVersion | str | None = None, + ): + super().__init__() self.conn_id = conn_id self.timeout = timeout self.proxies = proxies @@ -83,7 +84,7 @@ def resolve_api_version_from_value( default, ) - def get_api_version(self, config: Dict) -> APIVersion: + def get_api_version(self, config: dict) -> APIVersion: if self._api_version is None: return self.resolve_api_version_from_value( api_version=config.get("api_version"), default=APIVersion.v1 @@ -97,7 +98,7 @@ def get_host(connection: Connection) -> str: return NationalClouds.Global.value @staticmethod - def to_httpx_proxies(proxies: Dict) -> Dict: + def to_httpx_proxies(proxies: dict) -> dict: proxies = proxies.copy() if proxies.get("http"): proxies["http://"] = proxies.pop("http") diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index f7ffa910579f7..b451b256860f8 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -17,18 +17,15 @@ # under the License. from __future__ import annotations +from types import NoneType from typing import ( - Dict, - Optional, Any, TYPE_CHECKING, Sequence, - Callable, - Type, + Callable, Optional, ) -from airflow import AirflowException -from airflow.exceptions import TaskDeferred +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.serialization.serializer import ( @@ -80,28 +77,28 @@ class MSGraphAsyncOperator(BaseOperator): def __init__( self, *, - url: Optional[str] = None, - response_type: Optional[ResponseType] = None, + url: str | None = None, + response_type: ResponseType | NoneType = None, response_handler: Callable[ - [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any ] = lambda response, error_map: response.json(), - path_parameters: Optional[Dict[str, Any]] = None, - url_template: Optional[str] = None, + path_parameters: Optional[dict[str, Any]] = None, + url_template: str | None = None, method: str = "GET", - query_parameters: Optional[Dict[str, QueryParams]] = None, - headers: Optional[Dict[str, str]] = None, - content: Optional[BytesIO] = None, + query_parameters: dict[str, QueryParams] | None = None, + headers: dict[str, str] | None = None, + content: BytesIO | NoneType = None, conn_id: str = KiotaRequestAdapterHook.default_conn_name, key: str = XCOM_RETURN_KEY, - timeout: Optional[float] = None, - proxies: Optional[Dict] = None, - api_version: Optional[APIVersion] = None, + timeout: float | None = None, + proxies: dict | None = None, + api_version: APIVersion | NoneType = None, result_processor: Callable[ [Context, Any], Any ] = lambda context, result: result, - serializer: Type[ResponseSerializer] = ResponseSerializer, + serializer: type[ResponseSerializer] = ResponseSerializer, **kwargs: Any, - ) -> None: + ): super().__init__(**kwargs) self.url = url self.response_type = response_type @@ -145,7 +142,7 @@ def execute(self, context: Context) -> None: def execute_complete( self, context: Context, - event: Optional[Dict[Any, Any]] = None, + event: dict[Any, Any] | None = None, ) -> Any: """ Callback for when the trigger fires - returns immediately. @@ -227,7 +224,7 @@ def push_xcom(self, context: Context, value) -> None: self.xcom_push(context=context, key=self.key, value=value) def pull_execute_complete( - self, context: Context, event: Optional[Dict[Any, Any]] = None + self, context: Context, event: dict[Any, Any] | None = None ) -> Any: self.results = list( self.xcom_pull( diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 71dfbcdec2c8a..66dde479f249f 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -1,14 +1,14 @@ -from typing import Optional, Dict, Any, Callable +from typing import Optional, Any, Callable from kiota_abstractions.response_handler import ResponseHandler, NativeResponseType -from kiota_abstractions.serialization import ParsableFactory # noqa: TC002 +from kiota_abstractions.serialization import ParsableFactory # type: ignore[TCH002] class CallableResponseHandler(ResponseHandler): def __init__( self, callable_function: Callable[ - [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any ], ): self.callable_function = callable_function @@ -16,6 +16,6 @@ def __init__( async def handle_response_async( self, response: NativeResponseType, - error_map: Optional[Dict[str, Optional[ParsableFactory]]], + error_map: Optional[dict[str, Optional[ParsableFactory]]], ) -> Any: return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 95b51937bf6da..35319d491b7c7 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -4,18 +4,18 @@ from contextlib import suppress from datetime import datetime from json import JSONDecodeError -from typing import Optional, Any +from typing import Any from uuid import UUID import pendulum class ResponseSerializer: - def __init__(self, encoding: Optional[str] = None): + def __init__(self, encoding: str | None = None): self.encoding = encoding or locale.getpreferredencoding() - def serialize(self, response) -> Optional[str]: - def convert(value) -> Optional[str]: + def serialize(self, response) -> str | None: + def convert(value) -> str | None: if value is not None: if isinstance(value, UUID): return str(value) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 07cc038a5150e..f26b1c5717fed 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -17,14 +17,13 @@ # under the License. from __future__ import annotations +from types import NoneType from typing import ( - Dict, Optional, Any, AsyncIterator, Sequence, Union, - Type, TYPE_CHECKING, Callable, ) @@ -87,22 +86,22 @@ class MSGraphTrigger(BaseTrigger): def __init__( self, - url: Optional[str] = None, - response_type: Optional[ResponseType] = None, + url: str | None = None, + response_type: ResponseType | NoneType = None, response_handler: Callable[ - [NativeResponseType, Optional[Dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any ] = lambda response, error_map: response.json(), - path_parameters: Optional[Dict[str, Any]] = None, - url_template: Optional[str] = None, + path_parameters: Optional[dict[str, Any]] = None, + url_template: str | None = None, method: str = "GET", - query_parameters: Optional[Dict[str, QueryParams]] = None, - headers: Optional[Dict[str, str]] = None, - content: Optional[BytesIO] = None, + query_parameters: dict[str, QueryParams] | None = None, + headers: dict[str, str] | None = None, + content: BytesIO | NoneType = None, conn_id: str = KiotaRequestAdapterHook.default_conn_name, - timeout: Optional[float] = None, - proxies: Optional[Dict] = None, - api_version: Union[APIVersion, str] = APIVersion.v1, - serializer: Union[str, Type[ResponseSerializer]] = ResponseSerializer, + timeout: float | None = None, + proxies: dict | None = None, + api_version: APIVersion | NoneType = None, + serializer: type[ResponseSerializer] = ResponseSerializer, ): super().__init__() self.hook = KiotaRequestAdapterHook( @@ -125,7 +124,7 @@ def __init__( )() @classmethod - def resolve_type(cls, value: Union[str, Type], default) -> Type: + def resolve_type(cls, value: str | type, default) -> type: if isinstance(value, str): try: return import_string(value) @@ -163,11 +162,11 @@ def conn_id(self) -> str: return self.hook.conn_id @property - def timeout(self) -> Optional[float]: + def timeout(self) -> float | None: return self.hook.timeout @property - def proxies(self) -> Optional[Dict]: + def proxies(self) -> dict | None: return self.hook.proxies @property @@ -242,7 +241,7 @@ def request_information(self) -> RequestInformation: return request_information @staticmethod - def error_mapping() -> Dict[str, Optional[ParsableFactory]]: + def error_mapping() -> dict[str, Optional[ParsableFactory]]: return { "4XX": APIError, "5XX": APIError, diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index ca96eb2ebea60..fbe474e20200e 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -18,17 +18,17 @@ from contextlib import contextmanager from copy import deepcopy from datetime import datetime -from typing import List, Tuple, Any, Iterable, Union, Optional +from typing import Any, Iterable, Union, Optional from unittest.mock import patch import pytest from kiota_http.httpx_request_adapter import HttpxRequestAdapter -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session # type: ignore[TCH002] from airflow.exceptions import TaskDeferred from airflow.models import Operator, TaskInstance from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook -from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.triggers.base import BaseTrigger, TriggerEvent # type: ignore[TCH001] from airflow.utils.session import NEW_SESSION from airflow.utils.state import TaskInstanceState from airflow.utils.xcom import XCOM_RETURN_KEY @@ -81,13 +81,13 @@ def patch_hook_and_request_adapter(self, response): yield @staticmethod - async def run_tigger(trigger: BaseTrigger) -> List[TriggerEvent]: + async def run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]: events = [] async for event in trigger.run(): events.append(event) return events - def execute_operator(self, operator: Operator) -> Tuple[Any, Any]: + def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) context = {"ti": task_instance} result = None diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index ab410e579f23a..703eba1ecfb78 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -21,7 +21,7 @@ import random import string from os.path import join, dirname -from typing import TypeVar, Iterable, Any, Optional, Dict +from typing import TypeVar, Iterable, Any, Optional from unittest.mock import MagicMock import pytest @@ -75,7 +75,7 @@ def mocked_connection(request, create_mock_connection): return create_mock_connection(request.param) -def mock_connection(schema: Optional[str] = None, host: Optional[str] = None) -> Connection: +def mock_connection(schema: str | None = None, host: str | None = None) -> Connection: connection = MagicMock(spec=Connection) connection.schema = schema connection.host = host @@ -115,7 +115,7 @@ def get_airflow_connection( login: str = "client_id", password: str = "client_secret", tenant_id: str = "tenant-id", - proxies: Optional[Dict] = None, + proxies: dict | None = None, api_version: APIVersion = APIVersion.v1): from airflow.models import Connection From 1f4f704e3af83d76863b317a48e6f7c95b33f36a Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 13 Mar 2024 19:54:13 +0100 Subject: [PATCH 004/105] refactor: Added some docstrings and fixed additional typing issues --- .../microsoft/azure/hooks/msgraph.py | 3 ++- .../microsoft/azure/operators/msgraph.py | 6 ++--- .../azure/serialization/response_handler.py | 24 ++++++++++++++++++- .../azure/serialization/serializer.py | 21 ++++++++++++++++ .../microsoft/azure/triggers/msgraph.py | 10 ++++---- tests/providers/microsoft/azure/base.py | 14 ++++++----- .../{test_trigger.py => test_msgraph.py} | 2 +- 7 files changed, 62 insertions(+), 18 deletions(-) rename tests/providers/microsoft/azure/triggers/{test_trigger.py => test_msgraph.py} (99%) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index a70f4f1143bdf..5e6e9211a1ac4 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -18,6 +18,7 @@ from __future__ import annotations import json +from types import NoneType from typing import Optional, Union, TYPE_CHECKING, Tuple from urllib.parse import urljoin @@ -75,7 +76,7 @@ def api_version(self) -> APIVersion: @staticmethod def resolve_api_version_from_value( - api_version: Union[APIVersion, str], default: Optional[APIVersion] = None + api_version: APIVersion | str, default: APIVersion | NoneType = None ) -> APIVersion: if isinstance(api_version, APIVersion): return api_version diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index b451b256860f8..682bca56ce0b6 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -145,9 +145,9 @@ def execute_complete( event: dict[Any, Any] | None = None, ) -> Any: """ - Callback for when the trigger fires - returns immediately. - Relies on trigger to throw an exception, otherwise it assumes execution was - successful. + Callback method that gets executed when MSGraphTrigger finishes execution. + + Relies on trigger to throw an exception, otherwise it assumes execution was successful. """ self.log.debug("context: %s", context) diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 66dde479f249f..caa4999fa1c7c 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -1,10 +1,32 @@ +# +# 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 typing import Optional, Any, Callable from kiota_abstractions.response_handler import ResponseHandler, NativeResponseType -from kiota_abstractions.serialization import ParsableFactory # type: ignore[TCH002] +# type: ignore[TCH002] +from kiota_abstractions.serialization import ParsableFactory class CallableResponseHandler(ResponseHandler): + """ + CallableResponseHandler executes the passed callable_function with the given response as parameter. + """ + def __init__( self, callable_function: Callable[ diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 35319d491b7c7..2dbdffd99238c 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -1,3 +1,20 @@ +# +# 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. import json import locale from base64 import b64encode @@ -11,6 +28,10 @@ class ResponseSerializer: + """ + ResponseSerializer serializes the response as a string. + """ + def __init__(self, encoding: str | None = None): self.encoding = encoding or locale.getpreferredencoding() diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index f26b1c5717fed..ffae95b878410 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -19,11 +19,9 @@ from types import NoneType from typing import ( - Optional, Any, AsyncIterator, Sequence, - Union, TYPE_CHECKING, Callable, ) @@ -32,7 +30,6 @@ from kiota_abstractions.method import Method from kiota_abstractions.request_information import RequestInformation from kiota_http.middleware.options import ResponseHandlerOption -from msgraph_core import APIVersion from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook from airflow.providers.microsoft.azure.serialization.response_handler import ( @@ -51,6 +48,7 @@ from kiota_abstractions.response_handler import NativeResponseType from kiota_abstractions.serialization import ParsableFactory from kiota_http.httpx_request_adapter import ResponseType + from msgraph_core import APIVersion class MSGraphTrigger(BaseTrigger): @@ -89,9 +87,9 @@ def __init__( url: str | None = None, response_type: ResponseType | NoneType = None, response_handler: Callable[ - [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any ] = lambda response, error_map: response.json(), - path_parameters: Optional[dict[str, Any]] = None, + path_parameters: dict[str, Any] | None = None, url_template: str | None = None, method: str = "GET", query_parameters: dict[str, QueryParams] | None = None, @@ -241,7 +239,7 @@ def request_information(self) -> RequestInformation: return request_information @staticmethod - def error_mapping() -> dict[str, Optional[ParsableFactory]]: + def error_mapping() -> dict[str, ParsableFactory | NoneType]: return { "4XX": APIError, "5XX": APIError, diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index fbe474e20200e..e507dfe528f74 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -18,17 +18,19 @@ from contextlib import contextmanager from copy import deepcopy from datetime import datetime -from typing import Any, Iterable, Union, Optional +from typing import Any, Iterable, Optional from unittest.mock import patch import pytest from kiota_http.httpx_request_adapter import HttpxRequestAdapter -from sqlalchemy.orm import Session # type: ignore[TCH002] +# type: ignore[TCH002] +from sqlalchemy.orm import Session from airflow.exceptions import TaskDeferred from airflow.models import Operator, TaskInstance from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook -from airflow.triggers.base import BaseTrigger, TriggerEvent # type: ignore[TCH001] +# type: ignore[TCH001] +from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.session import NEW_SESSION from airflow.utils.state import TaskInstanceState from airflow.utils.xcom import XCOM_RETURN_KEY @@ -40,13 +42,13 @@ class MockedTaskInstance(TaskInstance): def xcom_pull( self, - task_ids: Optional[Union[Iterable[str], str]] = None, + task_ids: Iterable[str] | str | None = None, dag_id: Optional[str] = None, key: str = XCOM_RETURN_KEY, include_prior_dates: bool = False, session: Session = NEW_SESSION, *, - map_indexes: Optional[Union[Iterable[int], int]] = None, + map_indexes: Iterable[int] | int | None = None, default: Optional[Any] = None, ) -> Any: self.task_id = task_ids @@ -57,7 +59,7 @@ def xcom_push( self, key: str, value: Any, - execution_date: Optional[datetime] = None, + execution_date: datetime | None = None, session: Session = NEW_SESSION, ) -> None: self.values[f"{self.task_id}_{self.dag_id}_{key}"] = value diff --git a/tests/providers/microsoft/azure/triggers/test_trigger.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py similarity index 99% rename from tests/providers/microsoft/azure/triggers/test_trigger.py rename to tests/providers/microsoft/azure/triggers/test_msgraph.py index ad6d12d267f06..1cb93e8e18252 100644 --- a/tests/providers/microsoft/azure/triggers/test_trigger.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -19,7 +19,7 @@ from base64 import b64encode from unittest.mock import patch -from airflow import AirflowException +from airflow.exceptions import AirflowException from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base From a32d7df91b80df110406e2f0890976b6916bf8fb Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 09:47:08 +0100 Subject: [PATCH 005/105] refactor: Fixed more static checks --- airflow/providers/microsoft/azure/operators/msgraph.py | 4 ++-- .../microsoft/azure/serialization/response_handler.py | 9 +++++---- airflow/providers/microsoft/azure/triggers/msgraph.py | 4 ++-- tests/providers/microsoft/azure/base.py | 6 +++--- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 682bca56ce0b6..d7adc1d720d18 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -80,9 +80,9 @@ def __init__( url: str | None = None, response_type: ResponseType | NoneType = None, response_handler: Callable[ - [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any ] = lambda response, error_map: response.json(), - path_parameters: Optional[dict[str, Any]] = None, + path_parameters: dict[str, Any] | None = None, url_template: str | None = None, method: str = "GET", query_parameters: dict[str, QueryParams] | None = None, diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index caa4999fa1c7c..9e9e04dd449a6 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -15,7 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import Optional, Any, Callable +from types import NoneType +from typing import Any, Callable from kiota_abstractions.response_handler import ResponseHandler, NativeResponseType # type: ignore[TCH002] @@ -24,13 +25,13 @@ class CallableResponseHandler(ResponseHandler): """ - CallableResponseHandler executes the passed callable_function with the given response as parameter. + CallableResponseHandler executes the passed callable_function with response as parameter. """ def __init__( self, callable_function: Callable[ - [NativeResponseType, Optional[dict[str, Optional[ParsableFactory]]]], Any + [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any ], ): self.callable_function = callable_function @@ -38,6 +39,6 @@ def __init__( async def handle_response_async( self, response: NativeResponseType, - error_map: Optional[dict[str, Optional[ParsableFactory]]], + error_map: dict[str, ParsableFactory | NoneType] | None, ) -> Any: return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index ffae95b878410..b300e7aa9a8a9 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -131,7 +131,7 @@ def resolve_type(cls, value: str | type, default) -> type: return value or default def serialize(self) -> tuple[str, dict[str, Any]]: - """Serializes HttpTrigger arguments and classpath.""" + """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__}", @@ -172,7 +172,7 @@ def api_version(self) -> APIVersion: return self.hook.api_version async def run(self) -> AsyncIterator[TriggerEvent]: - """Makes a series of asynchronous http calls via a KiotaRequestAdapterHook.""" + """ Make a series of asynchronous HTTP calls via a KiotaRequestAdapterHook.""" try: response = await self.execute() diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index e507dfe528f74..b0ac6277148f2 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from copy import deepcopy from datetime import datetime -from typing import Any, Iterable, Optional +from typing import Any, Iterable from unittest.mock import patch import pytest @@ -43,13 +43,13 @@ class MockedTaskInstance(TaskInstance): def xcom_pull( self, task_ids: Iterable[str] | str | None = None, - dag_id: Optional[str] = None, + dag_id: str | None = None, key: str = XCOM_RETURN_KEY, include_prior_dates: bool = False, session: Session = NEW_SESSION, *, map_indexes: Iterable[int] | int | None = None, - default: Optional[Any] = None, + default: Any | None = None, ) -> Any: self.task_id = task_ids self.dag_id = dag_id From 04d6b85494b0b9d3973564d4ac5abb718ac32cc7 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 10:13:41 +0100 Subject: [PATCH 006/105] refactor: Added license on top of test serializer and fixed import --- .../microsoft/azure/serializer/__init__.py | 16 +++++++++++++++ .../azure/serializer/test_serializer.py | 20 +++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/azure/serializer/__init__.py b/tests/providers/microsoft/azure/serializer/__init__.py index e69de29bb2d1d..13a83393a9124 100644 --- a/tests/providers/microsoft/azure/serializer/__init__.py +++ b/tests/providers/microsoft/azure/serializer/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/providers/microsoft/azure/serializer/test_serializer.py b/tests/providers/microsoft/azure/serializer/test_serializer.py index 3318666cc968a..17814b7bce4ec 100644 --- a/tests/providers/microsoft/azure/serializer/test_serializer.py +++ b/tests/providers/microsoft/azure/serializer/test_serializer.py @@ -1,12 +1,28 @@ +# 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. import locale from base64 import b64encode, b64decode from datetime import datetime from uuid import uuid4 import pendulum -from airflow.providers.microsoft.msgraph.serialization.serializer import ResponseSerializer -from tests.unit.conftest import load_json, load_file +from airflow.providers.microsoft.azure.serialization.serializer import ResponseSerializer +from tests.providers.microsoft.conftest import load_json, load_file class TestResponseSerializer: From b77e1527d8a484cbf7fd9c1406ab1ecbabf527f0 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 10:16:57 +0100 Subject: [PATCH 007/105] Revert "refactor: Added license on top of test serializer and fixed import" This reverts commit 04d6b85494b0b9d3973564d4ac5abb718ac32cc7. --- .../microsoft/azure/serializer/__init__.py | 16 --------------- .../azure/serializer/test_serializer.py | 20 ++----------------- 2 files changed, 2 insertions(+), 34 deletions(-) diff --git a/tests/providers/microsoft/azure/serializer/__init__.py b/tests/providers/microsoft/azure/serializer/__init__.py index 13a83393a9124..e69de29bb2d1d 100644 --- a/tests/providers/microsoft/azure/serializer/__init__.py +++ b/tests/providers/microsoft/azure/serializer/__init__.py @@ -1,16 +0,0 @@ -# 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. diff --git a/tests/providers/microsoft/azure/serializer/test_serializer.py b/tests/providers/microsoft/azure/serializer/test_serializer.py index 17814b7bce4ec..3318666cc968a 100644 --- a/tests/providers/microsoft/azure/serializer/test_serializer.py +++ b/tests/providers/microsoft/azure/serializer/test_serializer.py @@ -1,28 +1,12 @@ -# 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. import locale from base64 import b64encode, b64decode from datetime import datetime from uuid import uuid4 import pendulum +from airflow.providers.microsoft.msgraph.serialization.serializer import ResponseSerializer -from airflow.providers.microsoft.azure.serialization.serializer import ResponseSerializer -from tests.providers.microsoft.conftest import load_json, load_file +from tests.unit.conftest import load_json, load_file class TestResponseSerializer: From 90c498b739520adb936e6734f3128eed26f235e2 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 10:36:18 +0100 Subject: [PATCH 008/105] refactor: Added license on top of serializer files and fixed additional static checks --- .../microsoft/azure/serialization/__init__.py | 16 +++++++++ .../microsoft/azure/serializer/__init__.py | 16 +++++++++ .../azure/serializer/test_serializer.py | 35 ++++++++++++++++--- 3 files changed, 62 insertions(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/serialization/__init__.py b/airflow/providers/microsoft/azure/serialization/__init__.py index e69de29bb2d1d..13a83393a9124 100644 --- a/airflow/providers/microsoft/azure/serialization/__init__.py +++ b/airflow/providers/microsoft/azure/serialization/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/providers/microsoft/azure/serializer/__init__.py b/tests/providers/microsoft/azure/serializer/__init__.py index e69de29bb2d1d..13a83393a9124 100644 --- a/tests/providers/microsoft/azure/serializer/__init__.py +++ b/tests/providers/microsoft/azure/serializer/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/providers/microsoft/azure/serializer/test_serializer.py b/tests/providers/microsoft/azure/serializer/test_serializer.py index 3318666cc968a..445d9fe75ac01 100644 --- a/tests/providers/microsoft/azure/serializer/test_serializer.py +++ b/tests/providers/microsoft/azure/serializer/test_serializer.py @@ -1,12 +1,30 @@ +# 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 locale -from base64 import b64encode, b64decode +from base64 import b64decode, b64encode from datetime import datetime from uuid import uuid4 import pendulum -from airflow.providers.microsoft.msgraph.serialization.serializer import ResponseSerializer -from tests.unit.conftest import load_json, load_file +from airflow.providers.microsoft.azure.serialization.serializer import ResponseSerializer +from tests.providers.microsoft.conftest import load_file, load_json class TestResponseSerializer: @@ -21,12 +39,19 @@ def test_serialize_when_bytes_then_base64_encoded(self): def test_serialize_when_dict_with_uuid_datatime_and_pendulum_then_json(self): id = uuid4() - response = {"id": id, "creationDate": datetime(2024, 2, 5), "modificationTime": pendulum.datetime(2024, 2, 5)} + response = { + "id": id, + "creationDate": datetime(2024, 2, 5), + "modificationTime": pendulum.datetime(2024, 2, 5), + } actual = ResponseSerializer().serialize(response) assert isinstance(actual, str) - assert actual == f'{{"id": "{id}", "creationDate": "2024-02-05T00:00:00", "modificationTime": "2024-02-05T00:00:00+00:00"}}' + assert ( + actual + == f'{{"id": "{id}", "creationDate": "2024-02-05T00:00:00", "modificationTime": "2024-02-05T00:00:00+00:00"}}' + ) def test_deserialize_when_json(self): response = load_file("resources", "users.json") From 07034dbdc2b782c74b20d2caffc9ef288d7adb1c Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 10:36:59 +0100 Subject: [PATCH 009/105] refactor: Added new line at end of json test files --- tests/providers/microsoft/azure/resources/next_users.json | 2 +- tests/providers/microsoft/azure/resources/users.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/azure/resources/next_users.json b/tests/providers/microsoft/azure/resources/next_users.json index f9376541b6f15..3a88cf08b2c34 100644 --- a/tests/providers/microsoft/azure/resources/next_users.json +++ b/tests/providers/microsoft/azure/resources/next_users.json @@ -1 +1 @@ -{"@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Leonardo DiCaprio", "mailNickname": "LeoD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Meryl Streep", "mailNickname": "MerylS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Tom Hanks", "mailNickname": "TomH"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Jennifer Lawrence", "mailNickname": "JenniferL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Denzel Washington", "mailNickname": "DenzelW"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Angelina Jolie", "mailNickname": "AngelinaJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Johnny Depp", "mailNickname": "JohnnyD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Scarlett Johansson", "mailNickname": "ScarlettJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Daniel Craig", "mailNickname": "DanielC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Charlize Theron", "mailNickname": "CharlizeT"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "George Clooney", "mailNickname": "GeorgeC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/149136c3-62f0-4d27-94f1-8a27bfc4cf73/Microsoft.DirectoryServices.User", "displayName": "Julia Roberts", "mailNickname": "JuliaR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/08f9d41c-a653-4de3-85c2-324eb53bcbff/Microsoft.DirectoryServices.User", "displayName": "Ryan Reynolds", "mailNickname": "RyanR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/cf5ce2ea-132b-4699-bd8f-6bfcac619fc3/Microsoft.DirectoryServices.User", "displayName": "Nicole Kidman", "mailNickname": "NicoleK"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e255041c-28e5-4abc-bb03-09205096cd87/Microsoft.DirectoryServices.User", "displayName": "Sean Connery", "mailNickname": "SeanC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/52f8faad-6b39-4cbb-8614-e968f5af9e9e/Microsoft.DirectoryServices.User", "displayName": "Emma Stone", "mailNickname": "EmmaS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e9bf2db3-4149-4396-aafd-278a0212179a/Microsoft.DirectoryServices.User", "displayName": "Robert Downey Jr.", "mailNickname": "RobertDJr"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/f66aba7d-5284-418d-a025-6d6f0639350b/Microsoft.DirectoryServices.User", "displayName": "Cate Blanchett", "mailNickname": "CateB"}]} \ No newline at end of file +{"@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Leonardo DiCaprio", "mailNickname": "LeoD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Meryl Streep", "mailNickname": "MerylS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Tom Hanks", "mailNickname": "TomH"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Jennifer Lawrence", "mailNickname": "JenniferL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Denzel Washington", "mailNickname": "DenzelW"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Angelina Jolie", "mailNickname": "AngelinaJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Johnny Depp", "mailNickname": "JohnnyD"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Scarlett Johansson", "mailNickname": "ScarlettJ"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Daniel Craig", "mailNickname": "DanielC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Charlize Theron", "mailNickname": "CharlizeT"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "George Clooney", "mailNickname": "GeorgeC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/149136c3-62f0-4d27-94f1-8a27bfc4cf73/Microsoft.DirectoryServices.User", "displayName": "Julia Roberts", "mailNickname": "JuliaR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/08f9d41c-a653-4de3-85c2-324eb53bcbff/Microsoft.DirectoryServices.User", "displayName": "Ryan Reynolds", "mailNickname": "RyanR"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/cf5ce2ea-132b-4699-bd8f-6bfcac619fc3/Microsoft.DirectoryServices.User", "displayName": "Nicole Kidman", "mailNickname": "NicoleK"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e255041c-28e5-4abc-bb03-09205096cd87/Microsoft.DirectoryServices.User", "displayName": "Sean Connery", "mailNickname": "SeanC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/52f8faad-6b39-4cbb-8614-e968f5af9e9e/Microsoft.DirectoryServices.User", "displayName": "Emma Stone", "mailNickname": "EmmaS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/e9bf2db3-4149-4396-aafd-278a0212179a/Microsoft.DirectoryServices.User", "displayName": "Robert Downey Jr.", "mailNickname": "RobertDJr"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/f66aba7d-5284-418d-a025-6d6f0639350b/Microsoft.DirectoryServices.User", "displayName": "Cate Blanchett", "mailNickname": "CateB"}]} diff --git a/tests/providers/microsoft/azure/resources/users.json b/tests/providers/microsoft/azure/resources/users.json index abaceebcf9a54..617e6f8420c62 100644 --- a/tests/providers/microsoft/azure/resources/users.json +++ b/tests/providers/microsoft/azure/resources/users.json @@ -1 +1 @@ -{"@odata.nextLink": "https://graph.microsoft.com/v1.0/users/delta()?$skiptoken=qLMhmnoTon81CQ1VVWyx9MNESqEIMKNpEWZWAfnn5F7tBNFuSgWh_pXZOweu67nEThGR0yQewi_a3Ixe75S6PoB8pdllphCEev0fMe5Uc1lWMtn3byOS8_OPTzPGZIZ17x-dVyxaE_4I55YyLJ0cgBxg8wsBrkYgaNE9vy5Su2HeCKxJODDQk4zRgP8QGo0pZatReTpqisVbrW5Gl1H_Xgy4lhenv1SmoRcBQtWBa5iAh-MURoaTo7i0kQjFhH6SCrkjBkfkRFVy9dafOOt2Owbxfn5hKGfEnfmG0RBmgdUsZPgX-ap0mjjf7PjExoxMek4CDnb8Yv737oGkh9C_G0XTJGeGxPBbkD-w4SaQookde4yxOzceAw1MuamBy63uJdbXt1ul61tDvfPwrJVHq99FxGU1n-i_RfHh65nRCHju3N3ApOFKrAi6933l3VupyaXsvmm5pCPh0T70dYK01CYktBce8Mc1HaVqB7SR-R9-X4PHYfozPWeqv4hng4YEvqA41XjRPUzOaS1VTH08k-HhR9ENpqw5UTFnimAu5RbMT5fTbUAMTQC2XcWF_5aDgjuw8D2VQvr5VsB7qmu4mgGb_dNrHM47QyJCKY2QcgLmmsTnr5Z-7Qe2AGGy1b5DREBVoLnSL-aJ_6m9TAlYD9oGityZbDJ1ssVxS0XsGYxAwSa5z_E_lgedr_ikHc0zzDAdj3TgLGD5gWIjwvxnkNEXa0onk-jqGfANEDFh0vGuDz1mlgwiHGIKN-QjmfTX-Equ_uY2lbuPcIXTVpdrgQob5BaXONZ78uVh-AIlm96PPqbwNj1QMf3EE0DfgGaMIloFdByjSD5FjwQ_6COxueUJw2iUtKh8l4fpja7yOs4QCP5_tPvUUZT26ylLHjbPZO7I1TOaJb9OTBM-kplBJitW3SAm7DJM9jocA89Iqbm_ap6mlvEET2cyw6Dn7j2AP_0NsdoP1MJn6H0JE6HOJddbUVuzhba58QBujgHsEHo.WPJMTkIZstDbZSnvkZm1eX2ASvooRj4BgY-GW_pQEZo", "@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Sylvester Stallone", "mailNickname": "SylvesterS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Jason Statham", "mailNickname": "JasonS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Jet Li", "mailNickname": "JetL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Dolph Lundgren", "mailNickname": "DolphL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Randy Couture", "mailNickname": "RandyC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Terry Crews", "mailNickname": "TerryC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Arnold Schwarzenegger", "mailNickname": "ArnoldS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Wesley Snipes", "mailNickname": "WesleyS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Mel Gibson", "mailNickname": "MelG"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Harrison Ford", "mailNickname": "HarrisonF"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "Antonio Banderas", "mailNickname": "AntonioB"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/34c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Chuck Norris", "mailNickname": "ChuckN"}]} \ No newline at end of file +{"@odata.nextLink": "https://graph.microsoft.com/v1.0/users/delta()?$skiptoken=qLMhmnoTon81CQ1VVWyx9MNESqEIMKNpEWZWAfnn5F7tBNFuSgWh_pXZOweu67nEThGR0yQewi_a3Ixe75S6PoB8pdllphCEev0fMe5Uc1lWMtn3byOS8_OPTzPGZIZ17x-dVyxaE_4I55YyLJ0cgBxg8wsBrkYgaNE9vy5Su2HeCKxJODDQk4zRgP8QGo0pZatReTpqisVbrW5Gl1H_Xgy4lhenv1SmoRcBQtWBa5iAh-MURoaTo7i0kQjFhH6SCrkjBkfkRFVy9dafOOt2Owbxfn5hKGfEnfmG0RBmgdUsZPgX-ap0mjjf7PjExoxMek4CDnb8Yv737oGkh9C_G0XTJGeGxPBbkD-w4SaQookde4yxOzceAw1MuamBy63uJdbXt1ul61tDvfPwrJVHq99FxGU1n-i_RfHh65nRCHju3N3ApOFKrAi6933l3VupyaXsvmm5pCPh0T70dYK01CYktBce8Mc1HaVqB7SR-R9-X4PHYfozPWeqv4hng4YEvqA41XjRPUzOaS1VTH08k-HhR9ENpqw5UTFnimAu5RbMT5fTbUAMTQC2XcWF_5aDgjuw8D2VQvr5VsB7qmu4mgGb_dNrHM47QyJCKY2QcgLmmsTnr5Z-7Qe2AGGy1b5DREBVoLnSL-aJ_6m9TAlYD9oGityZbDJ1ssVxS0XsGYxAwSa5z_E_lgedr_ikHc0zzDAdj3TgLGD5gWIjwvxnkNEXa0onk-jqGfANEDFh0vGuDz1mlgwiHGIKN-QjmfTX-Equ_uY2lbuPcIXTVpdrgQob5BaXONZ78uVh-AIlm96PPqbwNj1QMf3EE0DfgGaMIloFdByjSD5FjwQ_6COxueUJw2iUtKh8l4fpja7yOs4QCP5_tPvUUZT26ylLHjbPZO7I1TOaJb9OTBM-kplBJitW3SAm7DJM9jocA89Iqbm_ap6mlvEET2cyw6Dn7j2AP_0NsdoP1MJn6H0JE6HOJddbUVuzhba58QBujgHsEHo.WPJMTkIZstDbZSnvkZm1eX2ASvooRj4BgY-GW_pQEZo", "@odata.context": "https://graph.microsoft.com/v1.0/$metadata#users(displayName,description,mailNickname)", "value": [{"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/c52a9941-e5cb-49cb-9972-6f45cd7cd447/Microsoft.DirectoryServices.User", "displayName": "Sylvester Stallone", "mailNickname": "SylvesterS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/abee3661-4525-4491-8e0f-e1589383aea7/Microsoft.DirectoryServices.User", "displayName": "Jason Statham", "mailNickname": "JasonS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6794f45f-949e-4e8b-b01c-03d777e1cbf8/Microsoft.DirectoryServices.User", "displayName": "Jet Li", "mailNickname": "JetL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/ad31af2e-b69a-4ff7-8aee-b962cc739210/Microsoft.DirectoryServices.User", "displayName": "Dolph Lundgren", "mailNickname": "DolphL"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/6aed18fc-f39b-4762-85d1-1525ccdf4823/Microsoft.DirectoryServices.User", "displayName": "Randy Couture", "mailNickname": "RandyC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/44c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Terry Crews", "mailNickname": "TerryC"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d1db76e3-06e7-430a-a97e-7b6572f5fb15/Microsoft.DirectoryServices.User", "displayName": "Arnold Schwarzenegger", "mailNickname": "ArnoldS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/2b62c4a9-3c0c-4d60-8e3f-cb698d8ba9fc/Microsoft.DirectoryServices.User", "displayName": "Wesley Snipes", "mailNickname": "WesleyS"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/d982d3c5-4ea8-47bb-81b1-77ed0b107f30/Microsoft.DirectoryServices.User", "displayName": "Mel Gibson", "mailNickname": "MelG"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/0388216c-76ca-4031-a620-3af3df529485/Microsoft.DirectoryServices.User", "displayName": "Harrison Ford", "mailNickname": "HarrisonF"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/bbde9cbb-9688-4f07-af76-660244830541/Microsoft.DirectoryServices.User", "displayName": "Antonio Banderas", "mailNickname": "AntonioB"}, {"@odata.type": "#microsoft.graph.user", "@odata.id": "https://graph.microsoft.com/v2/c34aa217-8e78-4c5f-91a7-29e01e0d97d8/directoryObjects/34c76cb8-7118-4211-99d0-dd6651ce2fe6/Microsoft.DirectoryServices.User", "displayName": "Chuck Norris", "mailNickname": "ChuckN"}]} From a03378f16ecbf99879286ac395a4f453da01bac1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 11:15:00 +0100 Subject: [PATCH 010/105] refactor: Try fixing docstrings on operator and serializer --- airflow/providers/microsoft/azure/operators/msgraph.py | 4 ++-- airflow/providers/microsoft/azure/serialization/serializer.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index d7adc1d720d18..d1f83950fb890 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -145,9 +145,9 @@ def execute_complete( event: dict[Any, Any] | None = None, ) -> Any: """ - Callback method that gets executed when MSGraphTrigger finishes execution. + Execute callback when MSGraphTrigger finishes execution. - Relies on trigger to throw an exception, otherwise it assumes execution was successful. + This method gets executed automatically when MSGraphTrigger completes its execution. """ self.log.debug("context: %s", context) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 2dbdffd99238c..56b63e65fd654 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -29,7 +29,7 @@ class ResponseSerializer: """ - ResponseSerializer serializes the response as a string. + ResponseSerializer serializes the response as a string. """ def __init__(self, encoding: str | None = None): From f1e341438e611286ca7d804cb0fea0b903b4c92c Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 11:21:16 +0100 Subject: [PATCH 011/105] refactor: Replaced NoneType with None --- airflow/providers/microsoft/azure/hooks/msgraph.py | 5 ++--- .../providers/microsoft/azure/operators/msgraph.py | 12 +++++------- .../providers/microsoft/azure/triggers/msgraph.py | 11 +++++------ 3 files changed, 12 insertions(+), 16 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index 5e6e9211a1ac4..f314d1091da29 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -18,8 +18,7 @@ from __future__ import annotations import json -from types import NoneType -from typing import Optional, Union, TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING from urllib.parse import urljoin import httpx @@ -76,7 +75,7 @@ def api_version(self) -> APIVersion: @staticmethod def resolve_api_version_from_value( - api_version: APIVersion | str, default: APIVersion | NoneType = None + api_version: APIVersion | str, default: APIVersion | None = None ) -> APIVersion: if isinstance(api_version, APIVersion): return api_version diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index d1f83950fb890..b0dd30c72b2e4 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -17,13 +17,11 @@ # under the License. from __future__ import annotations -from types import NoneType from typing import ( Any, TYPE_CHECKING, Sequence, - Callable, Optional, -) + Callable, ) from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import BaseOperator @@ -78,21 +76,21 @@ def __init__( self, *, url: str | None = None, - response_type: ResponseType | NoneType = None, + response_type: ResponseType | None = None, response_handler: Callable[ - [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any + [NativeResponseType, dict[str, ParsableFactory | None] | None], Any ] = lambda response, error_map: response.json(), 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, - content: BytesIO | NoneType = None, + content: 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 | NoneType = None, + api_version: APIVersion | None = None, result_processor: Callable[ [Context, Any], Any ] = lambda context, result: result, diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index b300e7aa9a8a9..037f9c103f210 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -from types import NoneType from typing import ( Any, AsyncIterator, @@ -85,20 +84,20 @@ class MSGraphTrigger(BaseTrigger): def __init__( self, url: str | None = None, - response_type: ResponseType | NoneType = None, + response_type: ResponseType | None = None, response_handler: Callable[ - [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any + [NativeResponseType, dict[str, ParsableFactory | None] | None], Any ] = lambda response, error_map: response.json(), 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, - content: BytesIO | NoneType = None, + content: BytesIO | None = None, conn_id: str = KiotaRequestAdapterHook.default_conn_name, timeout: float | None = None, proxies: dict | None = None, - api_version: APIVersion | NoneType = None, + api_version: APIVersion | None = None, serializer: type[ResponseSerializer] = ResponseSerializer, ): super().__init__() @@ -239,7 +238,7 @@ def request_information(self) -> RequestInformation: return request_information @staticmethod - def error_mapping() -> dict[str, ParsableFactory | NoneType]: + def error_mapping() -> dict[str, ParsableFactory | None]: return { "4XX": APIError, "5XX": APIError, From 3e01b891f262cf8e065930a6edefbbd52930041d Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 11:24:51 +0100 Subject: [PATCH 012/105] refactor: Made type unions Python 3.8 compatible --- .../azure/serialization/response_handler.py | 16 +++++++--------- .../microsoft/azure/serialization/serializer.py | 12 ++++++------ 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 9e9e04dd449a6..037339459fe0f 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -15,30 +15,28 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from types import NoneType +from __future__ import annotations + from typing import Any, Callable -from kiota_abstractions.response_handler import ResponseHandler, NativeResponseType +from kiota_abstractions.response_handler import NativeResponseType, ResponseHandler + # type: ignore[TCH002] from kiota_abstractions.serialization import ParsableFactory class CallableResponseHandler(ResponseHandler): """ - CallableResponseHandler executes the passed callable_function with response as parameter. + CallableResponseHandler executes the passed callable_function with response as parameter. """ def __init__( self, - callable_function: Callable[ - [NativeResponseType, dict[str, ParsableFactory | NoneType] | None], Any - ], + callable_function: Callable[[NativeResponseType, dict[str, (ParsableFactory, None)]], Any], ): self.callable_function = callable_function async def handle_response_async( - self, - response: NativeResponseType, - error_map: dict[str, ParsableFactory | NoneType] | None, + self, response: NativeResponseType, error_map: dict[str, (ParsableFactory, None)] = None ) -> Any: return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 56b63e65fd654..3e0b84f07590d 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -15,6 +15,8 @@ # 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 @@ -32,11 +34,11 @@ class ResponseSerializer: ResponseSerializer serializes the response as a string. """ - def __init__(self, encoding: str | None = None): + def __init__(self, encoding: (str, None) = None): self.encoding = encoding or locale.getpreferredencoding() - def serialize(self, response) -> str | None: - def convert(value) -> str | None: + def serialize(self, response) -> (str, None): + def convert(value) -> (str, None): if value is not None: if isinstance(value, UUID): return str(value) @@ -44,9 +46,7 @@ def convert(value) -> str | None: 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!" - ) + raise TypeError(f"Object of type {type(value)} is not JSON serializable!") return None if response is not None: From d7198db9019745c1a6b2047475e4dc9544c63797 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 12:08:24 +0100 Subject: [PATCH 013/105] refactor: Reformatted some files to comply with static checks formatting --- .../microsoft/azure/hooks/test_msgraph.py | 18 +++++++------ .../microsoft/azure/operators/test_msgraph.py | 6 +++-- .../microsoft/azure/triggers/test_msgraph.py | 25 +++++++++++++------ tests/providers/microsoft/conftest.py | 17 +++++++------ 4 files changed, 40 insertions(+), 26 deletions(-) diff --git a/tests/providers/microsoft/azure/hooks/test_msgraph.py b/tests/providers/microsoft/azure/hooks/test_msgraph.py index a3c8c7ac9dba7..c52db83f8cc35 100644 --- a/tests/providers/microsoft/azure/hooks/test_msgraph.py +++ b/tests/providers/microsoft/azure/hooks/test_msgraph.py @@ -14,6 +14,8 @@ # 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.mock import patch from kiota_http.httpx_request_adapter import HttpxRequestAdapter @@ -26,8 +28,8 @@ class TestKiotaRequestAdapterHook: def test_get_conn(self): with (patch( - "airflow.hooks.base.BaseHook.get_connection", - side_effect=get_airflow_connection, + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, )): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_conn() @@ -37,8 +39,8 @@ def test_get_conn(self): def test_api_version(self): with (patch( - "airflow.hooks.base.BaseHook.get_connection", - side_effect=get_airflow_connection, + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, )): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") @@ -46,8 +48,8 @@ def test_api_version(self): def test_get_api_version_when_empty_config_dict(self): with (patch( - "airflow.hooks.base.BaseHook.get_connection", - side_effect=get_airflow_connection, + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, )): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_api_version({}) @@ -56,8 +58,8 @@ def test_get_api_version_when_empty_config_dict(self): def test_get_api_version_when_api_version_in_config_dict(self): with (patch( - "airflow.hooks.base.BaseHook.get_connection", - side_effect=get_airflow_connection, + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, )): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_api_version({"api_version": "beta"}) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 1bd5602a46a09..54936b80740b1 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -14,6 +14,8 @@ # 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 @@ -24,7 +26,7 @@ from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base -from tests.providers.microsoft.conftest import load_json, mock_json_response, load_file, mock_response +from tests.providers.microsoft.conftest import load_file, load_json, mock_json_response, mock_response class TestMSGraphAsyncOperator(Base): @@ -38,7 +40,7 @@ def test_run_when_expression_is_valid(self): task_id="users_delta", conn_id="msgraph_api", url="users", - result_processor=lambda context, result: result.get("value") + result_processor=lambda context, result: result.get("value"), ) results, events = self.execute_operator(operator) diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index 1cb93e8e18252..ab95ddded1cf6 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -14,6 +14,8 @@ # 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 @@ -23,8 +25,13 @@ from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base -from tests.providers.microsoft.conftest import load_json, mock_json_response, get_airflow_connection, \ - load_file, mock_response +from tests.providers.microsoft.conftest import ( + get_airflow_connection, + load_file, + load_json, + mock_json_response, + mock_response, +) class TestMSGraphTrigger(Base): @@ -70,7 +77,9 @@ def test_run_when_url_with_response_type_bytes(self): response = mock_response(200, content) with self.patch_hook_and_request_adapter(response): - url = "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content" + url = ( + "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content" + ) trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) @@ -82,8 +91,8 @@ def test_run_when_url_with_response_type_bytes(self): def test_serialize(self): with patch( - "airflow.hooks.base.BaseHook.get_connection", - side_effect=get_airflow_connection, + "airflow.hooks.base.BaseHook.get_connection", + side_effect=get_airflow_connection, ): url = "https://graph.microsoft.com/v1.0/me/drive/items" trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") @@ -96,12 +105,12 @@ def test_serialize(self): "url": "https://graph.microsoft.com/v1.0/me/drive/items", "path_parameters": None, "url_template": None, - "method": 'GET', + "method": "GET", "query_parameters": None, "headers": None, "content": None, - "response_type": 'bytes', - "conn_id": 'msgraph_api', + "response_type": "bytes", + "conn_id": "msgraph_api", "timeout": None, "proxies": None, "api_version": "v1.0", diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index 703eba1ecfb78..022fde060f0c6 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -20,8 +20,8 @@ import json import random import string -from os.path import join, dirname -from typing import TypeVar, Iterable, Any, Optional +from os.path import dirname, join +from typing import Any, Iterable, TypeVar from unittest.mock import MagicMock import pytest @@ -111,12 +111,13 @@ def load_file(*locations: Iterable[str], mode="r", encoding="utf-8"): def get_airflow_connection( - conn_id: str, - login: str = "client_id", - password: str = "client_secret", - tenant_id: str = "tenant-id", - proxies: dict | None = None, - api_version: APIVersion = APIVersion.v1): + conn_id: str, + login: str = "client_id", + password: str = "client_secret", + tenant_id: str = "tenant-id", + proxies: (dict, None) = None, + api_version: APIVersion = APIVersion.v1, +): from airflow.models import Connection return Connection( From bcaa2ff4cde5669a2288c1094c15df518a10019d Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 12:18:08 +0100 Subject: [PATCH 014/105] refactor: Reformatted base to comply with static checks formatting --- tests/providers/microsoft/azure/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index b0ac6277148f2..335bda425363e 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -14,6 +14,8 @@ # 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 contextlib import contextmanager from copy import deepcopy @@ -74,8 +76,9 @@ def teardown_method(self, method): @contextmanager def patch_hook_and_request_adapter(self, response): - with patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection), \ - patch.object(HttpxRequestAdapter, "get_http_response_message") as mock_get_http_response: + with patch( + "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection + ), patch.object(HttpxRequestAdapter, "get_http_response_message") as mock_get_http_response: if isinstance(response, Exception): mock_get_http_response.side_effect = response else: From b319be45a230d27a84a77ad50af0eeee9cbccd3b Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 14:41:53 +0100 Subject: [PATCH 015/105] refactor: Added msgraph-core dependency to provider.yaml --- airflow/providers/microsoft/azure/provider.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 956fdba293593..c9084789afe0f 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -97,6 +97,7 @@ dependencies: - azure-mgmt-datafactory>=2.0.0 - azure-mgmt-containerregistry>=8.0.0 - azure-mgmt-containerinstance>=9.0.0 + - msgraph-core>=1.0.0 devel-dependencies: - pywinrm From 7dc3e7e9beef4c9f49352e4fdd896364f6bf5852 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 14:51:21 +0100 Subject: [PATCH 016/105] refactor: Added msgraph integration info to provider.yaml --- airflow/providers/microsoft/azure/provider.yaml | 15 +++++++++++++++ docs/integration-logos/azure/Microsoft-Graph.png | Bin 0 -> 3784 bytes 2 files changed, 15 insertions(+) create mode 100644 docs/integration-logos/azure/Microsoft-Graph.png diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index c9084789afe0f..c1518ad71880f 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -164,6 +164,10 @@ integrations: external-doc-url: https://azure.microsoft.com/en-us/products/storage/data-lake-storage/ logo: /integration-logos/azure/Data Lake Storage.svg tags: [azure] + - integration-name: Microsoft Graph API + external-doc-url: https://learn.microsoft.com/en-us/graph/use-the-api/ + logo: /integration-logos/azure/Microsoft-Graph.svg + tags: [azure] operators: - integration-name: Microsoft Azure Data Lake Storage @@ -193,6 +197,9 @@ operators: - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.operators.synapse + - integration-name: Microsoft Graph API + python-modules: + - airflow.providers.microsoft.azure.operators.msgraph sensors: - integration-name: Microsoft Azure Cosmos DB @@ -247,6 +254,9 @@ hooks: - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.hooks.synapse + - integration-name: Microsoft Graph API + python-modules: + - airflow.providers.microsoft.azure.hooks.msgraph triggers: - integration-name: Microsoft Azure Data Factory @@ -255,6 +265,9 @@ triggers: - integration-name: Microsoft Azure Blob Storage python-modules: - airflow.providers.microsoft.azure.triggers.wasb + - integration-name: Microsoft Graph API + python-modules: + - airflow.providers.microsoft.azure.triggers.msgraph transfers: - source-integration-name: Local @@ -308,6 +321,8 @@ connection-types: connection-type: adls - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook connection-type: azure_synapse_pipeline + - hook-class-name: airflow.providers.microsoft.azure.hooks.msgraph.KiotaRequestAdapterHook + connection-type: http secrets-backends: - airflow.providers.microsoft.azure.secrets.key_vault.AzureKeyVaultBackend diff --git a/docs/integration-logos/azure/Microsoft-Graph.png b/docs/integration-logos/azure/Microsoft-Graph.png new file mode 100644 index 0000000000000000000000000000000000000000..0724a1e09b495aa9831aad2f9daa37b86fc57c5c GIT binary patch literal 3784 zcmV;(4ma_MP)u0C?07GhP%#eO0yXBBkj`WrqMNSShOOQr!O^n&SX=)BtqT3VzoBUbIx- z|24Pp0Bg}PvhD(F!T_P^OVj*D%=rLt(mK8K0dmYCoZ=ji*$^*R@b>)K=l4Lx_5fzj zyV36zg3pPua-2F@ zif@6WEl-936=;Bwu1;f-03mFkuEY*xtN;LF02opL3r=x>meP@3&Hw-nib+I4RCt`# zoqJy!w-ScE_S)GVnKi`N#^#bW5CTn1xivs@PD41PNz>4i_WOUBN?Ko}U0p`fh+oe9 znS><|?`UT9N^2So&nVu!d;KhWyngql@ckcp#`N`QltqtZu-|HJ{rf*p2v(epMoIJ- zY>&I$7K>ktCnG!R5fYUu%SeBJtK@36y5#j!u^k<(uf;^mm|(31SgZBCKYIN%WM`8! zJ&nhek?s900gR#zveBj=NR&uclxlO-UeA&)YE?S!Q z_V8q*-Fdt9Tu0XG?(=^XxsuH8O3_edBZK{}9*m)}Z|v3arp(KcSSU`%S8=Y$mXAt z^C=MkMq|AAqG%1ExSNP%0jrIc`u(vz6Bn^H|FTiEG5K&at%N66rW`u#Xvrv+Y^S{n zS$^sv8|`80p#y-HWpT)^YNDk*S2Q+sfb+19Mk`4Svf`&2nY(QLEmLo{qa|xa1hn%> z9WCv_i)3aSIzVXIGkmhk0|gn~+`!NvLx(=P!jmFYl8ao4mW}{yeCT-Y3r2S}M0{3( zOmJ@0)6*M$2;jBJ9O3{;KBZDz2SF5fXnwK!GP5fHfcn`w{W2#k_ zGm@7ey9R36_U4s?IUU$st))34%7)@4T-v}=?f_x+Fc?~rU51lYT?2JW?m0LeC}^W+ zMKIbSfGkaq`yu16HOBSUiLe7$x6}|G!y)n8Wcs6pdkOX$lBPSuW|J2E?RdIvDF&WU zFQHIxKxMLKv$?^}QU>E#Y$~*5R0N@&+tZPeaWo3ok5nR+XklK$rM-r5vN;-UrE!E) zqJ>ow{RVUq!={MFvJz(>BS%+#O(j6KD;rrUnTSSbt<{B(*l3COC^@?@ue6;dnSjR2 z3AS(=wCtJhD5-Z1c;v&T5{(LnJ#uw5M|r+p!lhZX9f6FZN!wAj#l~dQqWO6VX4hb~ z%u2MH9XZ%XWRra#A^eV|4s%mJ4x2Js-9#XLyhfl!OA-&+X2yxXLZ(7vwGwNrtoT0S zH}UWi46Z?*HmjjRV{k|-PLJl{B}`2de~nCo#;nK)Mr7s3=hy&Du0b^!IhyWRJ0tsX z(VV@6OE)q-8k?EJo16(P8M(=(tsy%~G7}m`lVjHg^=M9`q_&1=$c$)g^a`9SYi^y) z=F$LkElmSP(O%k)qJr4WXh~9l$rN`D%w#q+o|OSBPGMWsQ@qGs7_{; z%n7YL4R(dIzIkQ$5)`fhfQ*LAJv5Ke%bJi@FQM)l=tiqb4Bf8SWdp|zf&B`1X4wH8 zR__|Xz?{+8eKaS30dhOtC|k`nfRZ_*@vMZOHR>!a>48Qsp)&E?D`Nl}4KjwBCtbbo zvkKD=U?$&Jav<|VgCFgH(6W(UHrX`*CnJjuAF{H8L*PJYh8@7vpUfAnJRty}nH&Ux zP1m=@`${g!$kDuCuwmDpu+Vg)7B>~eVYkqifo-KxP9b%r6oS2axv{HV3{GJ zv2ChaI5=A6X7co6;=jyFG_1GlvrM^@vzsHUCrDmG=^6x(SxkR7IMwLofF?Bn?izUX z={RVK>J{6?*Oz=%Qico1WJaF z#^J=#%Moqf0M0hU$?(y5Rtht-+sBo*NbMQD;nM0 z9HWpq9o0AoyH(Rqh+8BL95P>z70O2X*_vDIuKj?H&sy&QY|p(m(k43V{t7jSx{X zSG39b{CxDG(+N*XTr`rb9m`$Gg)Gkx7q-urUpk#o;ZV_7#h1y3PNE$x4dk2OUInuS zAJqt8$zE+nazV>4g!j?!DOzAsqMT=xWFy7!=Ylpl65Qk2ADk^jI3&}b){i~a;~mg? z7t#^iW2Y0yhJ;3&wDR#Ts$I#4P@i6vaNbFYf=0+o&9uG)7wud=glar|?t}@4fJVsf zP`3{Oqvh#gqoJ0~mkj}JKI4a>P2N|@8ncCLo=F*cz{V=RVe1H2O>yDp@$Hh$N4P*| zoa`DTb3ofXs*^PuE7>|H?!pD3(SgwTTWd1*f>f{{(m)ntv?*NzyDY`*p}Kr%4ef%dZh{acC6Eh+ubsAQn0gUHl1 zq?GMMLhGDFMIYdim!$MK&!9VL-CR`*i0d7?Xp{Ga5v}wRypqxd zjk!7wToAP0h4oYH?(1lzM=#iD`oh16=^rXKfNi~S10x76PeaLUI{+d6{4*Mv zV!Jl+oI}L;2u#UN`_|*Si7(pZt_VjndI>_@w%ae-ook!NjD?2qi<26K(Mu3SWATDY zhj6pristMZs6%y{s39xiZm+SGtD^Nzk$&O%BVe`Nw#B%;Zlcmv(KZ8yM=5Mxf+*Tj zyAw>f3YwQ|pbo2-P$gq0sp(*hZ-aibl#>||wL28>{0w??fYW$3&FfN0CO^&Fms zbPRJ1)KUB|M6{)LUW;yp_EQl7O?=0}myE%^@W{(mI^q`oWqekUxuoL~&bg#4$a?E> z$Qlod9e_JoyBxusq>EWYcog+}30`Czu6^P_x19KkNNA0lYD#*fqW}jU!RCX1o@2NM z>L~7i!4Yjv2W@r3MGYyUp;f(v0MV9k9Y8lI(E^T~AR&#C!DMCF-G(_D_V6f&^o|1# z8H(ppIn&Kt?mU4MTL- zKpkHqk+u1Jsg#b@kor5x_%ZhqB9OJq`LZsZ*CjKb;G<*=(dKaV1D_o{yVw% Date: Thu, 14 Mar 2024 14:51:51 +0100 Subject: [PATCH 017/105] refactor: Added init in resources --- .../microsoft/azure/resources/__init__.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 tests/providers/microsoft/azure/resources/__init__.py diff --git a/tests/providers/microsoft/azure/resources/__init__.py b/tests/providers/microsoft/azure/resources/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/providers/microsoft/azure/resources/__init__.py @@ -0,0 +1,16 @@ +# 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 4cef94eb779fa2ee2588fa296842d297a5689a1a Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 15:18:58 +0100 Subject: [PATCH 018/105] fix: Fixed typing of response_handler --- .../microsoft/azure/serialization/response_handler.py | 5 ++--- .../providers/microsoft/azure/serialization/serializer.py | 6 +++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 037339459fe0f..1bcf03650f63d 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -20,7 +20,6 @@ from typing import Any, Callable from kiota_abstractions.response_handler import NativeResponseType, ResponseHandler - # type: ignore[TCH002] from kiota_abstractions.serialization import ParsableFactory @@ -32,11 +31,11 @@ class CallableResponseHandler(ResponseHandler): def __init__( self, - callable_function: Callable[[NativeResponseType, dict[str, (ParsableFactory, None)]], Any], + callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any], ): self.callable_function = callable_function async def handle_response_async( - self, response: NativeResponseType, error_map: dict[str, (ParsableFactory, None)] = None + self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None ) -> Any: return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 3e0b84f07590d..2a4c5af64bcfb 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -34,11 +34,11 @@ class ResponseSerializer: ResponseSerializer serializes the response as a string. """ - def __init__(self, encoding: (str, None) = None): + def __init__(self, encoding: str | None = None): self.encoding = encoding or locale.getpreferredencoding() - def serialize(self, response) -> (str, None): - def convert(value) -> (str, None): + def serialize(self, response) -> str | None: + def convert(value) -> str | None: if value is not None: if isinstance(value, UUID): return str(value) From 88aa7dccd95b98585872ae9eb5cd339162a06bb9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 15:58:58 +0100 Subject: [PATCH 019/105] refactor: Added assertions on conn_id, tenant_id, client_id and client_secret --- .../microsoft/azure/hooks/msgraph.py | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index f314d1091da29..252539b0dc804 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -18,7 +18,7 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from urllib.parse import urljoin import httpx @@ -106,12 +106,19 @@ def to_httpx_proxies(proxies: dict) -> dict: proxies["https://"] = proxies.pop("https") return proxies - def get_conn(self) -> RequestAdapter: - if not self.conn_id: + @staticmethod + def assert_not_none(value: Any, message: str) -> None: + if not value: raise AirflowException( "Failed to create the KiotaRequestAdapterHook. No conn_id provided!" ) + def get_conn(self) -> RequestAdapter: + self.assert_not_none( + self.conn_id, + "Failed to create the KiotaRequestAdapterHook. No conn_id provided!", + ) + api_version, request_adapter = self.cached_request_adapters.get( self.conn_id, (None, None) ) @@ -146,6 +153,20 @@ def get_conn(self) -> RequestAdapter: self.log.info("Timeout: %s", self.timeout) self.log.info("Trust env: %s", trust_env) self.log.info("Proxies: %s", json.dumps(proxies)) + + self.assert_not_none( + tenant_id, + "Failed to create the KiotaRequestAdapterHook. No tenant_id provided!", + ) + self.assert_not_none( + connection.login, + "Failed to create the KiotaRequestAdapterHook. No client_id provided!", + ) + self.assert_not_none( + connection.password, + "Failed to create the KiotaRequestAdapterHook. No client_secret provided!", + ) + credentials = identity.ClientSecretCredential( tenant_id=tenant_id, client_id=connection.login, From c528f50760a4c9bbde3c5cc82bd1e0c3f03ada17 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 18:46:20 +0100 Subject: [PATCH 020/105] refactor: Fixed some static checks --- .../providers/microsoft/azure/operators/msgraph.py | 12 ++++-------- .../providers/microsoft/azure/triggers/msgraph.py | 12 ++++-------- 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index b0dd30c72b2e4..33e75575f795d 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -21,7 +21,7 @@ Any, TYPE_CHECKING, Sequence, - Callable, ) + Callable, List, ) from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import BaseOperator @@ -75,7 +75,7 @@ class MSGraphAsyncOperator(BaseOperator): def __init__( self, *, - url: str | None = None, + url: str, response_type: ResponseType | None = None, response_handler: Callable[ [NativeResponseType, dict[str, ParsableFactory | None] | None], Any @@ -114,7 +114,7 @@ def __init__( self.api_version = api_version self.result_processor = result_processor self.serializer: ResponseSerializer = serializer() - self.results = None + self.results: List[Any] | None = None def execute(self, context: Context) -> None: self.log.info("Executing url '%s' as '%s'", self.url, self.method) @@ -162,11 +162,9 @@ def execute_complete( self.log.info("response: %s", response) if response: - self.log.debug("response type: %s", type(response)) - response = self.serializer.deserialize(response) - self.log.debug("deserialized response type: %s", type(response)) + self.log.debug("deserialize response: %s", response) result = self.result_processor(context, response) @@ -174,8 +172,6 @@ def execute_complete( event["response"] = result - self.log.debug("parsed response type: %s", type(response)) - try: self.trigger_next_link( response, method_name="pull_execute_complete" diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 037f9c103f210..6146e3473f5b6 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -83,7 +83,7 @@ class MSGraphTrigger(BaseTrigger): def __init__( self, - url: str | None = None, + url: str, response_type: ResponseType | None = None, response_handler: Callable[ [NativeResponseType, dict[str, ParsableFactory | None] | None], Any @@ -180,17 +180,13 @@ async def run(self) -> AsyncIterator[TriggerEvent]: if response: response_type = type(response) - self.log.debug("response type: %s", type(response)) - - response = self.serializer.serialize(response) - - self.log.debug("serialized response type: %s", type(response)) + self.log.debug("response type: %s", response_type) yield TriggerEvent( { "status": "success", "type": f"{response_type.__module__}.{response_type.__name__}", - "response": response, + "response": self.serializer.serialize(response), } ) else: @@ -205,7 +201,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: self.log.exception("An error occurred: %s", e) yield TriggerEvent({"status": "failure", "message": str(e)}) - def normalize_url(self) -> str: + def normalize_url(self) -> str | None: if self.url.startswith("/"): return self.url.replace("/", "", 1) return self.url From 258abd6e4afb246df259ccaabb1c2a086ab52b51 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 18:50:01 +0100 Subject: [PATCH 021/105] Revert "refactor: Added assertions on conn_id, tenant_id, client_id and client_secret" This reverts commit 88aa7dccd95b98585872ae9eb5cd339162a06bb9. --- .../microsoft/azure/hooks/msgraph.py | 27 +++---------------- 1 file changed, 3 insertions(+), 24 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index 252539b0dc804..f314d1091da29 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -18,7 +18,7 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from urllib.parse import urljoin import httpx @@ -106,19 +106,12 @@ def to_httpx_proxies(proxies: dict) -> dict: proxies["https://"] = proxies.pop("https") return proxies - @staticmethod - def assert_not_none(value: Any, message: str) -> None: - if not value: + def get_conn(self) -> RequestAdapter: + if not self.conn_id: raise AirflowException( "Failed to create the KiotaRequestAdapterHook. No conn_id provided!" ) - def get_conn(self) -> RequestAdapter: - self.assert_not_none( - self.conn_id, - "Failed to create the KiotaRequestAdapterHook. No conn_id provided!", - ) - api_version, request_adapter = self.cached_request_adapters.get( self.conn_id, (None, None) ) @@ -153,20 +146,6 @@ def get_conn(self) -> RequestAdapter: self.log.info("Timeout: %s", self.timeout) self.log.info("Trust env: %s", trust_env) self.log.info("Proxies: %s", json.dumps(proxies)) - - self.assert_not_none( - tenant_id, - "Failed to create the KiotaRequestAdapterHook. No tenant_id provided!", - ) - self.assert_not_none( - connection.login, - "Failed to create the KiotaRequestAdapterHook. No client_id provided!", - ) - self.assert_not_none( - connection.password, - "Failed to create the KiotaRequestAdapterHook. No client_secret provided!", - ) - credentials = identity.ClientSecretCredential( tenant_id=tenant_id, client_id=connection.login, From 2304baff80e83948287967eddd2f714e49376835 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 18:51:52 +0100 Subject: [PATCH 022/105] refactor: Changed imports in hook as we don't use mockito anymore we don't need the module before constructor --- .../microsoft/azure/hooks/msgraph.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index f314d1091da29..ef31afe032855 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -22,10 +22,11 @@ from urllib.parse import urljoin import httpx -from azure import identity +from azure.identity import ClientSecretCredential from httpx import Timeout -from kiota_authentication_azure import azure_identity_authentication_provider -from kiota_http import httpx_request_adapter +from kiota_authentication_azure.azure_identity_authentication_provider import \ + AzureIdentityAuthenticationProvider +from kiota_http.httpx_request_adapter import HttpxRequestAdapter from msgraph_core import GraphClientFactory from msgraph_core._enums import APIVersion, NationalClouds @@ -146,10 +147,10 @@ def get_conn(self) -> RequestAdapter: self.log.info("Timeout: %s", self.timeout) self.log.info("Trust env: %s", trust_env) self.log.info("Proxies: %s", json.dumps(proxies)) - credentials = identity.ClientSecretCredential( - tenant_id=tenant_id, - client_id=connection.login, - client_secret=connection.password, + credentials = ClientSecretCredential( + tenant_id=tenant_id, # type: ignore + client_id=connection.login, # type: ignore + client_secret=connection.password, # type: ignore proxies=proxies, ) http_client = GraphClientFactory.create_with_default_middleware( @@ -162,10 +163,10 @@ def get_conn(self) -> RequestAdapter: ), host=host, ) - auth_provider = azure_identity_authentication_provider.AzureIdentityAuthenticationProvider( + auth_provider = AzureIdentityAuthenticationProvider( credentials=credentials, scopes=scopes ) - request_adapter = httpx_request_adapter.HttpxRequestAdapter( + request_adapter = HttpxRequestAdapter( authentication_provider=auth_provider, http_client=http_client, base_url=base_url, From 22906cff83f1095aac2e5ec9cff443dd3b6a8477 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 18:55:59 +0100 Subject: [PATCH 023/105] refactor: Renamed test methods --- tests/providers/microsoft/azure/operators/test_msgraph.py | 8 ++++---- tests/providers/microsoft/azure/triggers/test_msgraph.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 54936b80740b1..1520fe603e7fb 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -30,7 +30,7 @@ class TestMSGraphAsyncOperator(Base): - def test_run_when_expression_is_valid(self): + def test_execute(self): users = load_json("resources", "users.json") next_users = load_json("resources", "next_users.json") response = mock_json_response(200, users, next_users) @@ -57,7 +57,7 @@ def test_run_when_expression_is_valid(self): assert events[1].payload["type"] == "builtins.dict" assert events[1].payload["response"] == json.dumps(next_users) - def test_run_when_expression_is_valid_and_do_xcom_push_is_false(self): + def test_execute_when_do_xcom_push_is_false(self): users = load_json("resources", "users.json") users.pop("@odata.nextLink") response = mock_json_response(200, users) @@ -79,7 +79,7 @@ def test_run_when_expression_is_valid_and_do_xcom_push_is_false(self): assert events[0].payload["type"] == "builtins.dict" assert events[0].payload["response"] == json.dumps(users) - def test_run_when_an_exception_occurs(self): + def test_execute_when_an_exception_occurs(self): with self.patch_hook_and_request_adapter(AirflowException()): operator = MSGraphAsyncOperator( task_id="users_delta", @@ -91,7 +91,7 @@ def test_run_when_an_exception_occurs(self): with pytest.raises(AirflowException): self.execute_operator(operator) - def test_run_when_url_which_returns_bytes(self): + def test_execute_when_response_is_bytes(self): content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) drive_id = "82f9d24d-6891-4790-8b6d-f1b2a1d0ca22" diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index ab95ddded1cf6..e55771c11f8eb 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -71,7 +71,7 @@ def test_run_when_response_cannot_be_converted_to_json(self): assert actual.payload["status"] == "failure" assert actual.payload["message"] == "" - def test_run_when_url_with_response_type_bytes(self): + def test_run_when_response_is_bytes(self): content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) response = mock_response(200, content) From 8078c620f925d0b90b63d1658354787a616560ea Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 21:17:37 +0100 Subject: [PATCH 024/105] refactor: Replace List type with list --- airflow/providers/microsoft/azure/operators/msgraph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 33e75575f795d..00986ac2135d0 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -21,7 +21,7 @@ Any, TYPE_CHECKING, Sequence, - Callable, List, ) + Callable, ) from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import BaseOperator @@ -114,7 +114,7 @@ def __init__( self.api_version = api_version self.result_processor = result_processor self.serializer: ResponseSerializer = serializer() - self.results: List[Any] | None = None + self.results: list[Any] | None = None def execute(self, context: Context) -> None: self.log.info("Executing url '%s' as '%s'", self.url, self.method) From 967216e1107df2fd2907c976add7714343ca99fa Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 21:17:53 +0100 Subject: [PATCH 025/105] refactor: Moved docstring as one line --- .../microsoft/azure/serialization/response_handler.py | 4 +--- airflow/providers/microsoft/azure/serialization/serializer.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 1bcf03650f63d..abe0da8cae572 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -25,9 +25,7 @@ class CallableResponseHandler(ResponseHandler): - """ - CallableResponseHandler executes the passed callable_function with response as parameter. - """ + """CallableResponseHandler executes the passed callable_function with response as parameter.""" def __init__( self, diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py index 2a4c5af64bcfb..419f0ce005bf2 100644 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ b/airflow/providers/microsoft/azure/serialization/serializer.py @@ -30,9 +30,7 @@ class ResponseSerializer: - """ - ResponseSerializer serializes the response as a string. - """ + """ResponseSerializer serializes the response as a string.""" def __init__(self, encoding: str | None = None): self.encoding = encoding or locale.getpreferredencoding() From c8c46715bd0f0502654d5c9ff540b1dcfe5da90e Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 21:27:35 +0100 Subject: [PATCH 026/105] refactor: Fixed typing for tests and added test for response_handler --- .../azure/serialization/response_handler.py | 7 ++-- tests/providers/microsoft/azure/base.py | 10 +++--- .../azure/serializer/test_response_handler.py | 36 +++++++++++++++++++ 3 files changed, 45 insertions(+), 8 deletions(-) create mode 100644 tests/providers/microsoft/azure/serializer/test_response_handler.py diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index abe0da8cae572..d50a6893ba8e0 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -17,11 +17,12 @@ # under the License. from __future__ import annotations -from typing import Any, Callable +from typing import Any, Callable, TYPE_CHECKING from kiota_abstractions.response_handler import NativeResponseType, ResponseHandler -# type: ignore[TCH002] -from kiota_abstractions.serialization import ParsableFactory + +if TYPE_CHECKING: + from kiota_abstractions.serialization import ParsableFactory class CallableResponseHandler(ResponseHandler): diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 335bda425363e..2b3d275ea4f1b 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -20,24 +20,24 @@ from contextlib import contextmanager from copy import deepcopy from datetime import datetime -from typing import Any, Iterable +from typing import Any, Iterable, TYPE_CHECKING from unittest.mock import patch import pytest from kiota_http.httpx_request_adapter import HttpxRequestAdapter -# type: ignore[TCH002] -from sqlalchemy.orm import Session from airflow.exceptions import TaskDeferred from airflow.models import Operator, TaskInstance from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook -# type: ignore[TCH001] -from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.session import NEW_SESSION from airflow.utils.state import TaskInstanceState from airflow.utils.xcom import XCOM_RETURN_KEY from tests.providers.microsoft.conftest import get_airflow_connection +if TYPE_CHECKING: + from sqlalchemy.orm import Session + from airflow.triggers.base import BaseTrigger, TriggerEvent + class MockedTaskInstance(TaskInstance): values = {} diff --git a/tests/providers/microsoft/azure/serializer/test_response_handler.py b/tests/providers/microsoft/azure/serializer/test_response_handler.py new file mode 100644 index 0000000000000..77ec4ffd52b77 --- /dev/null +++ b/tests/providers/microsoft/azure/serializer/test_response_handler.py @@ -0,0 +1,36 @@ +# 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 airflow.providers.microsoft.azure.serialization.response_handler import CallableResponseHandler +from tests.providers.microsoft.azure.base import Base +from tests.providers.microsoft.conftest import load_json, mock_json_response + + +class TestResponseHandler(Base): + def test_handle_response_async(self): + users = load_json("resources", "users.json") + response = mock_json_response(200, users) + + actual = self._loop.run_until_complete( + CallableResponseHandler( + lambda response, error_map: response.json() + ).handle_response_async(response, None) + ) + + assert isinstance(actual, dict) + assert actual == users From 1d6ceea8bed63db362a9697a0e3247066e0362d4 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 14 Mar 2024 21:40:17 +0100 Subject: [PATCH 027/105] refactor: Refactored tests --- tests/providers/microsoft/azure/base.py | 10 ++++++++-- .../azure/serializer/test_response_handler.py | 2 +- .../providers/microsoft/azure/triggers/test_msgraph.py | 8 ++++---- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 2b3d275ea4f1b..2b6752159a4a0 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -86,12 +86,18 @@ def patch_hook_and_request_adapter(self, response): yield @staticmethod - async def run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]: + async def _run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]: events = [] async for event in trigger.run(): events.append(event) return events + def run_trigger(self, trigger: BaseTrigger) -> list[TriggerEvent]: + return self.run_async(self._run_tigger(trigger)) + + def run_async(self, future: Any) -> Any: + return self._loop.run_until_complete(future) + def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) context = {"ti": task_instance} @@ -104,7 +110,7 @@ def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task = deferred.value while task: - events = self._loop.run_until_complete(self.run_tigger(deferred.value.trigger)) + events = self.run_trigger(deferred.value.trigger) if not events: break diff --git a/tests/providers/microsoft/azure/serializer/test_response_handler.py b/tests/providers/microsoft/azure/serializer/test_response_handler.py index 77ec4ffd52b77..6b7a6447ad506 100644 --- a/tests/providers/microsoft/azure/serializer/test_response_handler.py +++ b/tests/providers/microsoft/azure/serializer/test_response_handler.py @@ -26,7 +26,7 @@ def test_handle_response_async(self): users = load_json("resources", "users.json") response = mock_json_response(200, users) - actual = self._loop.run_until_complete( + actual = self.run_async( CallableResponseHandler( lambda response, error_map: response.json() ).handle_response_async(response, None) diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index e55771c11f8eb..80b8e4ae30753 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -41,7 +41,7 @@ def test_run_when_valid_response(self): with self.patch_hook_and_request_adapter(response): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") - actual = self._loop.run_until_complete(self.run_tigger(trigger)) + actual = self.run_trigger(trigger) assert len(actual) == 1 assert isinstance(actual[0], TriggerEvent) @@ -54,7 +54,7 @@ def test_run_when_response_is_none(self): with self.patch_hook_and_request_adapter(response): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") - actual = self._loop.run_until_complete(self.run_tigger(trigger)) + actual = self.run_trigger(trigger) assert len(actual) == 1 assert isinstance(actual[0], TriggerEvent) @@ -65,7 +65,7 @@ def test_run_when_response_is_none(self): def test_run_when_response_cannot_be_converted_to_json(self): with self.patch_hook_and_request_adapter(AirflowException()): trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api") - actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) + actual = next(iter(self.run_trigger(trigger))) assert isinstance(actual, TriggerEvent) assert actual.payload["status"] == "failure" @@ -81,7 +81,7 @@ def test_run_when_response_is_bytes(self): "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content" ) trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api") - actual = next(iter(self._loop.run_until_complete(self.run_tigger(trigger)))) + actual = next(iter(self.run_trigger(trigger))) assert isinstance(actual, TriggerEvent) assert actual.payload["status"] == "success" From 2d1cc517157ef63adec1f9d69854d875bc9143a8 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 07:47:42 +0100 Subject: [PATCH 028/105] fix: Fixed MS Graph logo filename --- airflow/providers/microsoft/azure/provider.yaml | 2 +- ...{Microsoft-Graph.png => Microsoft-Graph-API.png} | Bin 2 files changed, 1 insertion(+), 1 deletion(-) rename docs/integration-logos/azure/{Microsoft-Graph.png => Microsoft-Graph-API.png} (100%) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index c1518ad71880f..4a53ec80ee9ad 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -166,7 +166,7 @@ integrations: tags: [azure] - integration-name: Microsoft Graph API external-doc-url: https://learn.microsoft.com/en-us/graph/use-the-api/ - logo: /integration-logos/azure/Microsoft-Graph.svg + logo: /integration-logos/azure/Microsoft-Graph-API.png tags: [azure] operators: diff --git a/docs/integration-logos/azure/Microsoft-Graph.png b/docs/integration-logos/azure/Microsoft-Graph-API.png similarity index 100% rename from docs/integration-logos/azure/Microsoft-Graph.png rename to docs/integration-logos/azure/Microsoft-Graph-API.png From 740c96ebe37daf2d4e4336d945c912e30e3ffc75 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 09:41:28 +0100 Subject: [PATCH 029/105] refactor: Fixed additional static checks remarks --- .../microsoft/azure/hooks/msgraph.py | 20 +++++------ .../microsoft/azure/operators/msgraph.py | 33 ++++++++----------- .../azure/serialization/response_handler.py | 2 +- .../microsoft/azure/triggers/msgraph.py | 25 +++++--------- tests/providers/microsoft/azure/base.py | 4 +-- .../microsoft/azure/hooks/test_msgraph.py | 16 ++++----- .../azure/serializer/test_response_handler.py | 6 ++-- .../microsoft/azure/triggers/test_msgraph.py | 2 +- 8 files changed, 44 insertions(+), 64 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index ef31afe032855..5ed148813bb65 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -24,8 +24,9 @@ import httpx from azure.identity import ClientSecretCredential from httpx import Timeout -from kiota_authentication_azure.azure_identity_authentication_provider import \ - AzureIdentityAuthenticationProvider +from kiota_authentication_azure.azure_identity_authentication_provider import ( + AzureIdentityAuthenticationProvider, +) from kiota_http.httpx_request_adapter import HttpxRequestAdapter from msgraph_core import GraphClientFactory from msgraph_core._enums import APIVersion, NationalClouds @@ -34,9 +35,10 @@ from airflow.hooks.base import BaseHook if TYPE_CHECKING: - from airflow.models import Connection from kiota_abstractions.request_adapter import RequestAdapter + from airflow.models import Connection + class KiotaRequestAdapterHook(BaseHook): """ @@ -109,13 +111,9 @@ def to_httpx_proxies(proxies: dict) -> dict: def get_conn(self) -> RequestAdapter: if not self.conn_id: - raise AirflowException( - "Failed to create the KiotaRequestAdapterHook. No conn_id provided!" - ) + raise AirflowException("Failed to create the KiotaRequestAdapterHook. No conn_id provided!") - api_version, request_adapter = self.cached_request_adapters.get( - self.conn_id, (None, None) - ) + api_version, request_adapter = self.cached_request_adapters.get(self.conn_id, (None, None)) if not request_adapter: connection = self.get_connection(conn_id=self.conn_id) @@ -163,9 +161,7 @@ def get_conn(self) -> RequestAdapter: ), host=host, ) - auth_provider = AzureIdentityAuthenticationProvider( - credentials=credentials, scopes=scopes - ) + auth_provider = AzureIdentityAuthenticationProvider(credentials=credentials, scopes=scopes) request_adapter = HttpxRequestAdapter( authentication_provider=auth_provider, http_client=http_client, diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 00986ac2135d0..08ba00ea64c22 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -18,10 +18,11 @@ from __future__ import annotations from typing import ( - Any, TYPE_CHECKING, + Any, + Callable, Sequence, - Callable, ) +) from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import BaseOperator @@ -33,13 +34,15 @@ from airflow.utils.xcom import XCOM_RETURN_KEY if TYPE_CHECKING: - from msgraph_core import APIVersion from io import BytesIO - from airflow.utils.context import Context + from kiota_abstractions.request_adapter import ResponseType from kiota_abstractions.request_information import QueryParams from kiota_abstractions.response_handler import NativeResponseType from kiota_abstractions.serialization import ParsableFactory + from msgraph_core import APIVersion + + from airflow.utils.context import Context class MSGraphAsyncOperator(BaseOperator): @@ -91,9 +94,7 @@ def __init__( timeout: float | None = None, proxies: dict | None = None, api_version: APIVersion | None = None, - result_processor: Callable[ - [Context, Any], Any - ] = lambda context, result: result, + result_processor: Callable[[Context, Any], Any] = lambda context, result: result, serializer: type[ResponseSerializer] = ResponseSerializer, **kwargs: Any, ): @@ -150,9 +151,7 @@ def execute_complete( self.log.debug("context: %s", context) if event: - self.log.info( - "%s completed with %s: %s", self.task_id, event.get("status"), event - ) + self.log.info("%s completed with %s: %s", self.task_id, event.get("status"), event) if event.get("status") == "failure": raise AirflowException(event.get("message")) @@ -173,9 +172,7 @@ def execute_complete( event["response"] = result try: - self.trigger_next_link( - response, method_name="pull_execute_complete" - ) + self.trigger_next_link(response, method_name="pull_execute_complete") except TaskDeferred as exception: self.append_result( result=result, @@ -217,9 +214,7 @@ def push_xcom(self, context: Context, value) -> None: 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: + def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = None) -> Any: self.results = list( self.xcom_pull( context=context, @@ -227,7 +222,7 @@ def pull_execute_complete( dag_id=self.dag_id, key=self.key, ) - or [] # noqa: W503 + or [] ) self.log.info( "Pulled XCom with task_id '%s' and dag_id '%s' and key '%s': %s", @@ -238,9 +233,7 @@ def pull_execute_complete( ) return self.execute_complete(context, event) - def trigger_next_link( - self, response, method_name="execute_complete" - ) -> None: + def trigger_next_link(self, response, method_name="execute_complete") -> None: if isinstance(response, dict): odata_next_link = response.get("@odata.nextLink") diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index d50a6893ba8e0..4503babe61d2d 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -17,7 +17,7 @@ # under the License. from __future__ import annotations -from typing import Any, Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable from kiota_abstractions.response_handler import NativeResponseType, ResponseHandler diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 6146e3473f5b6..b4e2bda1fd91a 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -18,11 +18,11 @@ from __future__ import annotations from typing import ( + TYPE_CHECKING, Any, AsyncIterator, - Sequence, - TYPE_CHECKING, Callable, + Sequence, ) from kiota_abstractions.api_error import APIError @@ -42,6 +42,7 @@ if TYPE_CHECKING: from io import BytesIO + from kiota_abstractions.request_adapter import RequestAdapter from kiota_abstractions.request_information import QueryParams from kiota_abstractions.response_handler import NativeResponseType @@ -116,9 +117,7 @@ def __init__( self.query_parameters = query_parameters self.headers = headers self.content = content - self.serializer: ResponseSerializer = self.resolve_type( - serializer, default=ResponseSerializer - )() + self.serializer: ResponseSerializer = self.resolve_type(serializer, default=ResponseSerializer)() @classmethod def resolve_type(cls, value: str | type, default) -> type: @@ -171,7 +170,7 @@ 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.""" + """Make a series of asynchronous HTTP calls via a KiotaRequestAdapterHook.""" try: response = await self.execute() @@ -216,21 +215,13 @@ def request_information(self) -> RequestInformation: request_information.http_method = Method(self.method.strip().upper()) request_information.query_parameters = self.query_parameters or {} if not self.response_type: - request_information.request_options[ - ResponseHandlerOption.get_key() - ] = ResponseHandlerOption( + request_information.request_options[ResponseHandlerOption.get_key()] = ResponseHandlerOption( response_handler=CallableResponseHandler(self.response_handler) ) request_information.content = self.content - headers = ( - {**self.DEFAULT_HEADERS, **self.headers} - if self.headers - else self.DEFAULT_HEADERS - ) + headers = {**self.DEFAULT_HEADERS, **self.headers} if self.headers else self.DEFAULT_HEADERS for header_name, header_value in headers.items(): - request_information.headers.try_add( - header_name=header_name, header_value=header_value - ) + request_information.headers.try_add(header_name=header_name, header_value=header_value) return request_information @staticmethod diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 2b6752159a4a0..537899251055d 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -20,7 +20,7 @@ from contextlib import contextmanager from copy import deepcopy from datetime import datetime -from typing import Any, Iterable, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Iterable from unittest.mock import patch import pytest @@ -77,7 +77,7 @@ def teardown_method(self, method): @contextmanager def patch_hook_and_request_adapter(self, response): with patch( - "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection + "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection ), patch.object(HttpxRequestAdapter, "get_http_response_message") as mock_get_http_response: if isinstance(response, Exception): mock_get_http_response.side_effect = response diff --git a/tests/providers/microsoft/azure/hooks/test_msgraph.py b/tests/providers/microsoft/azure/hooks/test_msgraph.py index c52db83f8cc35..1c1046e1fa4f3 100644 --- a/tests/providers/microsoft/azure/hooks/test_msgraph.py +++ b/tests/providers/microsoft/azure/hooks/test_msgraph.py @@ -27,10 +27,10 @@ class TestKiotaRequestAdapterHook: def test_get_conn(self): - with (patch( + with patch( "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection, - )): + ): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_conn() @@ -38,29 +38,29 @@ def test_get_conn(self): assert actual.base_url == "https://graph.microsoft.com/v1.0" def test_api_version(self): - with (patch( + with patch( "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection, - )): + ): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") assert hook.api_version == APIVersion.v1 def test_get_api_version_when_empty_config_dict(self): - with (patch( + with patch( "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection, - )): + ): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_api_version({}) assert actual == APIVersion.v1 def test_get_api_version_when_api_version_in_config_dict(self): - with (patch( + with patch( "airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection, - )): + ): hook = KiotaRequestAdapterHook(conn_id="msgraph_api") actual = hook.get_api_version({"api_version": "beta"}) diff --git a/tests/providers/microsoft/azure/serializer/test_response_handler.py b/tests/providers/microsoft/azure/serializer/test_response_handler.py index 6b7a6447ad506..f5575df5fdcbb 100644 --- a/tests/providers/microsoft/azure/serializer/test_response_handler.py +++ b/tests/providers/microsoft/azure/serializer/test_response_handler.py @@ -27,9 +27,9 @@ def test_handle_response_async(self): response = mock_json_response(200, users) actual = self.run_async( - CallableResponseHandler( - lambda response, error_map: response.json() - ).handle_response_async(response, None) + CallableResponseHandler(lambda response, error_map: response.json()).handle_response_async( + response, None + ) ) assert isinstance(actual, dict) diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index 80b8e4ae30753..4b943b4170944 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -114,5 +114,5 @@ def test_serialize(self): "timeout": None, "proxies": None, "api_version": "v1.0", - "serializer": "airflow.providers.microsoft.azure.serialization.serializer.ResponseSerializer" + "serializer": "airflow.providers.microsoft.azure.serialization.serializer.ResponseSerializer", } From 77dfaeef17de714654650c272dfd4b9918c2520b Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 10:07:39 +0100 Subject: [PATCH 030/105] refactor: Added white line in type checking block --- tests/providers/microsoft/azure/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 537899251055d..3cd980383a1a0 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from sqlalchemy.orm import Session + from airflow.triggers.base import BaseTrigger, TriggerEvent From 6ffe8a93d330c371bb45051bb263bf9c4a1b3f08 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 10:10:04 +0100 Subject: [PATCH 031/105] refactor: Added msgraph-core dependency to provider_dependencies.json --- generated/provider_dependencies.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index ef5039db27512..98da2aacd6453 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -701,7 +701,8 @@ "azure-storage-file-datalake>=12.9.1", "azure-storage-file-share", "azure-synapse-artifacts>=0.17.0", - "azure-synapse-spark" + "azure-synapse-spark", + "msgraph-core>=1.0.0" ], "devel-deps": [ "pywinrm" From 71fb85acf92dbdca3fd2542bd7699ec1a8d909ea Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 11:45:30 +0100 Subject: [PATCH 032/105] refactor: Updated docstring on response handler --- .../microsoft/azure/serialization/response_handler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py index 4503babe61d2d..58d18b9357ab8 100644 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ b/airflow/providers/microsoft/azure/serialization/response_handler.py @@ -26,7 +26,11 @@ class CallableResponseHandler(ResponseHandler): - """CallableResponseHandler executes the passed callable_function with response as parameter.""" + """ + CallableResponseHandler executes the passed callable_function with response as parameter. + + :param callable_function: Function which allows you to handle the response before returning. + """ def __init__( self, From 5dd317cdc9ff6c3bd088e5d8480f53a1c9cfbafa Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 12:53:09 +0100 Subject: [PATCH 033/105] refactor: Moved ResponseHandler and Serializer to triggers module --- .../microsoft/azure/operators/msgraph.py | 2 +- .../microsoft/azure/serialization/__init__.py | 16 ----- .../azure/serialization/response_handler.py | 44 ------------ .../azure/serialization/serializer.py | 62 ---------------- .../microsoft/azure/triggers/msgraph.py | 69 ++++++++++++++++-- .../microsoft/azure/serializer/__init__.py | 16 ----- .../azure/serializer/test_response_handler.py | 36 ---------- .../azure/serializer/test_serializer.py | 71 ------------------- .../microsoft/azure/triggers/test_msgraph.py | 70 +++++++++++++++++- 9 files changed, 130 insertions(+), 256 deletions(-) delete mode 100644 airflow/providers/microsoft/azure/serialization/__init__.py delete mode 100644 airflow/providers/microsoft/azure/serialization/response_handler.py delete mode 100644 airflow/providers/microsoft/azure/serialization/serializer.py delete mode 100644 tests/providers/microsoft/azure/serializer/__init__.py delete mode 100644 tests/providers/microsoft/azure/serializer/test_response_handler.py delete mode 100644 tests/providers/microsoft/azure/serializer/test_serializer.py diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 08ba00ea64c22..49f3d9b062829 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -27,7 +27,7 @@ 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.serialization.serializer import ( +from airflow.providers.microsoft.azure.triggers.msgraph import ( ResponseSerializer, ) from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger diff --git a/airflow/providers/microsoft/azure/serialization/__init__.py b/airflow/providers/microsoft/azure/serialization/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/airflow/providers/microsoft/azure/serialization/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/airflow/providers/microsoft/azure/serialization/response_handler.py b/airflow/providers/microsoft/azure/serialization/response_handler.py deleted file mode 100644 index 58d18b9357ab8..0000000000000 --- a/airflow/providers/microsoft/azure/serialization/response_handler.py +++ /dev/null @@ -1,44 +0,0 @@ -# -# 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 - -from kiota_abstractions.response_handler import NativeResponseType, ResponseHandler - -if TYPE_CHECKING: - from kiota_abstractions.serialization import ParsableFactory - - -class CallableResponseHandler(ResponseHandler): - """ - CallableResponseHandler executes the passed callable_function with response as parameter. - - :param callable_function: Function which allows you to handle the response before returning. - """ - - def __init__( - self, - callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any], - ): - self.callable_function = callable_function - - async def handle_response_async( - self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None - ) -> Any: - return self.callable_function(response, error_map) diff --git a/airflow/providers/microsoft/azure/serialization/serializer.py b/airflow/providers/microsoft/azure/serialization/serializer.py deleted file mode 100644 index 419f0ce005bf2..0000000000000 --- a/airflow/providers/microsoft/azure/serialization/serializer.py +++ /dev/null @@ -1,62 +0,0 @@ -# -# 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 Any -from uuid import UUID - -import pendulum - - -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 diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index b4e2bda1fd91a..130ffa71567ea 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -17,26 +17,29 @@ # 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 Any from typing import ( TYPE_CHECKING, - Any, AsyncIterator, Callable, Sequence, ) +from uuid import UUID +import pendulum from kiota_abstractions.api_error import APIError from kiota_abstractions.method import Method from kiota_abstractions.request_information import RequestInformation +from kiota_abstractions.response_handler import ResponseHandler from kiota_http.middleware.options import ResponseHandlerOption from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook -from airflow.providers.microsoft.azure.serialization.response_handler import ( - CallableResponseHandler, -) -from airflow.providers.microsoft.azure.serialization.serializer import ( - ResponseSerializer, -) from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.module_loading import import_string @@ -51,6 +54,58 @@ 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 CallableResponseHandler(ResponseHandler): + """ + CallableResponseHandler executes the passed callable_function with response as parameter. + + :param callable_function: Function which allows you to handle the response before returning. + """ + + def __init__( + self, + callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any], + ): + self.callable_function = callable_function + + async def handle_response_async( + self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None + ) -> Any: + return self.callable_function(response, error_map) + + class MSGraphTrigger(BaseTrigger): """ A Microsoft Graph API trigger which allows you to execute an async REST call to the Microsoft Graph API. diff --git a/tests/providers/microsoft/azure/serializer/__init__.py b/tests/providers/microsoft/azure/serializer/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/tests/providers/microsoft/azure/serializer/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/tests/providers/microsoft/azure/serializer/test_response_handler.py b/tests/providers/microsoft/azure/serializer/test_response_handler.py deleted file mode 100644 index f5575df5fdcbb..0000000000000 --- a/tests/providers/microsoft/azure/serializer/test_response_handler.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 airflow.providers.microsoft.azure.serialization.response_handler import CallableResponseHandler -from tests.providers.microsoft.azure.base import Base -from tests.providers.microsoft.conftest import load_json, mock_json_response - - -class TestResponseHandler(Base): - def test_handle_response_async(self): - users = load_json("resources", "users.json") - response = mock_json_response(200, users) - - actual = self.run_async( - CallableResponseHandler(lambda response, error_map: response.json()).handle_response_async( - response, None - ) - ) - - assert isinstance(actual, dict) - assert actual == users diff --git a/tests/providers/microsoft/azure/serializer/test_serializer.py b/tests/providers/microsoft/azure/serializer/test_serializer.py deleted file mode 100644 index 445d9fe75ac01..0000000000000 --- a/tests/providers/microsoft/azure/serializer/test_serializer.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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 locale -from base64 import b64decode, b64encode -from datetime import datetime -from uuid import uuid4 - -import pendulum - -from airflow.providers.microsoft.azure.serialization.serializer import ResponseSerializer -from tests.providers.microsoft.conftest import load_file, load_json - - -class TestResponseSerializer: - def test_serialize_when_bytes_then_base64_encoded(self): - response = load_file("resources", "dummy.pdf", mode="rb", encoding=None) - content = b64encode(response).decode(locale.getpreferredencoding()) - - actual = ResponseSerializer().serialize(response) - - assert isinstance(actual, str) - assert actual == content - - def test_serialize_when_dict_with_uuid_datatime_and_pendulum_then_json(self): - id = uuid4() - response = { - "id": id, - "creationDate": datetime(2024, 2, 5), - "modificationTime": pendulum.datetime(2024, 2, 5), - } - - actual = ResponseSerializer().serialize(response) - - assert isinstance(actual, str) - assert ( - actual - == f'{{"id": "{id}", "creationDate": "2024-02-05T00:00:00", "modificationTime": "2024-02-05T00:00:00+00:00"}}' - ) - - def test_deserialize_when_json(self): - response = load_file("resources", "users.json") - - actual = ResponseSerializer().deserialize(response) - - assert isinstance(actual, dict) - assert actual == load_json("resources", "users.json") - - def test_deserialize_when_base64_encoded_string(self): - content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) - response = b64encode(content).decode(locale.getpreferredencoding()) - - actual = ResponseSerializer().deserialize(response) - - assert actual == response - assert b64decode(actual) == content diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index 4b943b4170944..cafd4e79afa3c 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -18,11 +18,16 @@ import json import locale -from base64 import b64encode +from base64 import b64encode, b64decode +from datetime import datetime from unittest.mock import patch +from uuid import uuid4 + +import pendulum from airflow.exceptions import AirflowException -from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger +from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger, CallableResponseHandler, \ + ResponseSerializer from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base from tests.providers.microsoft.conftest import ( @@ -114,5 +119,64 @@ def test_serialize(self): "timeout": None, "proxies": None, "api_version": "v1.0", - "serializer": "airflow.providers.microsoft.azure.serialization.serializer.ResponseSerializer", + "serializer": "airflow.providers.microsoft.azure.triggers.msgraph.ResponseSerializer", } + + +class TestResponseHandler(Base): + def test_handle_response_async(self): + users = load_json("resources", "users.json") + response = mock_json_response(200, users) + + actual = self.run_async( + CallableResponseHandler(lambda response, error_map: response.json()).handle_response_async( + response, None + ) + ) + + assert isinstance(actual, dict) + assert actual == users + + +class TestResponseSerializer: + def test_serialize_when_bytes_then_base64_encoded(self): + response = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + content = b64encode(response).decode(locale.getpreferredencoding()) + + actual = ResponseSerializer().serialize(response) + + assert isinstance(actual, str) + assert actual == content + + def test_serialize_when_dict_with_uuid_datatime_and_pendulum_then_json(self): + id = uuid4() + response = { + "id": id, + "creationDate": datetime(2024, 2, 5), + "modificationTime": pendulum.datetime(2024, 2, 5), + } + + actual = ResponseSerializer().serialize(response) + + assert isinstance(actual, str) + assert ( + actual + == f'{{"id": "{id}", "creationDate": "2024-02-05T00:00:00", "modificationTime": "2024-02-05T00:00:00+00:00"}}' + ) + + def test_deserialize_when_json(self): + response = load_file("resources", "users.json") + + actual = ResponseSerializer().deserialize(response) + + assert isinstance(actual, dict) + assert actual == load_json("resources", "users.json") + + def test_deserialize_when_base64_encoded_string(self): + content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) + response = b64encode(content).decode(locale.getpreferredencoding()) + + actual = ResponseSerializer().deserialize(response) + + assert actual == response + assert b64decode(actual) == content From 04553ca20292926212464a9237665bf714969ae9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 13:33:25 +0100 Subject: [PATCH 034/105] docs: Added documentation on how to use the MSGraphAsyncOperator --- .../microsoft/azure/operators/msgraph.py | 4 ++ .../operators/msgraph.rst | 57 +++++++++++++++++ .../microsoft/azure/example_msgraph.py | 64 +++++++++++++++++++ 3 files changed, 125 insertions(+) create mode 100644 docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst create mode 100644 tests/system/providers/microsoft/azure/example_msgraph.py diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 49f3d9b062829..49d6045f48419 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -51,6 +51,10 @@ class MSGraphAsyncOperator(BaseOperator): 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 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 url: The url being executed on the Microsoft Graph API (templated). diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst new file mode 100644 index 0000000000000..8ef239924acd0 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst @@ -0,0 +1,57 @@ + .. 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 Graph API Operators +============================= + +Prerequisite Tasks +^^^^^^^^^^^^^^^^^^ + +.. include:: /operators/_partials/prerequisite_tasks.rst + +.. _howto/operator:MSGraphAsyncOperator: + +MSGraphAsyncOperator +---------------------------------- +Use the +:class:`~airflow.providers.microsoft.azure.operators.msgraph.MSGraphAsyncOperator` to call Microsoft Graph API. + + +Below is an example of using this operator to get a Sharepoint site. + +.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_graph_site] + :end-before: [END howto_operator_graph_site] + +Below is an example of using this operator to get a Sharepoint site pages. + +.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_graph_site_pages] + :end-before: [END howto_operator_graph_site_pages] + + +Reference +--------- + +For further information, look at: + +* `Use the Microsoft Graph API `__ diff --git a/tests/system/providers/microsoft/azure/example_msgraph.py b/tests/system/providers/microsoft/azure/example_msgraph.py new file mode 100644 index 0000000000000..ee8e9f66a9978 --- /dev/null +++ b/tests/system/providers/microsoft/azure/example_msgraph.py @@ -0,0 +1,64 @@ +# 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 datetime + +from airflow import models +from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator + +DAG_ID = "example_sharepoint_site" + +with models.DAG( + DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + tags=["example"], +) as dag: + # [START howto_operator_graph_site] + site_task = MSGraphAsyncOperator( + task_id="news_site", + conn_id="msgraph_api", + url="sites/850v1v.sharepoint.com:/sites/news", + result_processor=lambda context, response: response["id"].split(",")[1], # only keep site_id + ) + # [END howto_operator_graph_site] + + # [START howto_operator_graph_site_pages] + site_pages_task = MSGraphAsyncOperator( + task_id="news_pages", + conn_id="msgraph_api", + api_version="beta", + url=( + "sites/%s/pages" + % "{{ ti.xcom_pull(task_ids='news_site') }}" + ), + ) + # [START howto_operator_graph_site_pages] + + site_task >> site_pages_task + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) From d835e34aa61468070af4d87c556cb020680e9bd9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 14:12:04 +0100 Subject: [PATCH 035/105] docs: Fixed END tag in examples --- tests/system/providers/microsoft/azure/example_msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/system/providers/microsoft/azure/example_msgraph.py b/tests/system/providers/microsoft/azure/example_msgraph.py index ee8e9f66a9978..e3f30c4343d56 100644 --- a/tests/system/providers/microsoft/azure/example_msgraph.py +++ b/tests/system/providers/microsoft/azure/example_msgraph.py @@ -48,7 +48,7 @@ % "{{ ti.xcom_pull(task_ids='news_site') }}" ), ) - # [START howto_operator_graph_site_pages] + # [END howto_operator_graph_site_pages] site_task >> site_pages_task From 6a14ebe01936ca31ab188ab0fcbb40ba1960c3ba Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 14:13:15 +0100 Subject: [PATCH 036/105] refactor: Removed docstring from CallableResponseHandler --- airflow/providers/microsoft/azure/triggers/msgraph.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 130ffa71567ea..111ff80899362 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -88,12 +88,6 @@ def deserialize(self, response) -> Any: class CallableResponseHandler(ResponseHandler): - """ - CallableResponseHandler executes the passed callable_function with response as parameter. - - :param callable_function: Function which allows you to handle the response before returning. - """ - def __init__( self, callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any], From 6cbdd9444f2e09e3cc605308cc6d8052a06e1ab9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 15:06:41 +0100 Subject: [PATCH 037/105] refactor: Ignore UP031 Use format specifiers instead of percent format as this is not possible here the way the DAG is evaluated in Airflow (due to XCom's) --- tests/system/providers/microsoft/azure/example_msgraph.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/system/providers/microsoft/azure/example_msgraph.py b/tests/system/providers/microsoft/azure/example_msgraph.py index e3f30c4343d56..6688f7080d8ed 100644 --- a/tests/system/providers/microsoft/azure/example_msgraph.py +++ b/tests/system/providers/microsoft/azure/example_msgraph.py @@ -44,6 +44,7 @@ conn_id="msgraph_api", api_version="beta", url=( + # noqa: UP031 "sites/%s/pages" % "{{ ti.xcom_pull(task_ids='news_site') }}" ), From 001f4927d25287e31610218043a772959df836a9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 15:07:01 +0100 Subject: [PATCH 038/105] Revert "refactor: Removed docstring from CallableResponseHandler" This reverts commit 6a14ebe01936ca31ab188ab0fcbb40ba1960c3ba. --- airflow/providers/microsoft/azure/triggers/msgraph.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 111ff80899362..130ffa71567ea 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -88,6 +88,12 @@ def deserialize(self, response) -> Any: class CallableResponseHandler(ResponseHandler): + """ + CallableResponseHandler executes the passed callable_function with response as parameter. + + :param callable_function: Function which allows you to handle the response before returning. + """ + def __init__( self, callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any], From 61a6ec8dccbd3f36e491a62fd9eadc6180dcd598 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 15:07:58 +0100 Subject: [PATCH 039/105] refactor: Simplified docstring on CallableResponseHandler --- airflow/providers/microsoft/azure/triggers/msgraph.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 130ffa71567ea..0141af0ade3f1 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -88,11 +88,7 @@ def deserialize(self, response) -> Any: class CallableResponseHandler(ResponseHandler): - """ - CallableResponseHandler executes the passed callable_function with response as parameter. - - :param callable_function: Function which allows you to handle the response before returning. - """ + """CallableResponseHandler executes the passed callable_function with response as parameter.""" def __init__( self, From 5beceef58ae4c96818a555c408a58496feb96c59 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 15 Mar 2024 17:57:40 +0100 Subject: [PATCH 040/105] refactor: Updated provider.yaml to add reference of msgraph to how-to-guide --- airflow/providers/microsoft/azure/provider.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 31cf0bbeca727..093535b3c79de 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -168,6 +168,8 @@ integrations: - integration-name: Microsoft Graph API external-doc-url: https://learn.microsoft.com/en-us/graph/use-the-api/ logo: /integration-logos/azure/Microsoft-Graph-API.png + how-to-guide: + - /docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst tags: [azure] operators: From 7170983653b8ae8fc1e77842cd110a646319c22d Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 18 Mar 2024 14:11:59 +0100 Subject: [PATCH 041/105] refactor: Updated docstrings on operator and trigger --- .../providers/microsoft/azure/operators/msgraph.py | 14 +++++++------- .../providers/microsoft/azure/triggers/msgraph.py | 7 ++++++- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 49d6045f48419..c60115e89b1d3 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -55,9 +55,15 @@ class MSGraphAsyncOperator(BaseOperator): 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is + lambda response, error_map: response.json()). The default expression will convert the native response + to JSON. If response_type parameter is specified, then the response_handler will be ignored. + :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 url: The url being executed on the Microsoft Graph API (templated). :param timeout: The HTTP timeout being used by the KiotaRequestAdapter (default is None). When no timeout is specified or set to None then no HTTP timeout is applied on each request. :param proxies: A Dict defining the HTTP proxies to be used (default is None). @@ -67,12 +73,6 @@ class MSGraphAsyncOperator(BaseOperator): :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 GraphServiceClientHook are bytes, then those will be base64 encoded into a string. - :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 response_handler: Function to convert the native HTTPX response returned by the hook (default is - lambda response, error_map: response.json()). The default expression will convert the native response - to JSON. If response_type parameter is specified, then the response_handler will be ignored. :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. """ diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 0141af0ade3f1..945e6704764ff 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -111,14 +111,19 @@ class MSGraphTrigger(BaseTrigger): :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 response_handler: Function to convert the native HTTPX response returned by the hook (default is + lambda response, error_map: response.json()). The default expression will convert the native response + to JSON. If response_type parameter is specified, then the response_handler will be ignored. :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 trigger against (templated). + :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 no HTTP timeout is applied 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. """ DEFAULT_HEADERS = {"Accept": "application/json;q=1"} From aa09267335e109b43efe15f62fbc2e86785c89e8 Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 18 Mar 2024 17:36:44 +0100 Subject: [PATCH 042/105] refactor: Fixed additional static checks --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- airflow/providers/microsoft/azure/triggers/msgraph.py | 2 +- tests/providers/microsoft/azure/triggers/test_msgraph.py | 9 ++++++--- .../system/providers/microsoft/azure/example_msgraph.py | 6 +----- 4 files changed, 9 insertions(+), 10 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index c60115e89b1d3..3ac98d81953c9 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -28,9 +28,9 @@ 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.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger from airflow.utils.xcom import XCOM_RETURN_KEY if TYPE_CHECKING: diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 945e6704764ff..2da65d7cee1f0 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -23,9 +23,9 @@ from contextlib import suppress from datetime import datetime from json import JSONDecodeError -from typing import Any from typing import ( TYPE_CHECKING, + Any, AsyncIterator, Callable, Sequence, diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index cafd4e79afa3c..fc5bc9bb0e685 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -18,7 +18,7 @@ import json import locale -from base64 import b64encode, b64decode +from base64 import b64decode, b64encode from datetime import datetime from unittest.mock import patch from uuid import uuid4 @@ -26,8 +26,11 @@ import pendulum from airflow.exceptions import AirflowException -from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger, CallableResponseHandler, \ - ResponseSerializer +from airflow.providers.microsoft.azure.triggers.msgraph import ( + CallableResponseHandler, + MSGraphTrigger, + ResponseSerializer, +) from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base from tests.providers.microsoft.conftest import ( diff --git a/tests/system/providers/microsoft/azure/example_msgraph.py b/tests/system/providers/microsoft/azure/example_msgraph.py index 6688f7080d8ed..4520cb74eaf55 100644 --- a/tests/system/providers/microsoft/azure/example_msgraph.py +++ b/tests/system/providers/microsoft/azure/example_msgraph.py @@ -43,11 +43,7 @@ task_id="news_pages", conn_id="msgraph_api", api_version="beta", - url=( - # noqa: UP031 - "sites/%s/pages" - % "{{ ti.xcom_pull(task_ids='news_site') }}" - ), + url=("sites/%s/pages" % "{{ ti.xcom_pull(task_ids='news_site') }}"), ) # [END howto_operator_graph_site_pages] From 8fec9a678293f17f796749a18975952583cb6326 Mon Sep 17 00:00:00 2001 From: David Blain Date: Tue, 19 Mar 2024 10:12:53 +0100 Subject: [PATCH 043/105] refactor: Ignore UP031 Use format specifiers instead of percent format as this is not possible here the way the DAG is evaluated in Airflow (due to XCom's) --- tests/system/providers/microsoft/azure/example_msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/system/providers/microsoft/azure/example_msgraph.py b/tests/system/providers/microsoft/azure/example_msgraph.py index 4520cb74eaf55..5ff7ba6f88835 100644 --- a/tests/system/providers/microsoft/azure/example_msgraph.py +++ b/tests/system/providers/microsoft/azure/example_msgraph.py @@ -43,7 +43,7 @@ task_id="news_pages", conn_id="msgraph_api", api_version="beta", - url=("sites/%s/pages" % "{{ ti.xcom_pull(task_ids='news_site') }}"), + url=("sites/%s/pages" % "{{ ti.xcom_pull(task_ids='news_site') }}"), # noqa: UP031 ) # [END howto_operator_graph_site_pages] From 52a2152c1b7f32f25b9c8408ef22728e5ee788ef Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 21 Mar 2024 08:18:04 +0100 Subject: [PATCH 044/105] refactor: Added param to docstring ResponseHandler --- airflow/providers/microsoft/azure/triggers/msgraph.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 2da65d7cee1f0..b6e4e0810c094 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -88,7 +88,11 @@ def deserialize(self, response) -> Any: class CallableResponseHandler(ResponseHandler): - """CallableResponseHandler executes the passed callable_function with response as parameter.""" + """ + CallableResponseHandler executes the passed callable_function with response as parameter. + + param callable_function: Function that is applied to the response. + """ def __init__( self, From 8e535b22fed35c33480bd366baaa37ba9c080c19 Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 25 Mar 2024 08:43:46 +0100 Subject: [PATCH 045/105] refactor: Updated pyproject.toml as main --- pyproject.toml | 1241 ++++-------------------------------------------- 1 file changed, 83 insertions(+), 1158 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d38053889fb73..77b7f9e2aee77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ requires = [ "GitPython==3.1.42", "gitdb==4.0.11", - "hatchling==1.22.3", + "hatchling==1.22.4", "packaging==24.0", "pathspec==0.12.1", "pluggy==1.4.0", @@ -36,8 +36,6 @@ build-backend = "hatchling.build" [project] name = "apache-airflow" -dynamic = ["version"] - description = "Programmatically author, schedule and monitor data pipelines" readme = { file = "generated/PYPI_README.md", content-type = "text/markdown" } license-files.globs = ["LICENSE", "3rd-party-licenses/*.txt"] @@ -63,1150 +61,88 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Topic :: System :: Monitoring", ] -# When you remove a dependency from the list, you should also make sure to add the dependency to be removed -# in the scripts/docker/install_airflow_dependencies_from_branch_tip.sh script DEPENDENCIES_TO_REMOVE -# in order to make sure the dependency is not installed in the CI image build process from the main -# of Airflow branch. After your PR is merged, you should remove it from the list there. -dependencies = [ - # Alembic is important to handle our migrations in predictable and performant way. It is developed - # together with SQLAlchemy. Our experience with Alembic is that it very stable in minor version - # The 1.13.0 of alembic marked some migration code as SQLAlchemy 2+ only so we limit it to 1.13.1 - "alembic>=1.13.1, <2.0", - "argcomplete>=1.10", - "asgiref", - "attrs>=22.1.0", - # Blinker use for signals in Flask, this is an optional dependency in Flask 2.2 and lower. - # In Flask 2.3 it becomes a mandatory dependency, and flask signals are always available. - "blinker>=1.6.2", - # Colorlog 6.x merges TTYColoredFormatter into ColoredFormatter, breaking backwards compatibility with 4.x - # Update CustomTTYColoredFormatter to remove - "colorlog>=4.0.2, <5.0", - "configupdater>=3.1.1", - # `airflow/www/extensions/init_views` imports `connexion.decorators.validation.RequestBodyValidator` - # connexion v3 has refactored the entire module to middleware, see: /spec-first/connexion/issues/1525 - # Specifically, RequestBodyValidator was removed in: /spec-first/connexion/pull/1595 - # The usage was added in #30596, seemingly only to override and improve the default error message. - # Either revert that change or find another way, preferably without using connexion internals. - # This limit can be removed after https://github.com/apache/airflow/issues/35234 is fixed - "connexion[flask]>=2.10.0,<3.0", - "cron-descriptor>=1.2.24", - "croniter>=2.0.2", - "cryptography>=39.0.0", - "deprecated>=1.2.13", - "dill>=0.2.2", - "flask-caching>=1.5.0", - # Flask-Session 0.6 add new arguments into the SqlAlchemySessionInterface constructor as well as - # all parameters now are mandatory which make AirflowDatabaseSessionInterface incopatible with this version. - "flask-session>=0.4.0,<0.6", - "flask-wtf>=0.15", - # Flask 2.3 is scheduled to introduce a number of deprecation removals - some of them might be breaking - # for our dependencies - notably `_app_ctx_stack` and `_request_ctx_stack` removals. - # We should remove the limitation after 2.3 is released and our dependencies are updated to handle it - "flask>=2.2,<2.3", - "fsspec>=2023.10.0", - "google-re2>=1.0", - "gunicorn>=20.1.0", - "httpx", - "importlib_metadata>=1.7;python_version<\"3.9\"", - # Importib_resources 6.2.0-6.3.1 break pytest_rewrite - # see https://github.com/python/importlib_resources/issues/299 - "importlib_resources>=5.2,!=6.2.0,!=6.3.0,!=6.3.1;python_version<\"3.9\"", - "itsdangerous>=2.0", - "jinja2>=3.0.0", - "jsonschema>=4.18.0", - "lazy-object-proxy", - "linkify-it-py>=2.0.0", - "lockfile>=0.12.2", - "markdown-it-py>=2.1.0", - "markupsafe>=1.1.1", - "marshmallow-oneofschema>=2.0.1", - "mdit-py-plugins>=0.3.0", - "opentelemetry-api>=1.15.0", - "opentelemetry-exporter-otlp", - "packaging>=14.0", - "pathspec>=0.9.0", - "pendulum>=2.1.2,<4.0", - "pluggy>=1.0", - "psutil>=4.2.0", - "pygments>=2.0.1", - "pyjwt>=2.0.0", - "python-daemon>=3.0.0", - "python-dateutil>=2.3", - "python-nvd3>=0.15.0", - "python-slugify>=5.0", - # Requests 3 if it will be released, will be heavily breaking. - "requests>=2.27.0,<3", - "rfc3339-validator>=0.1.4", - "rich-argparse>=1.0.0", - "rich>=12.4.4", - "setproctitle>=1.1.8", - # We use some deprecated features of sqlalchemy 2.0 and we should replace them before we can upgrade - # See https://sqlalche.me/e/b8d9 for details of deprecated features - # you can set environment variable SQLALCHEMY_WARN_20=1 to show all deprecation warnings. - # The issue tracking it is https://github.com/apache/airflow/issues/28723 - "sqlalchemy>=1.4.36,<2.0", - "sqlalchemy-jsonfield>=1.0", - "tabulate>=0.7.5", - "tenacity>=6.2.0,!=8.2.0", - "termcolor>=1.1.0", - # We should remove this dependency when Providers are limited to Airflow 2.7+ - # as we replaced the usage of unicodecsv with csv in Airflow 2.7 - # See https://github.com/apache/airflow/pull/31693 - # We should also remove "licenses/LICENSE-unicodecsv.txt" file when we remove this dependency - "unicodecsv>=0.14.1", - # The Universal Pathlib provides Pathlib-like interface for FSSPEC - "universal-pathlib>=0.2.2", - # Werkzug 3 breaks Flask-Login 0.6.2, also connexion needs to be updated to >= 3.0 - # we should remove this limitation when FAB supports Flask 2.3 and we migrate connexion to 3+ - "werkzeug>=2.0,<3", -] -[project.optional-dependencies] -# Here manually managed extras start -# Those extras are manually managed and should be updated when needed +dynamic = ["version", "optional-dependencies", "dependencies"] + +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! +# !!! YOU MIGHT BE SURPRISED NOT SEEING THE DEPENDENCIES AS `project.dependencies` !!!!!!!!! +# !!! AND EXTRAS AS `project.optional-dependencies` !!!!!!!!! +# !!! THEY ARE marked as `dynamic` GENERATED by `hatch_build.py` !!!!!!!!! +# !!! SEE COMMENTS BELOW TO FIND WHERE DEPENDENCIES ARE MAINTAINED !!!!!!!!! +# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # -# START OF core extras +# !!!!!! Those provuders are defined in `hatch_build.py` and should be maintained there !!!!!!! # -# This required for AWS deferrable operators. -# There is conflict between boto3 and aiobotocore dependency botocore. -# TODO: We can remove it once boto3 and aiobotocore both have compatible botocore version or -# boto3 have native aync support and we move away from aio aiobotocore +# Those extras are available as regular core airflow extras - they install optional features of Airflow. # -aiobotocore = [ - "aiobotocore>=2.7.0", -] -async = [ - "eventlet>=0.33.3", - "gevent>=0.13", - "greenlet>=0.4.9", -] -cgroups = [ - # Cgroupspy 0.2.2 added Python 3.10 compatibility - "cgroupspy>=0.2.2", -] -deprecated-api = [ - "requests>=2.27.0,<3", -] -github-enterprise = [ - "apache-airflow[fab]", - "authlib>=1.0.0", -] -google-auth = [ - "apache-airflow[fab]", - "authlib>=1.0.0", -] -graphviz = [ - "graphviz>=0.12", -] -kerberos = [ - "pykerberos>=1.1.13", - "requests-kerberos>=0.10.0", - "thrift-sasl>=0.2.0", -] -ldap = [ - "ldap3>=2.5.1", - "python-ldap", -] -leveldb = [ - "plyvel", -] -otel = [ - "opentelemetry-exporter-prometheus", -] -pandas = [ - # In pandas 2.2 minimal version of the sqlalchemy is 2.0 - # https://pandas.pydata.org/docs/whatsnew/v2.2.0.html#increased-minimum-versions-for-dependencies - # However Airflow not fully supports it yet: https://github.com/apache/airflow/issues/28723 - # In addition FAB also limit sqlalchemy to < 2.0 - "pandas>=1.2.5,<2.2", -] -password = [ - "bcrypt>=2.0.0", - "flask-bcrypt>=0.7.1", -] -pydantic = [ - "pydantic>=2.3.0", -] -rabbitmq = [ - "amqp", -] -s3fs = [ - # This is required for support of S3 file system which uses aiobotocore - # which can have a conflict with boto3 as mentioned in aiobotocore extra - "s3fs>=2023.10.0", -] -saml = [ - # This is required for support of SAML which might be used by some providers (e.g. Amazon) - "python3-saml>=1.16.0", -] -sentry = [ - "blinker>=1.1", - # Sentry SDK 1.33 is broken when greenlets are installed and fails to import - # See https://github.com/getsentry/sentry-python/issues/2473 - "sentry-sdk>=1.32.0,!=1.33.0", -] -statsd = [ - "statsd>=3.3.0", -] -uv = [ - "uv>=0.1.22", -] -virtualenv = [ - "virtualenv", -] -# END OF core extras -# START OF Apache no provider extras -apache-atlas = [ - "atlasclient>=0.1.2", -] -apache-webhdfs = [ - "hdfs[avro,dataframe,kerberos]>=2.0.4", -] -# END OF Apache no provider extras -all-core = [ - "apache-airflow[aiobotocore]", - "apache-airflow[apache-atlas]", - "apache-airflow[async]", - "apache-airflow[cgroups]", - "apache-airflow[deprecated-api]", - "apache-airflow[github-enterprise]", - "apache-airflow[google-auth]", - "apache-airflow[graphviz]", - "apache-airflow[kerberos]", - "apache-airflow[ldap]", - "apache-airflow[leveldb]", - "apache-airflow[otel]", - "apache-airflow[pandas]", - "apache-airflow[password]", - "apache-airflow[pydantic]", - "apache-airflow[rabbitmq]", - "apache-airflow[s3fs]", - "apache-airflow[saml]", - "apache-airflow[sentry]", - "apache-airflow[statsd]", - "apache-airflow[apache-webhdfs]", - "apache-airflow[virtualenv]", -] -# START OF devel extras -devel-debuggers = [ - "ipdb>=0.13.13", -] -devel-devscripts = [ - "click>=8.0", - "gitpython>=3.1.40", - "hatch>=1.9.1", - "pipdeptree>=2.13.1", - "pygithub>=2.1.1", - "restructuredtext-lint>=1.4.0", - "rich-click>=1.7.0", - "semver>=3.0.2", - "towncrier>=23.11.0", - "twine>=4.0.2", -] -devel-duckdb = [ - # Python 3.12 support was added in 0.10.0 - "duckdb>=0.10.0; python_version >= '3.12'", - "duckdb>=0.9.0; python_version < '3.12'", -] -# Mypy 0.900 and above ships only with stubs from stdlib so if we need other stubs, we need to install them -# manually as `types-*`. See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports -# for details. We want to install them explicitly because we want to eventually move to -# mypyd which does not support installing the types dynamically with --install-types -devel-mypy = [ - # TODO: upgrade to newer versions of MyPy continuously as they are released - # Make sure to upgrade the mypy version in update-common-sql-api-stubs in .pre-commit-config.yaml - # when you upgrade it here !!!! - "mypy==1.9.0", - "types-Deprecated", - "types-Markdown", - "types-PyMySQL", - "types-PyYAML", - "types-aiofiles", - "types-certifi", - "types-croniter", - "types-docutils", - "types-paramiko", - "types-protobuf", - "types-python-dateutil", - "types-python-slugify", - "types-pytz", - "types-redis", - "types-requests", - "types-setuptools", - "types-tabulate", - "types-termcolor", - "types-toml", -] -devel-sentry = [ - "blinker>=1.7.0", -] -devel-static-checks = [ - "black>=23.12.0", - "pre-commit>=3.5.0", - "ruff==0.3.3", - "yamllint>=1.33.0", -] -devel-tests = [ - "aiofiles>=23.2.0", - "aioresponses>=0.7.6", - "backports.zoneinfo>=0.2.1;python_version<'3.9'", - "beautifulsoup4>=4.7.1", - # Coverage 7.4.0 added experimental support for Python 3.12 PEP669 which we use in Airflow - "coverage>=7.4.0", - "pytest-asyncio>=0.23.3", - "pytest-cov>=4.1.0", - "pytest-icdiff>=0.9", - "pytest-instafail>=0.5.0", - "pytest-mock>=3.12.0", - "pytest-rerunfailures>=13.0", - "pytest-timeouts>=1.2.1", - "pytest-xdist>=3.5.0", - # Temporary upper limmit to <8, not all dependencies at that moment ready to use 8.0 - # Internal meta-task for track https://github.com/apache/airflow/issues/37156 - "pytest>=7.4.4,<8.0", - "requests_mock>=1.11.0", - "time-machine>=2.13.0", - "wheel>=0.42.0", -] -# END OF devel extras -# START OF doc extras -doc = [ - "astroid>=2.12.3,<3.0", - "checksumdir>=1.2.0", - # click 8.1.4 and 8.1.5 generate mypy errors due to typing issue in the upstream package: - # https://github.com/pallets/click/issues/2558 - "click>=8.0,!=8.1.4,!=8.1.5", - # Docutils 0.17.0 converts generated
into
and breaks our doc formatting - # By adding a lot of whitespace separation. This limit can be lifted when we update our doc to handle - #
tags for sections - "docutils<0.17,>=0.16", - "sphinx-airflow-theme>=0.0.12", - "sphinx-argparse>=0.4.0", - # sphinx-autoapi fails with astroid 3.0, see: https://github.com/readthedocs/sphinx-autoapi/issues/407 - # This was fixed in sphinx-autoapi 3.0, however it has requirement sphinx>=6.1, but we stuck on 5.x - "sphinx-autoapi>=2.1.1", - "sphinx-copybutton>=0.5.2", - "sphinx-design>=0.5.0", - "sphinx-jinja>=2.0.2", - "sphinx-rtd-theme>=2.0.0", - # Currently we are using sphinx 5 but we need to migrate to Sphinx 7 - "sphinx>=5.3.0,<6.0.0", - "sphinxcontrib-applehelp>=1.0.4", - "sphinxcontrib-devhelp>=1.0.2", - "sphinxcontrib-htmlhelp>=2.0.1", - "sphinxcontrib-httpdomain>=1.8.1", - "sphinxcontrib-jquery>=4.1", - "sphinxcontrib-jsmath>=1.0.1", - "sphinxcontrib-qthelp>=1.0.3", - "sphinxcontrib-redoc>=1.6.0", - "sphinxcontrib-serializinghtml==1.1.5", - "sphinxcontrib-spelling>=8.0.0", -] -doc-gen = [ - "apache-airflow[doc]", - "eralchemy2>=1.3.8", -] -# END OF doc extras -# START OF bundle extras -all-dbs = [ - "apache-airflow[apache-cassandra]", - "apache-airflow[apache-drill]", - "apache-airflow[apache-druid]", - "apache-airflow[apache-hdfs]", - "apache-airflow[apache-hive]", - "apache-airflow[apache-impala]", - "apache-airflow[apache-pinot]", - "apache-airflow[arangodb]", - "apache-airflow[cloudant]", - "apache-airflow[databricks]", - "apache-airflow[exasol]", - "apache-airflow[influxdb]", - "apache-airflow[microsoft-mssql]", - "apache-airflow[mongo]", - "apache-airflow[mysql]", - "apache-airflow[neo4j]", - "apache-airflow[postgres]", - "apache-airflow[presto]", - "apache-airflow[trino]", - "apache-airflow[vertica]", -] -devel = [ - "apache-airflow[celery]", - "apache-airflow[cncf-kubernetes]", - "apache-airflow[common-io]", - "apache-airflow[common-sql]", - "apache-airflow[devel-debuggers]", - "apache-airflow[devel-devscripts]", - "apache-airflow[devel-duckdb]", - "apache-airflow[devel-mypy]", - "apache-airflow[devel-sentry]", - "apache-airflow[devel-static-checks]", - "apache-airflow[devel-tests]", - "apache-airflow[fab]", - "apache-airflow[ftp]", - "apache-airflow[http]", - "apache-airflow[imap]", - "apache-airflow[sqlite]", -] -devel-all-dbs = [ - "apache-airflow[apache-cassandra]", - "apache-airflow[apache-drill]", - "apache-airflow[apache-druid]", - "apache-airflow[apache-hdfs]", - "apache-airflow[apache-hive]", - "apache-airflow[apache-impala]", - "apache-airflow[apache-pinot]", - "apache-airflow[arangodb]", - "apache-airflow[cloudant]", - "apache-airflow[databricks]", - "apache-airflow[exasol]", - "apache-airflow[influxdb]", - "apache-airflow[microsoft-mssql]", - "apache-airflow[mongo]", - "apache-airflow[mysql]", - "apache-airflow[neo4j]", - "apache-airflow[postgres]", - "apache-airflow[presto]", - "apache-airflow[trino]", - "apache-airflow[vertica]", -] -devel-ci = [ - "apache-airflow[devel-all]", -] -devel-hadoop = [ - "apache-airflow[apache-hdfs]", - "apache-airflow[apache-hive]", - "apache-airflow[apache-impala]", - "apache-airflow[devel]", - "apache-airflow[hdfs]", - "apache-airflow[kerberos]", - "apache-airflow[presto]", -] -# END OF bundle extras -############################################################################################################# -# The whole section can be removed in Airflow 3.0 as those old aliases are deprecated in 2.* series -############################################################################################################# -# START OF deprecated extras -atlas = [ - "apache-airflow[apache-atlas]", -] -aws = [ - "apache-airflow[amazon]", -] -azure = [ - "apache-airflow[microsoft-azure]", -] -cassandra = [ - "apache-airflow[apache-cassandra]", -] -# Empty alias extra just for backward compatibility with Airflow 1.10 -crypto = [ -] -druid = [ - "apache-airflow[apache-druid]", -] -gcp = [ - "apache-airflow[google]", -] -gcp_api = [ - "apache-airflow[google]", -] -hdfs = [ - "apache-airflow[apache-hdfs]", -] -hive = [ - "apache-airflow[apache-hive]", -] -kubernetes = [ - "apache-airflow[cncf-kubernetes]", -] -mssql = [ - "apache-airflow[microsoft-mssql]", -] -pinot = [ - "apache-airflow[apache-pinot]", -] -s3 = [ - "apache-airflow[amazon]", -] -spark = [ - "apache-airflow[apache-spark]", -] -webhdfs = [ - "apache-airflow[apache-webhdfs]", -] -winrm = [ - "apache-airflow[microsoft-winrm]", -] -# END OF deprecated extras -############################################################################################################# -# The whole section below is automatically generated by `update-providers-dependencies` pre-commit based -# on `provider.yaml` files present in the `providers` subdirectories. The `provider.yaml` files are -# A single source of truth for provider dependencies, +# START CORE EXTRAS HERE # -# PLEASE DO NOT MODIFY THIS SECTION MANUALLY. IT WILL BE OVERWRITTEN BY PRE-COMMIT !! -# If you want to modify these - modify the corresponding provider.yaml instead. -############################################################################################################# -# START OF GENERATED DEPENDENCIES -airbyte = [ # source: airflow/providers/airbyte/provider.yaml - "apache-airflow[http]", -] -alibaba = [ # source: airflow/providers/alibaba/provider.yaml - "alibabacloud_adb20211201>=1.0.0", - "alibabacloud_tea_openapi>=0.3.7", - "oss2>=2.14.0", -] -amazon = [ # source: airflow/providers/amazon/provider.yaml - "PyAthena>=3.0.10", - "apache-airflow[common_sql]", - "apache-airflow[http]", - "asgiref", - "boto3>=1.33.0", - "botocore>=1.33.0", - "inflection>=0.5.1", - "jsonpath_ng>=1.5.3", - "redshift_connector>=2.0.918", - "sqlalchemy_redshift>=0.8.6", - "watchtower>=2.0.1,<4", - # Devel dependencies for the amazon provider - "aiobotocore>=2.7.0", - "aws_xray_sdk>=2.12.0", - "moto[cloudformation,glue]>=5.0.0", - "mypy-boto3-appflow>=1.33.0", - "mypy-boto3-rds>=1.33.0", - "mypy-boto3-redshift-data>=1.33.0", - "mypy-boto3-s3>=1.33.0", - "s3fs>=2023.10.0", - "openapi-schema-validator>=0.6.2", - "openapi-spec-validator>=0.7.1", -] -apache-beam = [ # source: airflow/providers/apache/beam/provider.yaml - "apache-beam>=2.53.0;python_version != \"3.12\"", - "pyarrow>=14.0.1;python_version != \"3.12\"", -] -apache-cassandra = [ # source: airflow/providers/apache/cassandra/provider.yaml - "cassandra-driver>=3.29.1", -] -apache-drill = [ # source: airflow/providers/apache/drill/provider.yaml - "apache-airflow[common_sql]", - "sqlalchemy-drill>=1.1.0", -] -apache-druid = [ # source: airflow/providers/apache/druid/provider.yaml - "apache-airflow[common_sql]", - "pydruid>=0.4.1", -] -apache-flink = [ # source: airflow/providers/apache/flink/provider.yaml - "apache-airflow[cncf_kubernetes]", - "cryptography>=2.0.0", -] -apache-hdfs = [ # source: airflow/providers/apache/hdfs/provider.yaml - "hdfs[avro,dataframe,kerberos]>=2.0.4", -] -apache-hive = [ # source: airflow/providers/apache/hive/provider.yaml - "apache-airflow[common_sql]", - "hmsclient>=0.1.0", - "pandas>=1.2.5,<2.2", - "pyhive[hive_pure_sasl]>=0.7.0", - "thrift>=0.9.2", -] -apache-impala = [ # source: airflow/providers/apache/impala/provider.yaml - "impyla>=0.18.0,<1.0", -] -apache-kafka = [ # source: airflow/providers/apache/kafka/provider.yaml - "asgiref", - "confluent-kafka>=1.8.2", -] -apache-kylin = [ # source: airflow/providers/apache/kylin/provider.yaml - "kylinpy>=2.6", -] -apache-livy = [ # source: airflow/providers/apache/livy/provider.yaml - "aiohttp>=3.9.2", - "apache-airflow[http]", - "asgiref", -] -apache-pig = [] # source: airflow/providers/apache/pig/provider.yaml -apache-pinot = [ # source: airflow/providers/apache/pinot/provider.yaml - "apache-airflow[common_sql]", - "pinotdb>=5.1.0", -] -apache-spark = [ # source: airflow/providers/apache/spark/provider.yaml - "grpcio-status>=1.59.0", - "pyspark", -] -apprise = [ # source: airflow/providers/apprise/provider.yaml - "apprise", -] -arangodb = [ # source: airflow/providers/arangodb/provider.yaml - "python-arango>=7.3.2", -] -asana = [ # source: airflow/providers/asana/provider.yaml - "asana>=0.10,<4.0.0", -] -atlassian-jira = [ # source: airflow/providers/atlassian/jira/provider.yaml - "atlassian-python-api>=1.14.2,!=3.41.6", - "beautifulsoup4", -] -celery = [ # source: airflow/providers/celery/provider.yaml - "celery>=5.3.0,<6,!=5.3.3,!=5.3.2", - "flower>=1.0.0", - "google-re2>=1.0", -] -cloudant = [ # source: airflow/providers/cloudant/provider.yaml - "cloudant>=2.0", -] -cncf-kubernetes = [ # source: airflow/providers/cncf/kubernetes/provider.yaml - "aiofiles>=23.2.0", - "asgiref>=3.5.2", - "cryptography>=2.0.0", - "google-re2>=1.0", - "kubernetes>=28.1.0,<=29.0.0", - "kubernetes_asyncio>=28.1.0,<=29.0.0", -] -cohere = [ # source: airflow/providers/cohere/provider.yaml - "cohere>=4.37,<5", -] -common-io = [] # source: airflow/providers/common/io/provider.yaml -common-sql = [ # source: airflow/providers/common/sql/provider.yaml - "more-itertools>=9.0.0", - "sqlparse>=0.4.2", -] -databricks = [ # source: airflow/providers/databricks/provider.yaml - "aiohttp>=3.9.2, <4", - "apache-airflow[common_sql]", - "databricks-sql-connector>=2.0.0, <3.0.0, !=2.9.0", - "requests>=2.27.0,<3", - # Devel dependencies for the databricks provider - "deltalake>=0.12.0", -] -datadog = [ # source: airflow/providers/datadog/provider.yaml - "datadog>=0.14.0", -] -dbt-cloud = [ # source: airflow/providers/dbt/cloud/provider.yaml - "aiohttp>=3.9.2", - "apache-airflow[http]", - "asgiref", -] -dingding = [ # source: airflow/providers/dingding/provider.yaml - "apache-airflow[http]", -] -discord = [ # source: airflow/providers/discord/provider.yaml - "apache-airflow[http]", -] -docker = [ # source: airflow/providers/docker/provider.yaml - "docker>=6", - "python-dotenv>=0.21.0", -] -elasticsearch = [ # source: airflow/providers/elasticsearch/provider.yaml - "apache-airflow[common_sql]", - "elasticsearch>=8.10,<9", -] -exasol = [ # source: airflow/providers/exasol/provider.yaml - "apache-airflow[common_sql]", - "pandas>=1.2.5,<2.2", - "pyexasol>=0.5.1", -] -fab = [ # source: airflow/providers/fab/provider.yaml - "flask-appbuilder==4.4.1", - "flask-login>=0.6.2", - "flask>=2.2,<2.3", - "google-re2>=1.0", -] -facebook = [ # source: airflow/providers/facebook/provider.yaml - "facebook-business>=6.0.2", -] -ftp = [] # source: airflow/providers/ftp/provider.yaml -github = [ # source: airflow/providers/github/provider.yaml - "PyGithub!=1.58", -] -google = [ # source: airflow/providers/google/provider.yaml - "PyOpenSSL", - "apache-airflow[common_sql]", - "asgiref>=3.5.2", - "gcloud-aio-auth>=4.0.0,<5.0.0", - "gcloud-aio-bigquery>=6.1.2", - "gcloud-aio-storage>=9.0.0", - "gcsfs>=2023.10.0", - "google-ads>=23.1.0", - "google-analytics-admin", - "google-api-core>=2.11.0,!=2.16.0", - "google-api-python-client>=1.6.0", - "google-auth-httplib2>=0.0.1", - "google-auth>=1.0.0", - "google-cloud-aiplatform>=1.42.1", - "google-cloud-automl>=2.12.0", - "google-cloud-batch>=0.13.0", - "google-cloud-bigquery-datatransfer>=3.13.0", - "google-cloud-bigtable>=2.17.0", - "google-cloud-build>=3.22.0", - "google-cloud-compute>=1.10.0", - "google-cloud-container>=2.17.4", - "google-cloud-datacatalog>=3.11.1", - "google-cloud-dataflow-client>=0.8.6", - "google-cloud-dataform>=0.5.0", - "google-cloud-dataplex>=1.10.0", - "google-cloud-dataproc-metastore>=1.12.0", - "google-cloud-dataproc>=5.8.0", - "google-cloud-dlp>=3.12.0", - "google-cloud-kms>=2.15.0", - "google-cloud-language>=2.9.0", - "google-cloud-logging>=3.5.0", - "google-cloud-memcache>=1.7.0", - "google-cloud-monitoring>=2.18.0", - "google-cloud-orchestration-airflow>=1.10.0", - "google-cloud-os-login>=2.9.1", - "google-cloud-pubsub>=2.19.0", - "google-cloud-redis>=2.12.0", - "google-cloud-run>=0.9.0", - "google-cloud-secret-manager>=2.16.0", - "google-cloud-spanner>=3.11.1", - "google-cloud-speech>=2.18.0", - "google-cloud-storage-transfer>=1.4.1", - "google-cloud-storage>=2.7.0", - "google-cloud-tasks>=2.13.0", - "google-cloud-texttospeech>=2.14.1", - "google-cloud-translate>=3.11.0", - "google-cloud-videointelligence>=2.11.0", - "google-cloud-vision>=3.4.0", - "google-cloud-workflows>=1.10.0", - "grpcio-gcp>=0.2.2", - "httpx", - "json-merge-patch>=0.2", - "looker-sdk>=22.2.0", - "pandas-gbq", - "pandas>=1.2.5,<2.2", - "proto-plus>=1.19.6", - "python-slugify>=5.0", - "sqlalchemy-bigquery>=1.2.1", - "sqlalchemy-spanner>=1.6.2", -] -grpc = [ # source: airflow/providers/grpc/provider.yaml - "google-auth-httplib2>=0.0.1", - "google-auth>=1.0.0, <3.0.0", - "grpcio>=1.15.0", -] -hashicorp = [ # source: airflow/providers/hashicorp/provider.yaml - "hvac>=1.1.0", -] -http = [ # source: airflow/providers/http/provider.yaml - "aiohttp>=3.9.2", - "asgiref", - "requests>=2.27.0,<3", - "requests_toolbelt", -] -imap = [] # source: airflow/providers/imap/provider.yaml -influxdb = [ # source: airflow/providers/influxdb/provider.yaml - "influxdb-client>=1.19.0", - "requests>=2.27.0,<3", -] -jdbc = [ # source: airflow/providers/jdbc/provider.yaml - "apache-airflow[common_sql]", - "jaydebeapi>=1.1.1", -] -jenkins = [ # source: airflow/providers/jenkins/provider.yaml - "python-jenkins>=1.0.0", -] -microsoft-azure = [ # source: airflow/providers/microsoft/azure/provider.yaml - "adal>=1.2.7", - "adlfs>=2023.10.0", - "azure-batch>=8.0.0", - "azure-cosmos>=4.0.0,<4.6.0", - "azure-datalake-store>=0.0.45", - "azure-identity>=1.3.1", - "azure-keyvault-secrets>=4.1.0", - "azure-kusto-data>=4.1.0", - "azure-mgmt-containerinstance>=9.0.0", - "azure-mgmt-containerregistry>=8.0.0", - "azure-mgmt-cosmosdb", - "azure-mgmt-datafactory>=2.0.0", - "azure-mgmt-datalake-store>=0.5.0", - "azure-mgmt-resource>=2.2.0", - "azure-mgmt-storage>=16.0.0", - "azure-servicebus>=7.6.1", - "azure-storage-blob>=12.14.0", - "azure-storage-file-datalake>=12.9.1", - "azure-storage-file-share", - "azure-synapse-artifacts>=0.17.0", - "azure-synapse-spark", - "msgraph-core>=1.0.0", - # Devel dependencies for the microsoft.azure provider - "pywinrm", -] -microsoft-mssql = [ # source: airflow/providers/microsoft/mssql/provider.yaml - "apache-airflow[common_sql]", - "pymssql>=2.1.8", -] -microsoft-psrp = [ # source: airflow/providers/microsoft/psrp/provider.yaml - "pypsrp>=0.8.0", -] -microsoft-winrm = [ # source: airflow/providers/microsoft/winrm/provider.yaml - "pywinrm>=0.4", -] -mongo = [ # source: airflow/providers/mongo/provider.yaml - "dnspython>=1.13.0", - "pymongo>=3.6.0", - # Devel dependencies for the mongo provider - "mongomock", -] -mysql = [ # source: airflow/providers/mysql/provider.yaml - "apache-airflow[common_sql]", - "mysql-connector-python>=8.0.29", - "mysqlclient>=1.3.6", -] -neo4j = [ # source: airflow/providers/neo4j/provider.yaml - "neo4j>=4.2.1", -] -odbc = [ # source: airflow/providers/odbc/provider.yaml - "apache-airflow[common_sql]", - "pyodbc", -] -openai = [ # source: airflow/providers/openai/provider.yaml - "openai[datalib]>=1.0", -] -openfaas = [] # source: airflow/providers/openfaas/provider.yaml -openlineage = [ # source: airflow/providers/openlineage/provider.yaml - "apache-airflow[common_sql]", - "attrs>=22.2", - "openlineage-integration-common>=0.28.0", - "openlineage-python>=0.28.0", -] -opensearch = [ # source: airflow/providers/opensearch/provider.yaml - "opensearch-py>=2.2.0", -] -opsgenie = [ # source: airflow/providers/opsgenie/provider.yaml - "opsgenie-sdk>=2.1.5", -] -oracle = [ # source: airflow/providers/oracle/provider.yaml - "apache-airflow[common_sql]", - "oracledb>=1.0.0", -] -pagerduty = [ # source: airflow/providers/pagerduty/provider.yaml - "pdpyras>=4.1.2", -] -papermill = [ # source: airflow/providers/papermill/provider.yaml - "ipykernel;python_version != \"3.12\"", - "papermill[all]>=2.4.0;python_version != \"3.12\"", - "scrapbook[all];python_version != \"3.12\"", -] -pgvector = [ # source: airflow/providers/pgvector/provider.yaml - "apache-airflow[postgres]", - "pgvector>=0.2.3", -] -pinecone = [ # source: airflow/providers/pinecone/provider.yaml - "pinecone-client>=2.2.4,<3.0", -] -postgres = [ # source: airflow/providers/postgres/provider.yaml - "apache-airflow[common_sql]", - "psycopg2-binary>=2.8.0", -] -presto = [ # source: airflow/providers/presto/provider.yaml - "apache-airflow[common_sql]", - "pandas>=1.2.5,<2.2", - "presto-python-client>=0.8.4", -] -qdrant = [ # source: airflow/providers/qdrant/provider.yaml - "qdrant_client>=1.7.0", -] -redis = [ # source: airflow/providers/redis/provider.yaml - "redis>=4.5.2,<5.0.0,!=4.5.5", -] -salesforce = [ # source: airflow/providers/salesforce/provider.yaml - "pandas>=1.2.5,<2.2", - "simple-salesforce>=1.0.0", -] -samba = [ # source: airflow/providers/samba/provider.yaml - "smbprotocol>=1.5.0", -] -segment = [ # source: airflow/providers/segment/provider.yaml - "analytics-python>=1.2.9", -] -sendgrid = [ # source: airflow/providers/sendgrid/provider.yaml - "sendgrid>=6.0.0", -] -sftp = [ # source: airflow/providers/sftp/provider.yaml - "apache-airflow[ssh]", - "asyncssh>=2.12.0", - "paramiko>=2.8.0", -] -singularity = [ # source: airflow/providers/singularity/provider.yaml - "spython>=0.0.56", -] -slack = [ # source: airflow/providers/slack/provider.yaml - "apache-airflow[common_sql]", - "slack_sdk>=3.19.0", -] -smtp = [] # source: airflow/providers/smtp/provider.yaml -snowflake = [ # source: airflow/providers/snowflake/provider.yaml - "apache-airflow[common_sql]", - "snowflake-connector-python>=2.7.8", - "snowflake-sqlalchemy>=1.1.0", -] -sqlite = [ # source: airflow/providers/sqlite/provider.yaml - "apache-airflow[common_sql]", -] -ssh = [ # source: airflow/providers/ssh/provider.yaml - "paramiko>=2.6.0", - "sshtunnel>=0.3.2", -] -tableau = [ # source: airflow/providers/tableau/provider.yaml - "tableauserverclient", -] -tabular = [ # source: airflow/providers/tabular/provider.yaml - # Devel dependencies for the tabular provider - "pyiceberg>=0.5.0", -] -telegram = [ # source: airflow/providers/telegram/provider.yaml - "python-telegram-bot>=20.2", -] -teradata = [ # source: airflow/providers/teradata/provider.yaml - "apache-airflow[common_sql]", - "teradatasql>=17.20.0.28", - "teradatasqlalchemy>=17.20.0.0", -] -trino = [ # source: airflow/providers/trino/provider.yaml - "apache-airflow[common_sql]", - "pandas>=1.2.5,<2.2", - "trino>=0.318.0", -] -vertica = [ # source: airflow/providers/vertica/provider.yaml - "apache-airflow[common_sql]", - "vertica-python>=0.5.1", -] -weaviate = [ # source: airflow/providers/weaviate/provider.yaml - "pandas>=1.2.5,<2.2", - "weaviate-client>=3.24.2", -] -yandex = [ # source: airflow/providers/yandex/provider.yaml - "python-dateutil>=2.8.0", - "requests>=2.27.0,<3", - "yandex-query-client>=0.1.2", - "yandexcloud>=0.228.0", -] -zendesk = [ # source: airflow/providers/zendesk/provider.yaml - "zenpy>=2.0.40", -] -all = [ - # core extras - "apache-airflow[aiobotocore]", - "apache-airflow[async]", - "apache-airflow[cgroups]", - "apache-airflow[deprecated-api]", - "apache-airflow[github-enterprise]", - "apache-airflow[google-auth]", - "apache-airflow[graphviz]", - "apache-airflow[kerberos]", - "apache-airflow[ldap]", - "apache-airflow[leveldb]", - "apache-airflow[otel]", - "apache-airflow[pandas]", - "apache-airflow[password]", - "apache-airflow[pydantic]", - "apache-airflow[rabbitmq]", - "apache-airflow[s3fs]", - "apache-airflow[saml]", - "apache-airflow[sentry]", - "apache-airflow[statsd]", - "apache-airflow[uv]", - "apache-airflow[virtualenv]", - # Apache no provider extras - "apache-airflow[apache-atlas]", - "apache-airflow[apache-webhdfs]", - "apache-airflow[all-core]", - # Provider extras - "apache-airflow[airbyte]", - "apache-airflow[alibaba]", - "apache-airflow[amazon]", - "apache-airflow[apache-beam]", - "apache-airflow[apache-cassandra]", - "apache-airflow[apache-drill]", - "apache-airflow[apache-druid]", - "apache-airflow[apache-flink]", - "apache-airflow[apache-hdfs]", - "apache-airflow[apache-hive]", - "apache-airflow[apache-impala]", - "apache-airflow[apache-kafka]", - "apache-airflow[apache-kylin]", - "apache-airflow[apache-livy]", - "apache-airflow[apache-pig]", - "apache-airflow[apache-pinot]", - "apache-airflow[apache-spark]", - "apache-airflow[apprise]", - "apache-airflow[arangodb]", - "apache-airflow[asana]", - "apache-airflow[atlassian-jira]", - "apache-airflow[celery]", - "apache-airflow[cloudant]", - "apache-airflow[cncf-kubernetes]", - "apache-airflow[cohere]", - "apache-airflow[common-io]", - "apache-airflow[common-sql]", - "apache-airflow[databricks]", - "apache-airflow[datadog]", - "apache-airflow[dbt-cloud]", - "apache-airflow[dingding]", - "apache-airflow[discord]", - "apache-airflow[docker]", - "apache-airflow[elasticsearch]", - "apache-airflow[exasol]", - "apache-airflow[fab]", - "apache-airflow[facebook]", - "apache-airflow[ftp]", - "apache-airflow[github]", - "apache-airflow[google]", - "apache-airflow[grpc]", - "apache-airflow[hashicorp]", - "apache-airflow[http]", - "apache-airflow[imap]", - "apache-airflow[influxdb]", - "apache-airflow[jdbc]", - "apache-airflow[jenkins]", - "apache-airflow[microsoft-azure]", - "apache-airflow[microsoft-mssql]", - "apache-airflow[microsoft-psrp]", - "apache-airflow[microsoft-winrm]", - "apache-airflow[mongo]", - "apache-airflow[mysql]", - "apache-airflow[neo4j]", - "apache-airflow[odbc]", - "apache-airflow[openai]", - "apache-airflow[openfaas]", - "apache-airflow[openlineage]", - "apache-airflow[opensearch]", - "apache-airflow[opsgenie]", - "apache-airflow[oracle]", - "apache-airflow[pagerduty]", - "apache-airflow[papermill]", - "apache-airflow[pgvector]", - "apache-airflow[pinecone]", - "apache-airflow[postgres]", - "apache-airflow[presto]", - "apache-airflow[qdrant]", - "apache-airflow[redis]", - "apache-airflow[salesforce]", - "apache-airflow[samba]", - "apache-airflow[segment]", - "apache-airflow[sendgrid]", - "apache-airflow[sftp]", - "apache-airflow[singularity]", - "apache-airflow[slack]", - "apache-airflow[smtp]", - "apache-airflow[snowflake]", - "apache-airflow[sqlite]", - "apache-airflow[ssh]", - "apache-airflow[tableau]", - "apache-airflow[tabular]", - "apache-airflow[telegram]", - "apache-airflow[teradata]", - "apache-airflow[trino]", - "apache-airflow[vertica]", - "apache-airflow[weaviate]", - "apache-airflow[yandex]", - "apache-airflow[zendesk]", -] -devel-all = [ - "apache-airflow[all]", - "apache-airflow[devel]", - "apache-airflow[doc]", - "apache-airflow[doc-gen]", - "apache-airflow[saml]", - # Apache no provider extras - "apache-airflow[apache-atlas]", - "apache-airflow[apache-webhdfs]", - "apache-airflow[all-core]", - # Include all provider deps - "apache-airflow[airbyte]", - "apache-airflow[alibaba]", - "apache-airflow[amazon]", - "apache-airflow[apache-beam]", - "apache-airflow[apache-cassandra]", - "apache-airflow[apache-drill]", - "apache-airflow[apache-druid]", - "apache-airflow[apache-flink]", - "apache-airflow[apache-hdfs]", - "apache-airflow[apache-hive]", - "apache-airflow[apache-impala]", - "apache-airflow[apache-kafka]", - "apache-airflow[apache-kylin]", - "apache-airflow[apache-livy]", - "apache-airflow[apache-pig]", - "apache-airflow[apache-pinot]", - "apache-airflow[apache-spark]", - "apache-airflow[apprise]", - "apache-airflow[arangodb]", - "apache-airflow[asana]", - "apache-airflow[atlassian-jira]", - "apache-airflow[celery]", - "apache-airflow[cloudant]", - "apache-airflow[cncf-kubernetes]", - "apache-airflow[cohere]", - "apache-airflow[common-io]", - "apache-airflow[common-sql]", - "apache-airflow[databricks]", - "apache-airflow[datadog]", - "apache-airflow[dbt-cloud]", - "apache-airflow[dingding]", - "apache-airflow[discord]", - "apache-airflow[docker]", - "apache-airflow[elasticsearch]", - "apache-airflow[exasol]", - "apache-airflow[fab]", - "apache-airflow[facebook]", - "apache-airflow[ftp]", - "apache-airflow[github]", - "apache-airflow[google]", - "apache-airflow[grpc]", - "apache-airflow[hashicorp]", - "apache-airflow[http]", - "apache-airflow[imap]", - "apache-airflow[influxdb]", - "apache-airflow[jdbc]", - "apache-airflow[jenkins]", - "apache-airflow[microsoft-azure]", - "apache-airflow[microsoft-mssql]", - "apache-airflow[microsoft-psrp]", - "apache-airflow[microsoft-winrm]", - "apache-airflow[mongo]", - "apache-airflow[mysql]", - "apache-airflow[neo4j]", - "apache-airflow[odbc]", - "apache-airflow[openai]", - "apache-airflow[openfaas]", - "apache-airflow[openlineage]", - "apache-airflow[opensearch]", - "apache-airflow[opsgenie]", - "apache-airflow[oracle]", - "apache-airflow[pagerduty]", - "apache-airflow[papermill]", - "apache-airflow[pgvector]", - "apache-airflow[pinecone]", - "apache-airflow[postgres]", - "apache-airflow[presto]", - "apache-airflow[qdrant]", - "apache-airflow[redis]", - "apache-airflow[salesforce]", - "apache-airflow[samba]", - "apache-airflow[segment]", - "apache-airflow[sendgrid]", - "apache-airflow[sftp]", - "apache-airflow[singularity]", - "apache-airflow[slack]", - "apache-airflow[smtp]", - "apache-airflow[snowflake]", - "apache-airflow[sqlite]", - "apache-airflow[ssh]", - "apache-airflow[tableau]", - "apache-airflow[tabular]", - "apache-airflow[telegram]", - "apache-airflow[teradata]", - "apache-airflow[trino]", - "apache-airflow[vertica]", - "apache-airflow[weaviate]", - "apache-airflow[yandex]", - "apache-airflow[zendesk]", -] -# END OF GENERATED DEPENDENCIES -############################################################################################################# -# The rest of the pyproject.toml file should be manually maintained -############################################################################################################# +# aiobotocore, apache-atlas, apache-webhdfs, async, cgroups, deprecated-api, github-enterprise, +# google-auth, graphviz, kerberos, ldap, leveldb, otel, pandas, password, pydantic, rabbitmq, s3fs, +# saml, sentry, statsd, uv, virtualenv +# +# END CORE EXTRAS HERE +# +# The ``devel`` extras are not available in the released packages. They are only available when you install +# Airflow from sources in ``editable`` installation - i.e. one that you are usually using to contribute to +# Airflow. They provide tools such as ``pytest`` and ``mypy`` for general purpose development and testing. +# +# START DEVEL EXTRAS HERE +# +# devel, devel-all-dbs, devel-ci, devel-debuggers, devel-devscripts, devel-duckdb, devel-hadoop, +# devel-mypy, devel-sentry, devel-static-checks, devel-tests +# +# END DEVEL EXTRAS HERE +# +# Those extras are bundles dynamically generated from other extras. +# +# START BUNDLE EXTRAS HERE +# +# all, all-core, all-dbs, devel-all, devel-ci +# +# END BUNDLE EXTRAS HERE +# +# The ``doc`` extras are not available in the released packages. They are only available when you install +# Airflow from sources in ``editable`` installation - i.e. one that you are usually using to contribute to +# Airflow. They provide tools needed when you want to build Airflow documentation (note that you also need +# ``devel`` extras installed for airflow and providers in order to build documentation for airflow and +# provider packages respectively). The ``doc`` package is enough to build regular documentation, where +# ``doc_gen`` is needed to generate ER diagram we have describing our database. +# +# START DOC EXTRAS HERE +# +# doc, doc-gen +# +# END DOC EXTRAS HERE +# +# The `deprecated` extras are deprecated extras from Airflow 1 that will be removed in future versions. +# +# START DEPRECATED EXTRAS HERE +# +# atlas, aws, azure, cassandra, crypto, druid, gcp, gcp-api, hdfs, hive, kubernetes, mssql, pinot, s3, +# spark, webhdfs, winrm +# +# END DEPRECATED EXTRAS HERE +# +# !!!!!! Those provuders are defined in the `airflow/providers//provider.yaml` files !!!!!!! +# +# Those extras are available as regular Airflow extras, they install provider packages in standard builds +# or dependencies that are necessary to enable the feature in editable build. +# START PROVIDER EXTRAS HERE +# +# airbyte, alibaba, amazon, apache.beam, apache.cassandra, apache.drill, apache.druid, apache.flink, +# apache.hdfs, apache.hive, apache.impala, apache.kafka, apache.kylin, apache.livy, apache.pig, +# apache.pinot, apache.spark, apprise, arangodb, asana, atlassian.jira, celery, cloudant, +# cncf.kubernetes, cohere, common.io, common.sql, databricks, datadog, dbt.cloud, dingding, discord, +# docker, elasticsearch, exasol, fab, facebook, ftp, github, google, grpc, hashicorp, http, imap, +# influxdb, jdbc, jenkins, microsoft.azure, microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, +# mysql, neo4j, odbc, openai, openfaas, openlineage, opensearch, opsgenie, oracle, pagerduty, +# papermill, pgvector, pinecone, postgres, presto, qdrant, redis, salesforce, samba, segment, +# sendgrid, sftp, singularity, slack, smtp, snowflake, sqlite, ssh, tableau, tabular, telegram, +# teradata, trino, vertica, weaviate, yandex, zendesk +# +# END PROVIDER EXTRAS HERE + [project.scripts] airflow = "airflow.__main__:main" [project.urls] @@ -1224,7 +160,7 @@ YouTube = "https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/" python = "3.8" platforms = ["linux", "macos"] description = "Default environment with Python 3.8 for maximum compatibility" -features = ["devel"] +features = [] [tool.hatch.envs.airflow-38] python = "3.8" @@ -1282,7 +218,6 @@ artifacts = [ "/airflow/www/static/dist/", "/airflow/git_version", "/generated/", - "/airflow_pre_installed_providers.txt", ] @@ -1333,6 +268,7 @@ extend-select = [ "G", # flake8-logging-format rules "LOG", # flake8-logging rules, most of them autofixable "PT", # flake8-pytest-style rules + "TID25", # flake8-tidy-imports rules # Per rule enables "RUF100", # Unused noqa (auto-fixable) # We ignore more pydocstyle than we enable, so be more selective at what we enable @@ -1350,8 +286,6 @@ extend-select = [ "D403", "D412", "D419", - "TID251", # Specific modules or module members that may not be imported or accessed - "TID253", # Ban certain modules from being imported at module level "PGH004", # Use specific rule codes when using noqa "PGH005", # Invalid unittest.mock.Mock methods/attributes/properties "B006", # Checks for uses of mutable objects as function argument defaults. @@ -1371,7 +305,6 @@ ignore = [ "PT006", # Wrong type of names in @pytest.mark.parametrize "PT007", # Wrong type of values in @pytest.mark.parametrize "PT011", # pytest.raises() is too broad, set the match parameter - "PT018", # assertion should be broken down into multiple parts "PT019", # fixture without value is injected as parameter, use @pytest.mark.usefixtures instead ] unfixable = [ @@ -1403,7 +336,7 @@ combine-as-imports = true # Ignore pydoc style from these "*.pyi" = ["D"] -"scripts/*" = ["D"] +"scripts/*" = ["D", "PT"] # In addition ignore pytest specific rules "docs/*" = ["D"] "provider_packages/*" = ["D"] "*/example_dags/*" = ["D"] @@ -1520,13 +453,10 @@ combine-as-imports = true "tests/providers/cncf/kubernetes/operators/test_pod.py" = ["PT012"] "tests/providers/cncf/kubernetes/utils/test_k8s_resource_iterator.py" = ["PT012"] "tests/providers/cncf/kubernetes/utils/test_pod_manager.py" = ["PT012"] -"tests/providers/common/sql/operators/test_sql.py" = ["PT012"] "tests/providers/databricks/hooks/test_databricks.py" = ["PT012"] "tests/providers/databricks/operators/test_databricks.py" = ["PT012"] "tests/providers/databricks/operators/test_databricks_repos.py" = ["PT012"] "tests/providers/databricks/sensors/test_databricks_partition.py" = ["PT012"] -"tests/providers/datadog/sensors/test_datadog.py" = ["PT012"] -"tests/providers/dbt/cloud/operators/test_dbt.py" = ["PT012"] "tests/providers/google/cloud/hooks/test_bigquery.py" = ["PT012"] "tests/providers/google/cloud/hooks/test_cloud_storage_transfer_service.py" = ["PT012"] "tests/providers/google/cloud/hooks/test_dataflow.py" = ["PT012"] @@ -1550,21 +480,16 @@ combine-as-imports = true "tests/providers/google/cloud/transfers/test_gcs_to_bigquery.py" = ["PT012"] "tests/providers/google/cloud/utils/test_credentials_provider.py" = ["PT012"] "tests/providers/google/common/hooks/test_base_google.py" = ["PT012"] -"tests/providers/jenkins/sensors/test_jenkins.py" = ["PT012"] -"tests/providers/microsoft/azure/hooks/test_data_factory.py" = ["PT012"] -"tests/providers/microsoft/azure/hooks/test_wasb.py" = ["PT012"] -"tests/providers/microsoft/psrp/hooks/test_psrp.py" = ["PT012"] -"tests/providers/oracle/hooks/test_oracle.py" = ["PT012"] "tests/providers/sftp/hooks/test_sftp.py" = ["PT012"] "tests/providers/sftp/operators/test_sftp.py" = ["PT012"] "tests/providers/sftp/sensors/test_sftp.py" = ["PT012"] "tests/providers/sftp/triggers/test_sftp.py" = ["PT012"] "tests/providers/ssh/hooks/test_ssh.py" = ["PT012"] "tests/providers/ssh/operators/test_ssh.py" = ["PT012"] -"tests/providers/telegram/hooks/test_telegram.py" = ["PT012"] -"tests/providers/telegram/operators/test_telegram.py" = ["PT012"] [tool.ruff.lint.flake8-tidy-imports] +# Disallow all relative imports. +ban-relative-imports = "all" # Ban certain modules from being imported at module level, instead requiring # that they're imported lazily (e.g., within a function definition). banned-module-level-imports = ["numpy", "pandas"] From 4e1919576dd26e6cad44619137469757fa1cb870 Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 25 Mar 2024 10:30:15 +0100 Subject: [PATCH 046/105] refactor: Reformatted docstrings in trigger --- airflow/providers/microsoft/azure/triggers/msgraph.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index b6e4e0810c094..6f74a04a072ce 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -110,8 +110,6 @@ class MSGraphTrigger(BaseTrigger): """ A Microsoft Graph API trigger which allows you to execute an async REST call to the Microsoft Graph API. - https://github.com/microsoftgraph/msgraph-sdk-python - :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). From ce4330607582c654ac582ab5cac584ae4533298c Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 27 Mar 2024 11:18:46 +0100 Subject: [PATCH 047/105] refactor: Removed unused serialization module --- .../microsoft/azure/serialization/__init__.py | 16 ---------------- 1 file changed, 16 deletions(-) delete mode 100644 airflow/providers/microsoft/azure/serialization/__init__.py diff --git a/airflow/providers/microsoft/azure/serialization/__init__.py b/airflow/providers/microsoft/azure/serialization/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/airflow/providers/microsoft/azure/serialization/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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 8a55cd6f6e8b12d7e6918e02dff39591eb4bed2f Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 08:18:33 +0200 Subject: [PATCH 048/105] fix: Fixed execution of consecutive tasks in execute_operator method --- tests/providers/microsoft/azure/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 3cd980383a1a0..b331cbf2c4f4f 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -111,7 +111,7 @@ def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task = deferred.value while task: - events = self.run_trigger(deferred.value.trigger) + events = self.run_trigger(task.trigger) if not events: break From 40a8cce3a96857dcc44e32c8cb6b5a5e1c2f72d9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 08:40:46 +0200 Subject: [PATCH 049/105] refactor: Added customizable pagination_function parameter to Operator and made operator PowerBI compatible --- .../microsoft/azure/hooks/msgraph.py | 47 +++++++++++++++---- .../microsoft/azure/operators/msgraph.py | 33 +++++++++++-- .../microsoft/azure/triggers/msgraph.py | 15 ++++-- 3 files changed, 80 insertions(+), 15 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index 5ed148813bb65..70f944237114d 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -19,7 +19,7 @@ import json from typing import TYPE_CHECKING -from urllib.parse import urljoin +from urllib.parse import urljoin, urlparse import httpx from azure.identity import ClientSecretCredential @@ -101,12 +101,32 @@ def get_host(connection: Connection) -> str: return NationalClouds.Global.value @staticmethod - def to_httpx_proxies(proxies: dict) -> dict: + def format_no_proxy_url(url: str) -> str: + if "://" not in url: + url = f"all://{url}" + return url + + @classmethod + def to_httpx_proxies(cls, proxies: dict) -> dict: proxies = proxies.copy() if proxies.get("http"): proxies["http://"] = proxies.pop("http") if proxies.get("https"): proxies["https://"] = proxies.pop("https") + if proxies.get("no_proxy"): + for url in proxies.pop("no_proxy", "").split(","): + proxies[cls.format_no_proxy_url(url.strip())] = None + return proxies + + @classmethod + def to_msal_proxies(cls, authority: str | None, proxies: dict): + if authority: + no_proxies = proxies.get("no_proxy") + if no_proxies: + for url in no_proxies.split(","): + domain_name = urlparse(url).path.replace("*", "") + if authority.endswith(domain_name): + return None return proxies def get_conn(self) -> RequestAdapter: @@ -128,6 +148,9 @@ def get_conn(self) -> RequestAdapter: scopes = config.get("scopes", ["https://graph.microsoft.com/.default"]) verify = config.get("verify", True) trust_env = config.get("trust_env", False) + authority = config.get("authority") + disable_instance_discovery = config.get("disable_instance_discovery", False) + allowed_hosts = (config.get("allowed_hosts", authority) or "").split(",") self.log.info( "Creating Microsoft Graph SDK client %s for conn_id: %s", @@ -144,24 +167,32 @@ def get_conn(self) -> RequestAdapter: self.log.info("Verify: %s", verify) self.log.info("Timeout: %s", self.timeout) self.log.info("Trust env: %s", trust_env) + self.log.info("Authority: %s", authority) self.log.info("Proxies: %s", json.dumps(proxies)) credentials = ClientSecretCredential( - tenant_id=tenant_id, # type: ignore - client_id=connection.login, # type: ignore - client_secret=connection.password, # type: ignore - proxies=proxies, + tenant_id=tenant_id, + client_id=connection.login, + client_secret=connection.password, + authority=authority, + proxies=self.to_msal_proxies(authority=authority, proxies=proxies), + disable_instance_discovery=disable_instance_discovery, + connection_verify=verify, ) http_client = GraphClientFactory.create_with_default_middleware( api_version=api_version, client=httpx.AsyncClient( - proxies=self.to_httpx_proxies(proxies), + proxies=self.to_httpx_proxies(proxies=proxies), timeout=Timeout(timeout=self.timeout), verify=verify, trust_env=trust_env, ), host=host, ) - auth_provider = AzureIdentityAuthenticationProvider(credentials=credentials, scopes=scopes) + auth_provider = AzureIdentityAuthenticationProvider( + credentials=credentials, + scopes=scopes, + allowed_hosts=allowed_hosts, + ) request_adapter = HttpxRequestAdapter( authentication_provider=auth_provider, http_client=http_client, diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 3ac98d81953c9..c6f23c6caad33 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +from copy import deepcopy from typing import ( TYPE_CHECKING, Any, @@ -98,6 +99,7 @@ def __init__( timeout: float | None = None, proxies: dict | None = None, api_version: APIVersion | None = None, + pagination_function: Callable[[MSGraphAsyncOperator, dict], Any] | None = None, result_processor: Callable[[Context, Any], Any] = lambda context, result: result, serializer: type[ResponseSerializer] = ResponseSerializer, **kwargs: Any, @@ -117,6 +119,7 @@ def __init__( 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 @@ -237,16 +240,38 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = ) return self.execute_complete(context, event) + @staticmethod + def paginate(operator: MSGraphAsyncOperator, response: dict): + odata_count = response.get("@odata.count") + if odata_count: + query_parameters = deepcopy(operator.query_parameters) + top: int = query_parameters.get("$top") + odata_count: int = 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): - odata_next_link = response.get("@odata.nextLink") + url, query_parameters = self.pagination_function(self, response) - self.log.debug("odata_next_link: %s", odata_next_link) + self.log.debug("url: %s", url) + self.log.debug("query_parameters: %s", query_parameters) - if odata_next_link: + if url: self.defer( trigger=MSGraphTrigger( - url=odata_next_link, + url=url, + query_parameters=query_parameters, response_type=self.response_type, response_handler=self.response_handler, conn_id=self.conn_id, diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 6f74a04a072ce..ba996f08216f9 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -30,6 +30,7 @@ Callable, Sequence, ) +from urllib.parse import quote from uuid import UUID import pendulum @@ -263,15 +264,23 @@ def normalize_url(self) -> str | None: return self.url.replace("/", "", 1) return self.url + @staticmethod + def encode_query_parameters(query_parameters: dict) -> dict: + return {quote(key): quote(str(value)) for key, value in query_parameters.items()} + def request_information(self) -> RequestInformation: request_information = RequestInformation() + + request_information.path_parameters = self.path_parameters or {} + request_information.http_method = Method(self.method.strip().upper()) + request_information.query_parameters = self.encode_query_parameters(self.query_parameters or {}) if self.url.startswith("http"): request_information.url = self.url + elif request_information.query_parameters.keys(): + query = ','.join(request_information.query_parameters.keys()) + request_information.url_template = f"{{+baseurl}}/{self.normalize_url()}{{?{query}}}" else: request_information.url_template = f"{{+baseurl}}/{self.normalize_url()}" - request_information.path_parameters = self.path_parameters or {} - request_information.http_method = Method(self.method.strip().upper()) - request_information.query_parameters = self.query_parameters or {} if not self.response_type: request_information.request_options[ResponseHandlerOption.get_key()] = ResponseHandlerOption( response_handler=CallableResponseHandler(self.response_handler) From 3b097324d53eb6ca02e7dbcb46789ed03f1c318a Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 09:53:32 +0200 Subject: [PATCH 050/105] refactor: Reformatted operator and trigger --- airflow/providers/microsoft/azure/operators/msgraph.py | 3 +-- airflow/providers/microsoft/azure/triggers/msgraph.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index c6f23c6caad33..8f7c910235479 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -251,8 +251,7 @@ def paginate(operator: MSGraphAsyncOperator, response: dict): if top and odata_count: if len(response.get("value")) == top: skip = ( - sum(map(lambda result: len(result["value"]), operator.results)) - + top + sum(map(lambda result: len(result["value"]), operator.results)) + top if operator.results else top ) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index ba996f08216f9..7ef5a3e8bc663 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -277,7 +277,7 @@ def request_information(self) -> RequestInformation: if self.url.startswith("http"): request_information.url = self.url elif request_information.query_parameters.keys(): - query = ','.join(request_information.query_parameters.keys()) + query = ",".join(request_information.query_parameters.keys()) request_information.url_template = f"{{+baseurl}}/{self.normalize_url()}{{?{query}}}" else: request_information.url_template = f"{{+baseurl}}/{self.normalize_url()}" From b28253a8a6e13368e78494733b4bacccef85ff27 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 10:14:27 +0200 Subject: [PATCH 051/105] refactor: Added check if query_parameters is not None --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 8f7c910235479..776e0e28954bd 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -243,7 +243,7 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = @staticmethod def paginate(operator: MSGraphAsyncOperator, response: dict): odata_count = response.get("@odata.count") - if odata_count: + if odata_count and operator.query_parameters: query_parameters = deepcopy(operator.query_parameters) top: int = query_parameters.get("$top") odata_count: int = response.get("@odata.count") From 445f5cd5a92b5247844bcc35308dce97e4eb4b82 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 10:44:27 +0200 Subject: [PATCH 052/105] refactor: Removed typing of top and odata_count --- airflow/providers/microsoft/azure/operators/msgraph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 776e0e28954bd..c20b13e0000a9 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -245,11 +245,11 @@ def paginate(operator: MSGraphAsyncOperator, response: dict): odata_count = response.get("@odata.count") if odata_count and operator.query_parameters: query_parameters = deepcopy(operator.query_parameters) - top: int = query_parameters.get("$top") - odata_count: int = response.get("@odata.count") + top = query_parameters.get("$top") + odata_count = response.get("@odata.count") if top and odata_count: - if len(response.get("value")) == top: + if len(response.get("value", [])) == top: skip = ( sum(map(lambda result: len(result["value"]), operator.results)) + top if operator.results From dcd322c12f6d1730a4b9ccbab08673bef2d3a591 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 10:46:19 +0200 Subject: [PATCH 053/105] refactor: Ignore type for tenant_id (this is an issue in the ClientSecretCredential class) --- airflow/providers/microsoft/azure/hooks/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index 70f944237114d..cdfc1f0866265 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -170,7 +170,7 @@ def get_conn(self) -> RequestAdapter: self.log.info("Authority: %s", authority) self.log.info("Proxies: %s", json.dumps(proxies)) credentials = ClientSecretCredential( - tenant_id=tenant_id, + tenant_id=tenant_id, # type: ignore client_id=connection.login, client_secret=connection.password, authority=authority, From 8fca1623b7245d13d0da206028b88dd05c4a279b Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 12:43:49 +0200 Subject: [PATCH 054/105] refactor: Changed docstring on MSGraphTrigger --- airflow/providers/microsoft/azure/triggers/msgraph.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 7ef5a3e8bc663..7ae381eda76af 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -112,19 +112,19 @@ 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 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is lambda response, error_map: response.json()). The default expression will convert the native response to JSON. If response_type parameter is specified, then the response_handler will be ignored. :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 no HTTP timeout is applied on each request. + :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". + 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. """ From f3fc9cafe75b7808b7c45a6b847a0279cf5d70a5 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 12:46:05 +0200 Subject: [PATCH 055/105] refactor: Changed docstring on MSGraphTrigger --- .../providers/microsoft/azure/operators/msgraph.py | 14 +++++++------- .../providers/microsoft/azure/triggers/msgraph.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index c20b13e0000a9..81ee67f2c6de5 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -57,20 +57,20 @@ class MSGraphAsyncOperator(BaseOperator): :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 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is lambda response, error_map: response.json()). The default expression will convert the native response to JSON. If response_type parameter is specified, then the response_handler will be ignored. :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 no HTTP timeout is applied on each request. - :param proxies: A Dict defining the HTTP proxies to be used (default is None). + :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". + 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 GraphServiceClientHook are bytes, then those will be base64 encoded into a string. diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 7ae381eda76af..3c349b9e16c03 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -121,7 +121,7 @@ class MSGraphTrigger(BaseTrigger): :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 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`. From 568dffc25f9f0b79e73a8354458fbb32525b2dd1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 17:23:43 +0200 Subject: [PATCH 056/105] refactor: Added docstring to handle_response_async method --- airflow/providers/microsoft/azure/triggers/msgraph.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 3c349b9e16c03..6629503e02d28 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -104,6 +104,12 @@ def __init__( async def handle_response_async( self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None ) -> Any: + """ + Callback method that is invoked when a response is received. + + param response: The type of the native response object. + param error_map: The error dict to use in case of a failed request. + """ return self.callable_function(response, error_map) From a6648670820745c013dc78ed518ce8eda743a31f Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 3 Apr 2024 18:08:28 +0200 Subject: [PATCH 057/105] refactor: Fixed docstring to imperative for handle_response_async method --- airflow/providers/microsoft/azure/triggers/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 6629503e02d28..808da2fdc3928 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -105,7 +105,7 @@ async def handle_response_async( self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None ) -> Any: """ - Callback method that is invoked when a response is received. + Invoke this callback method when a response is received. param response: The type of the native response object. param error_map: The error dict to use in case of a failed request. From 2f0018cdf440d9ec16d1ef9138af7291f15e0f1e Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 10:35:21 +0200 Subject: [PATCH 058/105] refactor: Try quoting Sharepoint so it doesn't get spell checked --- .../operators/msgraph.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst index 8ef239924acd0..f6b8d5173070d 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst @@ -32,7 +32,7 @@ Use the :class:`~airflow.providers.microsoft.azure.operators.msgraph.MSGraphAsyncOperator` to call Microsoft Graph API. -Below is an example of using this operator to get a Sharepoint site. +Below is an example of using this operator to get a `Sharepoint` site. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python @@ -40,7 +40,7 @@ Below is an example of using this operator to get a Sharepoint site. :start-after: [START howto_operator_graph_site] :end-before: [END howto_operator_graph_site] -Below is an example of using this operator to get a Sharepoint site pages. +Below is an example of using this operator to get a `Sharepoint` site pages. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python From 6cd7f5e02da0e9b44b6ba1b18ac4eed9bf5981a6 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 11:06:06 +0200 Subject: [PATCH 059/105] refactor: Try double quoting Sharepoint so it doesn't get spell checked --- .../operators/msgraph.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst index f6b8d5173070d..f8d87556affed 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst @@ -32,7 +32,7 @@ Use the :class:`~airflow.providers.microsoft.azure.operators.msgraph.MSGraphAsyncOperator` to call Microsoft Graph API. -Below is an example of using this operator to get a `Sharepoint` site. +Below is an example of using this operator to get a "Sharepoint" site. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python @@ -40,7 +40,7 @@ Below is an example of using this operator to get a `Sharepoint` site. :start-after: [START howto_operator_graph_site] :end-before: [END howto_operator_graph_site] -Below is an example of using this operator to get a `Sharepoint` site pages. +Below is an example of using this operator to get a "Sharepoint" site pages. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python From 0551bd0921e75b697c06362f31a69c984a6797fc Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 12:01:54 +0200 Subject: [PATCH 060/105] refactor: Always get a new event loop and close it after test is done --- tests/providers/microsoft/azure/base.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index b331cbf2c4f4f..d22de934165e5 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -17,7 +17,7 @@ from __future__ import annotations import asyncio -from contextlib import contextmanager +from contextlib import contextmanager, closing from copy import deepcopy from datetime import datetime from typing import TYPE_CHECKING, Any, Iterable @@ -69,8 +69,6 @@ def xcom_push( class Base: - _loop = asyncio.get_event_loop() - def teardown_method(self, method): KiotaRequestAdapterHook.cached_request_adapters.clear() MockedTaskInstance.values.clear() @@ -97,7 +95,11 @@ def run_trigger(self, trigger: BaseTrigger) -> list[TriggerEvent]: return self.run_async(self._run_tigger(trigger)) def run_async(self, future: Any) -> Any: - return self._loop.run_until_complete(future) + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + with closing(loop): + return loop.run_until_complete(future) def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) From 45ea05dcca2381af5fc0363a1767f0205d49e3c1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 13:21:07 +0200 Subject: [PATCH 061/105] refactor: Reordered imports from contextlib --- tests/providers/microsoft/azure/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index d22de934165e5..3ad48f99ec612 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -17,7 +17,7 @@ from __future__ import annotations import asyncio -from contextlib import contextmanager, closing +from contextlib import closing, contextmanager from copy import deepcopy from datetime import datetime from typing import TYPE_CHECKING, Any, Iterable From a2ba077da8d8311b743a4eeb7aed7efd0ab29a8d Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 14:32:42 +0200 Subject: [PATCH 062/105] refactor: Added Sharepoint to spelling_wordlist.txt --- .../operators/msgraph.rst | 4 ++-- docs/spelling_wordlist.txt | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst index f8d87556affed..8ef239924acd0 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst @@ -32,7 +32,7 @@ Use the :class:`~airflow.providers.microsoft.azure.operators.msgraph.MSGraphAsyncOperator` to call Microsoft Graph API. -Below is an example of using this operator to get a "Sharepoint" site. +Below is an example of using this operator to get a Sharepoint site. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python @@ -40,7 +40,7 @@ Below is an example of using this operator to get a "Sharepoint" site. :start-after: [START howto_operator_graph_site] :end-before: [END howto_operator_graph_site] -Below is an example of using this operator to get a "Sharepoint" site pages. +Below is an example of using this operator to get a Sharepoint site pages. .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_msgraph.py :language: python diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index ebff77ea37115..29191acd4bd00 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1439,6 +1439,7 @@ setted sftp SFTPClient sharded +Sharepoint shellcheck shellcmd shm From ae8a195a77cf928c4fcb1b757711cb4d134784b3 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 15:25:50 +0200 Subject: [PATCH 063/105] refactor: Removed connection-type for KiotaRequestAdapterHook --- airflow/providers/microsoft/azure/provider.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 2556d73e9b5d9..dec96d9a8915d 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -324,8 +324,6 @@ connection-types: connection-type: adls - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapsePipelineHook connection-type: azure_synapse_pipeline - - hook-class-name: airflow.providers.microsoft.azure.hooks.msgraph.KiotaRequestAdapterHook - connection-type: http secrets-backends: - airflow.providers.microsoft.azure.secrets.key_vault.AzureKeyVaultBackend From 4010b38f05d471b597e57e51d2a89b347ec5d5cc Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 16:53:34 +0200 Subject: [PATCH 064/105] refactor: Refactored encoded_query_parameters --- airflow/providers/microsoft/azure/triggers/msgraph.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 808da2fdc3928..148b688d75e92 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -270,16 +270,16 @@ def normalize_url(self) -> str | None: return self.url.replace("/", "", 1) return self.url - @staticmethod - def encode_query_parameters(query_parameters: dict) -> dict: - return {quote(key): quote(str(value)) for key, value in query_parameters.items()} + def encoded_query_parameters(self) -> dict: + if self.query_parameters: + return {quote(key): quote(str(value)) for key, value in self.query_parameters.items()} + return {} def request_information(self) -> RequestInformation: request_information = RequestInformation() - request_information.path_parameters = self.path_parameters or {} request_information.http_method = Method(self.method.strip().upper()) - request_information.query_parameters = self.encode_query_parameters(self.query_parameters or {}) + request_information.query_parameters = self.encoded_query_parameters() if self.url.startswith("http"): request_information.url = self.url elif request_information.query_parameters.keys(): From 09e478cb2d6b3908ad40d20641525be8cdd68a1d Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 16:55:16 +0200 Subject: [PATCH 065/105] refactor: Suppress ImportError --- airflow/providers/microsoft/azure/triggers/msgraph.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 148b688d75e92..57efcef8719ba 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -187,10 +187,9 @@ def __init__( @classmethod def resolve_type(cls, value: str | type, default) -> type: if isinstance(value, str): - try: + with suppress(ImportError): return import_string(value) - except ImportError: - return default + return default return value or default def serialize(self) -> tuple[str, dict[str, Any]]: From 180b52ed25a10f3a5587b2dca69228369a44abe8 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 16:56:45 +0200 Subject: [PATCH 066/105] refactor: Added return type to paginate method --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 81ee67f2c6de5..4b82e6865e976 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -241,7 +241,7 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = return self.execute_complete(context, event) @staticmethod - def paginate(operator: MSGraphAsyncOperator, response: dict): + def paginate(operator: MSGraphAsyncOperator, response: dict) -> tuple[str, dict]: odata_count = response.get("@odata.count") if odata_count and operator.query_parameters: query_parameters = deepcopy(operator.query_parameters) From 3d3b65b81d7630fe103b250ff839385b7de90a12 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 17:00:48 +0200 Subject: [PATCH 067/105] refactor: Updated paging_function type in MSGraphAsyncOperator --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 4b82e6865e976..a7e4eae07bbc2 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -99,7 +99,7 @@ def __init__( timeout: float | None = None, proxies: dict | None = None, api_version: APIVersion | None = None, - pagination_function: Callable[[MSGraphAsyncOperator, dict], Any] | 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, From 68814d0ece09e2130cf30422f4a73da2768cc91f Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 19:09:16 +0200 Subject: [PATCH 068/105] refactor: Pass the method name from method reference instead of hard coded string which is re-factor friendly --- airflow/providers/microsoft/azure/operators/msgraph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index a7e4eae07bbc2..ff89045950a97 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -142,7 +142,7 @@ def execute(self, context: Context) -> None: api_version=self.api_version, serializer=type(self.serializer), ), - method_name="execute_complete", + method_name=self.execute_complete.__name__, ) def execute_complete( @@ -179,7 +179,7 @@ def execute_complete( event["response"] = result try: - self.trigger_next_link(response, method_name="pull_execute_complete") + self.trigger_next_link(response, method_name=self.pull_execute_complete.__name__) except TaskDeferred as exception: self.append_result( result=result, From 9f12dcac0541b17c4ff593107bdeba22aadff31e Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 4 Apr 2024 19:10:33 +0200 Subject: [PATCH 069/105] refactor: Changed return type of paginate method --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index ff89045950a97..094fb76cb8bf1 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -241,7 +241,7 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = return self.execute_complete(context, event) @staticmethod - def paginate(operator: MSGraphAsyncOperator, response: dict) -> tuple[str, dict]: + 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) From aa50a99c6dc296be5cab2a128ae6fc1ea950fd18 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 10 Apr 2024 17:36:41 +0200 Subject: [PATCH 070/105] refactor: Added MSGraphSensor which easily allows us to poll PowerBI statuses --- .../microsoft/azure/operators/msgraph.py | 6 +- .../microsoft/azure/sensors/msgraph.py | 117 ++++++++++++++++++ .../microsoft/azure/triggers/msgraph.py | 17 ++- .../microsoft/azure/resources/status.json | 1 + .../microsoft/azure/sensors/test_msgraph.py | 47 +++++++ 5 files changed, 179 insertions(+), 9 deletions(-) create mode 100644 airflow/providers/microsoft/azure/sensors/msgraph.py create mode 100644 tests/providers/microsoft/azure/resources/status.json create mode 100644 tests/providers/microsoft/azure/sensors/test_msgraph.py diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 094fb76cb8bf1..9353048744141 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -93,7 +93,7 @@ def __init__( method: str = "GET", query_parameters: dict[str, QueryParams] | None = None, headers: dict[str, str] | None = None, - content: BytesIO | 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, @@ -113,7 +113,7 @@ def __init__( self.method = method self.query_parameters = query_parameters self.headers = headers - self.content = content + self.data = data self.conn_id = conn_id self.key = key self.timeout = timeout @@ -135,7 +135,7 @@ def execute(self, context: Context) -> None: method=self.method, query_parameters=self.query_parameters, headers=self.headers, - content=self.content, + data=self.data, conn_id=self.conn_id, timeout=self.timeout, proxies=self.proxies, diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py new file mode 100644 index 0000000000000..8f81d264161a3 --- /dev/null +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -0,0 +1,117 @@ +# +# 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 +import json +from io import BytesIO +from typing import TYPE_CHECKING, Any, Callable + +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, PokeReturnValue +from airflow.utils.context import Context + +if TYPE_CHECKING: + from airflow.triggers.base import TriggerEvent + + from kiota_abstractions.request_information import QueryParams + from kiota_abstractions.response_handler import NativeResponseType + from kiota_abstractions.serialization import ParsableFactory + from kiota_http.httpx_request_adapter import ResponseType + from msgraph_core import APIVersion + + +def default_event_processor(context: Context, event: TriggerEvent) -> bool: + if event.payload["status"] == "success": + return json.loads(event.payload["response"])["status"] == "Succeeded" + return False + + +class MSGraphSensor(BaseSensorOperator): + def __init__( + self, + url: str, + response_type: ResponseType | None = None, + response_handler: Callable[ + [NativeResponseType, dict[str, ParsableFactory | None] | None], Any + ] = lambda response, error_map: response.json(), + 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, + event_processor: Callable[[Context, TriggerEvent], bool] = default_event_processor, + **kwargs, + ): + super().__init__(**kwargs) + self.url = url + self.response_type = response_type + self.response_handler = response_handler + 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.timeout = timeout + self.proxies = proxies + self.api_version = api_version + self.serializer = serializer + self.event_processor = event_processor + + @property + def trigger(self): + return MSGraphTrigger( + url=self.url, + response_type=self.response_type, + response_handler=self.response_handler, + 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=self.serializer, + ) + + async def async_poke(self, context: Context) -> bool | PokeReturnValue: + self.log.info("Sensor triggered") + + async for response in self.trigger.run(): + self.log.debug("response: %s", response) + + is_done = self.event_processor(context, response) + + self.log.debug("is_done: %s", is_done) + + return PokeReturnValue(is_done=is_done, xcom_value=response) + + def poke(self, context) -> bool | PokeReturnValue: + return asyncio.run(self.async_poke(context)) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 57efcef8719ba..e48e4d7e9cd2a 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -22,6 +22,7 @@ from base64 import b64encode from contextlib import suppress from datetime import datetime +from io import BytesIO from json import JSONDecodeError from typing import ( TYPE_CHECKING, @@ -45,8 +46,6 @@ 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_abstractions.response_handler import NativeResponseType @@ -159,7 +158,7 @@ def __init__( method: str = "GET", query_parameters: dict[str, QueryParams] | None = None, headers: dict[str, str] | None = None, - content: BytesIO | 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, @@ -181,7 +180,7 @@ def __init__( self.method = method self.query_parameters = query_parameters self.headers = headers - self.content = content + self.data = data self.serializer: ResponseSerializer = self.resolve_type(serializer, default=ResponseSerializer)() @classmethod @@ -209,7 +208,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "method": self.method, "query_parameters": self.query_parameters, "headers": self.headers, - "content": self.content, + "data": self.data, "response_type": self.response_type, }, ) @@ -290,10 +289,16 @@ def request_information(self) -> RequestInformation: request_information.request_options[ResponseHandlerOption.get_key()] = ResponseHandlerOption( response_handler=CallableResponseHandler(self.response_handler) ) - request_information.content = self.content headers = {**self.DEFAULT_HEADERS, **self.headers} if self.headers else self.DEFAULT_HEADERS for header_name, header_value in headers.items(): request_information.headers.try_add(header_name=header_name, header_value=header_value) + if isinstance(self.data, BytesIO) or isinstance(self.data, bytes) or isinstance(self.data, str): + request_information.content = self.data + elif self.data: + request_information.headers.try_add( + header_name=RequestInformation.CONTENT_TYPE_HEADER, header_value="application/json" + ) + request_information.content = json.dumps(self.data).encode("utf-8") return request_information @staticmethod diff --git a/tests/providers/microsoft/azure/resources/status.json b/tests/providers/microsoft/azure/resources/status.json new file mode 100644 index 0000000000000..ab9037dfc6d4a --- /dev/null +++ b/tests/providers/microsoft/azure/resources/status.json @@ -0,0 +1 @@ +{"id": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", "createdDateTime": "2024-04-10T15:05:17.357", "status": "Succeeded"} \ No newline at end of file diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py new file mode 100644 index 0000000000000..2b9c86dc9a9d2 --- /dev/null +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -0,0 +1,47 @@ +# 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 airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor +from airflow.triggers.base import TriggerEvent +from airflow.utils.state import TaskInstanceState +from tests.providers.microsoft.azure.base import Base, MockedTaskInstance +from tests.providers.microsoft.conftest import load_json, mock_json_response + + +class TestMSGraphSensor(Base): + def test_execute(self): + status = load_json("resources", "status.json") + response = mock_json_response(200, status) + + with self.patch_hook_and_request_adapter(response): + sensor = MSGraphSensor( + task_id="check_workspaces_status", + conn_id="powerbi", + url="myorg/admin/workspaces/scanStatus/0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", + timeout=350.0, + ) + actual = sensor.execute( + context=MockedTaskInstance(task=sensor, run_id="run_id", state=TaskInstanceState.RUNNING) + ) + + assert isinstance(actual, TriggerEvent) + assert actual.payload["status"] == "success" + assert actual.payload["type"] == "builtins.dict" + assert actual.payload["response"] == json.dumps(status) From c0126147485c79ba7a02850ab70e50122a0eb9c2 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 10 Apr 2024 17:58:30 +0200 Subject: [PATCH 071/105] refactor: Moved BytesIO and Context to type checking block for MSGraphSensor --- airflow/providers/microsoft/azure/sensors/msgraph.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 8f81d264161a3..988cae38f491d 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -19,7 +19,6 @@ import asyncio import json -from io import BytesIO from typing import TYPE_CHECKING, Any, Callable from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook @@ -28,7 +27,9 @@ from airflow.utils.context import Context if TYPE_CHECKING: + from io import BytesIO from airflow.triggers.base import TriggerEvent + from airflow.utils.context import Context from kiota_abstractions.request_information import QueryParams from kiota_abstractions.response_handler import NativeResponseType From ca6f92cae94edeed190df1c2c807324e510bbae3 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 10 Apr 2024 17:59:44 +0200 Subject: [PATCH 072/105] refactor: Added noqa check on pull_execute_complete method of MSGraphOperator --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 9353048744141..1d433e5017d3b 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -229,7 +229,7 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = dag_id=self.dag_id, key=self.key, ) - or [] + or [] # noqa: W503 ) self.log.info( "Pulled XCom with task_id '%s' and dag_id '%s' and key '%s': %s", From bf08a937c07c6ffbe08fe4712c8968ea6f1a183c Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 10 Apr 2024 18:45:09 +0200 Subject: [PATCH 073/105] fix: Fixed test_serialize of TestMSGraphTrigger --- tests/providers/microsoft/azure/triggers/test_msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index fc5bc9bb0e685..b155f0b7aabff 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -116,7 +116,7 @@ def test_serialize(self): "method": "GET", "query_parameters": None, "headers": None, - "content": None, + "data": None, "response_type": "bytes", "conn_id": "msgraph_api", "timeout": None, From 48c43486044e2f5191661573aeed915b9ccabb73 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 08:26:12 +0200 Subject: [PATCH 074/105] refactor: Added docstring to MSGraphSensor and updated the docstring of the MSGraphAsyncOperator --- .../microsoft/azure/operators/msgraph.py | 2 +- .../microsoft/azure/sensors/msgraph.py | 32 +++++++++++++++++-- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 1d433e5017d3b..95dc0c1f296ba 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -73,7 +73,7 @@ class MSGraphAsyncOperator(BaseOperator): 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 - GraphServiceClientHook are bytes, then those will be base64 encoded into a string. + `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. """ diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 988cae38f491d..86c685298e90a 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -19,7 +19,7 @@ import asyncio import json -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any, Callable, Sequence from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger, ResponseSerializer @@ -45,6 +45,32 @@ def default_event_processor(context: Context, event: TriggerEvent) -> bool: 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is + lambda response, error_map: response.json()). The default expression will convert the native response + to JSON. If response_type parameter is specified, then the response_handler will be ignored. + :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 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 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", "conn_id") + def __init__( self, url: str, @@ -62,8 +88,8 @@ def __init__( timeout: float | None = None, proxies: dict | None = None, api_version: APIVersion | None = None, - serializer: type[ResponseSerializer] = ResponseSerializer, event_processor: Callable[[Context, TriggerEvent], bool] = default_event_processor, + serializer: type[ResponseSerializer] = ResponseSerializer, **kwargs, ): super().__init__(**kwargs) @@ -80,8 +106,8 @@ def __init__( self.timeout = timeout self.proxies = proxies self.api_version = api_version - self.serializer = serializer self.event_processor = event_processor + self.serializer = serializer @property def trigger(self): From edf1ecd95bb715fe2bf5e5263aed0d510a6f39d4 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 08:53:00 +0200 Subject: [PATCH 075/105] refactor: Reformatted docstring of MSGraphSensor --- .../microsoft/azure/sensors/msgraph.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 86c685298e90a..cf13decdd7027 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -46,28 +46,28 @@ def default_event_processor(context: Context, event: TriggerEvent) -> bool: 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is - lambda response, error_map: response.json()). The default expression will convert the native response - to JSON. If response_type parameter is specified, then the response_handler will be ignored. - :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 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 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. - """ + 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 response_handler: Function to convert the native HTTPX response returned by the hook (default is + lambda response, error_map: response.json()). The default expression will convert the native response + to JSON. If response_type parameter is specified, then the response_handler will be ignored. + :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 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 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", "conn_id") From 0d955f5dfda17e36dac380f307e0bd05ab44a1e1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:00:07 +0200 Subject: [PATCH 076/105] refactor: Added white line at end of status.json file to keep static check happy --- tests/providers/microsoft/azure/resources/status.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/resources/status.json b/tests/providers/microsoft/azure/resources/status.json index ab9037dfc6d4a..6bff9e29afb41 100644 --- a/tests/providers/microsoft/azure/resources/status.json +++ b/tests/providers/microsoft/azure/resources/status.json @@ -1 +1 @@ -{"id": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", "createdDateTime": "2024-04-10T15:05:17.357", "status": "Succeeded"} \ No newline at end of file +{"id": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", "createdDateTime": "2024-04-10T15:05:17.357", "status": "Succeeded"} From c42133c28aa0117ea895acc37892ab8e19b0cc3c Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:04:14 +0200 Subject: [PATCH 077/105] refactor: Removed timeout parameter from constructor MSGraphSensor as it is already defined in the BaseSensorOperator --- airflow/providers/microsoft/azure/sensors/msgraph.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index cf13decdd7027..6846de2b34064 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -56,8 +56,6 @@ class MSGraphSensor(BaseSensorOperator): to JSON. If response_type parameter is specified, then the response_handler will be ignored. :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, @@ -85,7 +83,6 @@ def __init__( 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, event_processor: Callable[[Context, TriggerEvent], bool] = default_event_processor, @@ -103,7 +100,6 @@ def __init__( self.headers = headers self.data = data self.conn_id = conn_id - self.timeout = timeout self.proxies = proxies self.api_version = api_version self.event_processor = event_processor From 229bb600ba58d97069194e614d060b480190a199 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:06:15 +0200 Subject: [PATCH 078/105] fix: Added missing return for async_poke in MSGraphSensor --- airflow/providers/microsoft/azure/sensors/msgraph.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 6846de2b34064..92834c87f2ce9 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -135,6 +135,7 @@ async def async_poke(self, context: Context) -> bool | PokeReturnValue: self.log.debug("is_done: %s", is_done) return PokeReturnValue(is_done=is_done, xcom_value=response) + return PokeReturnValue(is_done=True) def poke(self, context) -> bool | PokeReturnValue: return asyncio.run(self.async_poke(context)) From f7045933f97c299c1f85bec3479ffa7a86a5896e Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:33:58 +0200 Subject: [PATCH 079/105] Revert "refactor: Added noqa check on pull_execute_complete method of MSGraphOperator" This reverts commit ca6f92cae94edeed190df1c2c807324e510bbae3. --- airflow/providers/microsoft/azure/operators/msgraph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 95dc0c1f296ba..59e5b809d09e8 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -229,7 +229,7 @@ def pull_execute_complete(self, context: Context, event: dict[Any, Any] | None = dag_id=self.dag_id, key=self.key, ) - or [] # noqa: W503 + or [] ) self.log.info( "Pulled XCom with task_id '%s' and dag_id '%s' and key '%s': %s", From e1e2ed4baf67faf2a5151325014471bf8a4354bf Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:36:12 +0200 Subject: [PATCH 080/105] refactor: Reorganised imports on MSGraphSensor --- airflow/providers/microsoft/azure/sensors/msgraph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 92834c87f2ce9..2836333bba375 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -24,12 +24,9 @@ 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, PokeReturnValue -from airflow.utils.context import Context if TYPE_CHECKING: from io import BytesIO - from airflow.triggers.base import TriggerEvent - from airflow.utils.context import Context from kiota_abstractions.request_information import QueryParams from kiota_abstractions.response_handler import NativeResponseType @@ -37,6 +34,9 @@ from kiota_http.httpx_request_adapter import ResponseType from msgraph_core import APIVersion + from airflow.triggers.base import TriggerEvent + from airflow.utils.context import Context + def default_event_processor(context: Context, event: TriggerEvent) -> bool: if event.payload["status"] == "success": From b520b95f572914c216e721fc9a1f42400588cbb7 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:36:54 +0200 Subject: [PATCH 081/105] refactor: Reformatted TestMSGraphSensor --- .../providers/microsoft/azure/sensors/test_msgraph.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index 2b9c86dc9a9d2..4ea855b5a699e 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -32,11 +32,11 @@ def test_execute(self): with self.patch_hook_and_request_adapter(response): sensor = MSGraphSensor( - task_id="check_workspaces_status", - conn_id="powerbi", - url="myorg/admin/workspaces/scanStatus/0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", - timeout=350.0, - ) + task_id="check_workspaces_status", + conn_id="powerbi", + url="myorg/admin/workspaces/scanStatus/0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", + timeout=350.0, + ) actual = sensor.execute( context=MockedTaskInstance(task=sensor, run_id="run_id", state=TaskInstanceState.RUNNING) ) From 5b15e5e13859d34687fc0c9d3768789e51672c60 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 09:40:01 +0200 Subject: [PATCH 082/105] refactor: Added MSGraph sensor integration name in provider.yaml --- airflow/providers/microsoft/azure/provider.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 070f55cf04516..5318af154b24f 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -214,6 +214,9 @@ sensors: - integration-name: Microsoft Azure Data Factory python-modules: - airflow.providers.microsoft.azure.sensors.data_factory + - integration-name: Microsoft Graph API + python-modules: + - airflow.providers.microsoft.azure.sensors.msgraph filesystems: - airflow.providers.microsoft.azure.fs.adls From c939111adecdcc92f3a681f3c9df728ad6823620 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 12:10:41 +0200 Subject: [PATCH 083/105] refactor: Updated apache-airflow version to at least 2.7.0 in provider.yaml of microsoft-azure provider --- airflow/providers/microsoft/azure/provider.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 5318af154b24f..f1fa058f86618 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -76,7 +76,7 @@ versions: - 1.0.0 dependencies: - - apache-airflow>=2.6.0 + - apache-airflow>=2.7.0 - adlfs>=2023.10.0 - azure-batch>=8.0.0 - azure-cosmos>=4.6.0 From 5a1bdfdee17a1f8d7d7d5cd11d0f1a6e7db528ed Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 12:20:17 +0200 Subject: [PATCH 084/105] refactor: Exclude microsoft-azure from compatibility check with airflow 2.6.0 as version 2.7.0 will at least be required --- dev/breeze/src/airflow_breeze/global_constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index e17c60ad0df9f..9eea6f4dc7083 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -471,7 +471,7 @@ def _exclusion(providers: Iterable[str]) -> str: { "python-version": "3.8", "airflow-version": "2.6.0", - "remove-providers": _exclusion(["openlineage", "common.io", "cohere", "fab", "qdrant"]), + "remove-providers": _exclusion(["openlineage", "common.io", "cohere", "fab", "qdrant", "microsoft-azure"]), }, { "python-version": "3.8", From 21232c8e4ce433e0136d1b1d293260cb5fe8f451 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 13:15:28 +0200 Subject: [PATCH 085/105] refactor: Also updated the apache-airflow dependency version from 2.6.0 to 2.7.0 for microsoft-azure provider in provider_dependencies.json --- generated/provider_dependencies.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 83b95c76bd00d..841f3764674e3 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -680,7 +680,7 @@ "deps": [ "adal>=1.2.7", "adlfs>=2023.10.0", - "apache-airflow>=2.6.0", + "apache-airflow>=2.7.0", "azure-batch>=8.0.0", "azure-cosmos>=4.6.0", "azure-datalake-store>=0.0.45", From 732a13bf6b6268e0bfb36f3a0406f6279e831741 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 14:07:37 +0200 Subject: [PATCH 086/105] refactor: Reformatted global_constants.py --- dev/breeze/src/airflow_breeze/global_constants.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 9eea6f4dc7083..316ff4049087b 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -471,7 +471,9 @@ def _exclusion(providers: Iterable[str]) -> str: { "python-version": "3.8", "airflow-version": "2.6.0", - "remove-providers": _exclusion(["openlineage", "common.io", "cohere", "fab", "qdrant", "microsoft-azure"]), + "remove-providers": _exclusion( + ["openlineage", "common.io", "cohere", "fab", "qdrant", "microsoft-azure"] + ), }, { "python-version": "3.8", From 7c8064aa8ded0e4047552fc458c5ac42ce3ffc42 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 15:56:05 +0200 Subject: [PATCH 087/105] refactor: Add logging statements for proxies and authority related stuff --- .../microsoft/azure/hooks/msgraph.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/msgraph.py b/airflow/providers/microsoft/azure/hooks/msgraph.py index cdfc1f0866265..7fcc328f8670a 100644 --- a/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import json from typing import TYPE_CHECKING from urllib.parse import urljoin, urlparse @@ -113,15 +112,15 @@ def to_httpx_proxies(cls, proxies: dict) -> dict: proxies["http://"] = proxies.pop("http") if proxies.get("https"): proxies["https://"] = proxies.pop("https") - if proxies.get("no_proxy"): - for url in proxies.pop("no_proxy", "").split(","): + if proxies.get("no"): + for url in proxies.pop("no", "").split(","): proxies[cls.format_no_proxy_url(url.strip())] = None return proxies @classmethod def to_msal_proxies(cls, authority: str | None, proxies: dict): if authority: - no_proxies = proxies.get("no_proxy") + no_proxies = proxies.get("no") if no_proxies: for url in no_proxies.split(","): domain_name = urlparse(url).path.replace("*", "") @@ -144,11 +143,13 @@ def get_conn(self) -> RequestAdapter: api_version = self.get_api_version(config) host = self.get_host(connection) base_url = config.get("base_url", urljoin(host, api_version.value)) + authority = config.get("authority") proxies = self.proxies or config.get("proxies", {}) + msal_proxies = self.to_msal_proxies(authority=authority, proxies=proxies) + httpx_proxies = self.to_httpx_proxies(proxies=proxies) scopes = config.get("scopes", ["https://graph.microsoft.com/.default"]) verify = config.get("verify", True) trust_env = config.get("trust_env", False) - authority = config.get("authority") disable_instance_discovery = config.get("disable_instance_discovery", False) allowed_hosts = (config.get("allowed_hosts", authority) or "").split(",") @@ -168,20 +169,24 @@ def get_conn(self) -> RequestAdapter: self.log.info("Timeout: %s", self.timeout) self.log.info("Trust env: %s", trust_env) self.log.info("Authority: %s", authority) - self.log.info("Proxies: %s", json.dumps(proxies)) + self.log.info("Disable instance discovery: %s", disable_instance_discovery) + self.log.info("Allowed hosts: %s", allowed_hosts) + self.log.info("Proxies: %s", proxies) + self.log.info("MSAL Proxies: %s", msal_proxies) + self.log.info("HTTPX Proxies: %s", httpx_proxies) credentials = ClientSecretCredential( tenant_id=tenant_id, # type: ignore client_id=connection.login, client_secret=connection.password, authority=authority, - proxies=self.to_msal_proxies(authority=authority, proxies=proxies), + proxies=msal_proxies, disable_instance_discovery=disable_instance_discovery, connection_verify=verify, ) http_client = GraphClientFactory.create_with_default_middleware( api_version=api_version, client=httpx.AsyncClient( - proxies=self.to_httpx_proxies(proxies=proxies), + proxies=httpx_proxies, timeout=Timeout(timeout=self.timeout), verify=verify, trust_env=trust_env, From 35915efc66ae938506f5776df6e7fc3d488c0527 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 16:52:12 +0200 Subject: [PATCH 088/105] fix: Fixed exclusion of microsoft.azure dependency in global_constants.py --- dev/breeze/src/airflow_breeze/global_constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 316ff4049087b..6abbb2dbc3661 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -472,7 +472,7 @@ def _exclusion(providers: Iterable[str]) -> str: "python-version": "3.8", "airflow-version": "2.6.0", "remove-providers": _exclusion( - ["openlineage", "common.io", "cohere", "fab", "qdrant", "microsoft-azure"] + ["openlineage", "common.io", "cohere", "fab", "qdrant", "microsoft.azure"] ), }, { From 2767828323efda1b3a0bc0f4af76e157f4e6cf04 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 20:17:20 +0200 Subject: [PATCH 089/105] refactor: Some Azure related imports should be ignored when running Airflow 2.6.0 or lower --- .../providers/amazon/aws/transfers/azure_blob_to_s3.py | 8 +++++++- airflow/providers/google/cloud/transfers/adls_to_gcs.py | 8 +++++++- .../providers/google/cloud/transfers/azure_blob_to_gcs.py | 8 +++++++- .../google/cloud/transfers/azure_fileshare_to_gcs.py | 8 +++++++- 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py b/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py index d62931a76510d..40b6447fb7d2e 100644 --- a/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py @@ -23,7 +23,6 @@ from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.s3 import S3Hook -from airflow.providers.microsoft.azure.hooks.wasb import WasbHook if TYPE_CHECKING: from airflow.utils.context import Context @@ -111,6 +110,13 @@ def __init__( def execute(self, context: Context) -> list[str]: # list all files in the Azure Blob Storage container + try: + from airflow.providers.microsoft.azure.hooks.wasb import WasbHook + except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + wasb_hook = WasbHook(wasb_conn_id=self.wasb_conn_id, **self.wasb_extra_args) s3_hook = S3Hook( aws_conn_id=self.aws_conn_id, diff --git a/airflow/providers/google/cloud/transfers/adls_to_gcs.py b/airflow/providers/google/cloud/transfers/adls_to_gcs.py index 7abbd9a9c3142..4cce9e940e410 100644 --- a/airflow/providers/google/cloud/transfers/adls_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/adls_to_gcs.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Sequence from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url -from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook from airflow.providers.microsoft.azure.operators.adls import ADLSListOperator if TYPE_CHECKING: @@ -120,6 +119,13 @@ def __init__( self.google_impersonation_chain = google_impersonation_chain def execute(self, context: Context): + try: + from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook + except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + # use the super to list all files in an Azure Data Lake path files = super().execute(context) g_hook = GCSHook( diff --git a/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py b/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py index 8ba6f2d6eb079..28645613ee99b 100644 --- a/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py @@ -22,7 +22,6 @@ from airflow.models import BaseOperator from airflow.providers.google.cloud.hooks.gcs import GCSHook -from airflow.providers.microsoft.azure.hooks.wasb import WasbHook if TYPE_CHECKING: from airflow.utils.context import Context @@ -88,6 +87,13 @@ def __init__( ) def execute(self, context: Context) -> str: + try: + from airflow.providers.microsoft.azure.hooks.wasb import WasbHook + except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + azure_hook = WasbHook(wasb_conn_id=self.wasb_conn_id) gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, diff --git a/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py b/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py index 9ba612979164c..10ada9fb64997 100644 --- a/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py @@ -24,7 +24,6 @@ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.models import BaseOperator from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url, gcs_object_is_directory -from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook if TYPE_CHECKING: from airflow.utils.context import Context @@ -115,6 +114,13 @@ def _check_inputs(self) -> None: ) def execute(self, context: Context): + try: + from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook + except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + self._check_inputs() azure_fileshare_hook = AzureFileShareHook( share_name=self.share_name, From a26ffd43fb0625b36f2fdee30d155d2c3cc4b03a Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 20:57:14 +0200 Subject: [PATCH 090/105] refactor: Import of ADLSListOperator should be ignored when running Airflow 2.6.0 or lower --- airflow/providers/google/cloud/transfers/adls_to_gcs.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/transfers/adls_to_gcs.py b/airflow/providers/google/cloud/transfers/adls_to_gcs.py index 4cce9e940e410..9a13f188e6805 100644 --- a/airflow/providers/google/cloud/transfers/adls_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/adls_to_gcs.py @@ -24,7 +24,13 @@ from typing import TYPE_CHECKING, Sequence from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url -from airflow.providers.microsoft.azure.operators.adls import ADLSListOperator + +try: + from airflow.providers.microsoft.azure.operators.adls import ADLSListOperator +except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) if TYPE_CHECKING: from airflow.utils.context import Context From 1259b326af93cd5a3113724aed7254a40fb4ca11 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 11 Apr 2024 23:03:53 +0200 Subject: [PATCH 091/105] refactor: Moved optional provider imports that should be ignored when running Airflow 2.6.0 or lower at top of file --- .../amazon/aws/transfers/azure_blob_to_s3.py | 14 +++++++------- .../google/cloud/transfers/adls_to_gcs.py | 8 +------- .../google/cloud/transfers/azure_blob_to_gcs.py | 14 +++++++------- .../cloud/transfers/azure_fileshare_to_gcs.py | 14 +++++++------- 4 files changed, 22 insertions(+), 28 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py b/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py index 40b6447fb7d2e..9af93e212b394 100644 --- a/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py @@ -24,6 +24,13 @@ from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.s3 import S3Hook +try: + from airflow.providers.microsoft.azure.hooks.wasb import WasbHook +except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + if TYPE_CHECKING: from airflow.utils.context import Context @@ -110,13 +117,6 @@ def __init__( def execute(self, context: Context) -> list[str]: # list all files in the Azure Blob Storage container - try: - from airflow.providers.microsoft.azure.hooks.wasb import WasbHook - except ModuleNotFoundError as e: - from airflow.exceptions import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException(e) - wasb_hook = WasbHook(wasb_conn_id=self.wasb_conn_id, **self.wasb_extra_args) s3_hook = S3Hook( aws_conn_id=self.aws_conn_id, diff --git a/airflow/providers/google/cloud/transfers/adls_to_gcs.py b/airflow/providers/google/cloud/transfers/adls_to_gcs.py index 9a13f188e6805..f11b6aa881b80 100644 --- a/airflow/providers/google/cloud/transfers/adls_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/adls_to_gcs.py @@ -26,6 +26,7 @@ from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url try: + from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook from airflow.providers.microsoft.azure.operators.adls import ADLSListOperator except ModuleNotFoundError as e: from airflow.exceptions import AirflowOptionalProviderFeatureException @@ -125,13 +126,6 @@ def __init__( self.google_impersonation_chain = google_impersonation_chain def execute(self, context: Context): - try: - from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook - except ModuleNotFoundError as e: - from airflow.exceptions import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException(e) - # use the super to list all files in an Azure Data Lake path files = super().execute(context) g_hook = GCSHook( diff --git a/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py b/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py index 28645613ee99b..1da9e82c09247 100644 --- a/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py @@ -23,6 +23,13 @@ from airflow.models import BaseOperator from airflow.providers.google.cloud.hooks.gcs import GCSHook +try: + from airflow.providers.microsoft.azure.hooks.wasb import WasbHook +except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + if TYPE_CHECKING: from airflow.utils.context import Context @@ -87,13 +94,6 @@ def __init__( ) def execute(self, context: Context) -> str: - try: - from airflow.providers.microsoft.azure.hooks.wasb import WasbHook - except ModuleNotFoundError as e: - from airflow.exceptions import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException(e) - azure_hook = WasbHook(wasb_conn_id=self.wasb_conn_id) gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, diff --git a/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py b/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py index 10ada9fb64997..cca318001c779 100644 --- a/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py @@ -25,6 +25,13 @@ from airflow.models import BaseOperator from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url, gcs_object_is_directory +try: + from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook +except ModuleNotFoundError as e: + from airflow.exceptions import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException(e) + if TYPE_CHECKING: from airflow.utils.context import Context @@ -114,13 +121,6 @@ def _check_inputs(self) -> None: ) def execute(self, context: Context): - try: - from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook - except ModuleNotFoundError as e: - from airflow.exceptions import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException(e) - self._check_inputs() azure_fileshare_hook = AzureFileShareHook( share_name=self.share_name, From 201f3718cbb945e092c32f6093cdcbcdc5a76c4b Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 08:46:49 +0200 Subject: [PATCH 092/105] refactor: Fixed the event loop closed issue when executing long running tests on the MSGraphOperator --- tests/providers/microsoft/azure/base.py | 47 ++++++++----------- .../microsoft/azure/triggers/test_msgraph.py | 5 +- 2 files changed, 23 insertions(+), 29 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 3ad48f99ec612..77d3f23c8b657 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -17,13 +17,12 @@ from __future__ import annotations import asyncio -from contextlib import closing, contextmanager +from contextlib import contextmanager from copy import deepcopy from datetime import datetime from typing import TYPE_CHECKING, Any, Iterable from unittest.mock import patch -import pytest from kiota_http.httpx_request_adapter import HttpxRequestAdapter from airflow.exceptions import TaskDeferred @@ -92,39 +91,33 @@ async def _run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]: return events def run_trigger(self, trigger: BaseTrigger) -> list[TriggerEvent]: - return self.run_async(self._run_tigger(trigger)) - - def run_async(self, future: Any) -> Any: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - with closing(loop): - return loop.run_until_complete(future) + return asyncio.run(self._run_tigger(trigger)) def execute_operator(self, operator: Operator) -> tuple[Any, Any]: task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) context = {"ti": task_instance} + return asyncio.run(self.deferrable_operator(context, operator)) + + async def deferrable_operator(self, context, operator): result = None triggered_events = [] + try: + result = operator.execute(context=context) + except TaskDeferred as deferred: + task = deferred - with pytest.raises(TaskDeferred) as deferred: - operator.execute(context=context) - - task = deferred.value - - while task: - events = self.run_trigger(task.trigger) - - if not events: - break + while task: + events = await self._run_tigger(task.trigger) - triggered_events.extend(deepcopy(events)) + if not events: + break - try: - method = getattr(operator, deferred.value.method_name) - result = method(context=context, event=next(iter(events)).payload) - task = None - except TaskDeferred as exception: - task = exception + triggered_events.extend(deepcopy(events)) + try: + method = getattr(operator, task.method_name) + result = method(context=context, event=next(iter(events)).payload) + task = None + except TaskDeferred as exception: + task = exception return result, triggered_events diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index b155f0b7aabff..107a5d993fcef 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import asyncio import json import locale from base64 import b64decode, b64encode @@ -126,12 +127,12 @@ def test_serialize(self): } -class TestResponseHandler(Base): +class TestResponseHandler: def test_handle_response_async(self): users = load_json("resources", "users.json") response = mock_json_response(200, users) - actual = self.run_async( + actual = asyncio.run( CallableResponseHandler(lambda response, error_map: response.json()).handle_response_async( response, None ) From aac5db715af9b2e970d3001b8b41f7b466c39211 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 09:33:08 +0200 Subject: [PATCH 093/105] refactor: Extracted reusable mock_context method --- tests/providers/microsoft/azure/base.py | 6 +-- .../microsoft/azure/sensors/test_msgraph.py | 9 ++-- tests/providers/microsoft/conftest.py | 41 +++++++++++++++++++ 3 files changed, 46 insertions(+), 10 deletions(-) diff --git a/tests/providers/microsoft/azure/base.py b/tests/providers/microsoft/azure/base.py index 77d3f23c8b657..4cda62858e815 100644 --- a/tests/providers/microsoft/azure/base.py +++ b/tests/providers/microsoft/azure/base.py @@ -29,9 +29,8 @@ from airflow.models import Operator, TaskInstance from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook from airflow.utils.session import NEW_SESSION -from airflow.utils.state import TaskInstanceState from airflow.utils.xcom import XCOM_RETURN_KEY -from tests.providers.microsoft.conftest import get_airflow_connection +from tests.providers.microsoft.conftest import get_airflow_connection, mock_context if TYPE_CHECKING: from sqlalchemy.orm import Session @@ -94,8 +93,7 @@ def run_trigger(self, trigger: BaseTrigger) -> list[TriggerEvent]: return asyncio.run(self._run_tigger(trigger)) def execute_operator(self, operator: Operator) -> tuple[Any, Any]: - task_instance = MockedTaskInstance(task=operator, run_id="run_id", state=TaskInstanceState.RUNNING) - context = {"ti": task_instance} + context = mock_context(task=operator) return asyncio.run(self.deferrable_operator(context, operator)) async def deferrable_operator(self, context, operator): diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index 4ea855b5a699e..2ada550daf92e 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -20,9 +20,8 @@ from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor from airflow.triggers.base import TriggerEvent -from airflow.utils.state import TaskInstanceState -from tests.providers.microsoft.azure.base import Base, MockedTaskInstance -from tests.providers.microsoft.conftest import load_json, mock_json_response +from tests.providers.microsoft.azure.base import Base +from tests.providers.microsoft.conftest import load_json, mock_context, mock_json_response class TestMSGraphSensor(Base): @@ -37,9 +36,7 @@ def test_execute(self): url="myorg/admin/workspaces/scanStatus/0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", timeout=350.0, ) - actual = sensor.execute( - context=MockedTaskInstance(task=sensor, run_id="run_id", state=TaskInstanceState.RUNNING) - ) + actual = sensor.execute(context=mock_context(task=sensor)) assert isinstance(actual, TriggerEvent) assert actual.payload["status"] == "success" diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index 022fde060f0c6..975eb214fa387 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -100,6 +100,47 @@ def mock_response(status_code, content: Any = None) -> Response: return response +def mock_context(task): + from datetime import datetime + from sqlalchemy.orm import Session + + from airflow.models import TaskInstance + from airflow.utils.session import NEW_SESSION + from airflow.utils.state import TaskInstanceState + from airflow.utils.xcom import XCOM_RETURN_KEY + + class MockedTaskInstance(TaskInstance): + def __init__(self): + super().__init__(task=task, run_id="run_id", state=TaskInstanceState.RUNNING) + self.values = {} + + def xcom_pull( + self, + task_ids: Iterable[str] | str | None = None, + dag_id: str | None = None, + key: str = XCOM_RETURN_KEY, + include_prior_dates: bool = False, + session: Session = NEW_SESSION, + *, + map_indexes: Iterable[int] | int | None = None, + default: Any | None = None, + ) -> Any: + self.task_id = task_ids + self.dag_id = dag_id + return self.values.get(f"{task_ids}_{dag_id}_{key}") + + def xcom_push( + self, + key: str, + value: Any, + execution_date: datetime | None = None, + session: Session = NEW_SESSION, + ) -> None: + self.values[f"{self.task_id}_{self.dag_id}_{key}"] = value + + return {"ti": MockedTaskInstance()} + + def load_json(*locations: Iterable[str]): with open(join(dirname(__file__), "azure", join(*locations)), encoding="utf-8") as file: return json.load(file) From 14e932713239606092e5aab7b694f8514672bd73 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 10:02:04 +0200 Subject: [PATCH 094/105] refactor: Moved import of Session into type checking block --- tests/providers/microsoft/conftest.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index 975eb214fa387..87eb7816ba98d 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -21,7 +21,7 @@ import random import string from os.path import dirname, join -from typing import Any, Iterable, TypeVar +from typing import TYPE_CHECKING, Any, Iterable, TypeVar from unittest.mock import MagicMock import pytest @@ -30,6 +30,9 @@ from airflow.models import Connection +if TYPE_CHECKING: + from sqlalchemy.orm import Session + T = TypeVar("T", dict, str, Connection) @@ -102,7 +105,6 @@ def mock_response(status_code, content: Any = None) -> Response: def mock_context(task): from datetime import datetime - from sqlalchemy.orm import Session from airflow.models import TaskInstance from airflow.utils.session import NEW_SESSION From aed317a4273db7b60670a18be8098531ab0180a5 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 10:02:21 +0200 Subject: [PATCH 095/105] refactor: Updated the TestMSGraphSensor --- tests/providers/microsoft/azure/sensors/test_msgraph.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index 2ada550daf92e..df55cb1d18f0c 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -33,7 +33,8 @@ def test_execute(self): sensor = MSGraphSensor( task_id="check_workspaces_status", conn_id="powerbi", - url="myorg/admin/workspaces/scanStatus/0a1b1bf3-37de-48f7-9863-ed4cda97a9ef", + url="myorg/admin/workspaces/scanStatus/{scanId}", + path_parameters={"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}, timeout=350.0, ) actual = sensor.execute(context=mock_context(task=sensor)) From d6f7b45902c27e769efd12da77140cbb485edec1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 10:03:04 +0200 Subject: [PATCH 096/105] refactor: Reformatted the mock_context method --- tests/providers/microsoft/conftest.py | 28 +++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/providers/microsoft/conftest.py b/tests/providers/microsoft/conftest.py index 87eb7816ba98d..78d8748a89e04 100644 --- a/tests/providers/microsoft/conftest.py +++ b/tests/providers/microsoft/conftest.py @@ -117,26 +117,26 @@ def __init__(self): self.values = {} def xcom_pull( - self, - task_ids: Iterable[str] | str | None = None, - dag_id: str | None = None, - key: str = XCOM_RETURN_KEY, - include_prior_dates: bool = False, - session: Session = NEW_SESSION, - *, - map_indexes: Iterable[int] | int | None = None, - default: Any | None = None, + self, + task_ids: Iterable[str] | str | None = None, + dag_id: str | None = None, + key: str = XCOM_RETURN_KEY, + include_prior_dates: bool = False, + session: Session = NEW_SESSION, + *, + map_indexes: Iterable[int] | int | None = None, + default: Any | None = None, ) -> Any: self.task_id = task_ids self.dag_id = dag_id return self.values.get(f"{task_ids}_{dag_id}_{key}") def xcom_push( - self, - key: str, - value: Any, - execution_date: datetime | None = None, - session: Session = NEW_SESSION, + self, + key: str, + value: Any, + execution_date: datetime | None = None, + session: Session = NEW_SESSION, ) -> None: self.values[f"{self.task_id}_{self.dag_id}_{key}"] = value From 693975eb8dbf8a2982f3b3d05a4d385b312009f9 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 11:04:40 +0200 Subject: [PATCH 097/105] refactor: Try implementing cached connections on MSGraphTrigger --- airflow/providers/microsoft/azure/triggers/msgraph.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index e48e4d7e9cd2a..61cf842a30056 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -135,6 +135,7 @@ class MSGraphTrigger(BaseTrigger): """ DEFAULT_HEADERS = {"Accept": "application/json;q=1"} + cached_connections: dict[str:RequestAdapter] = {} template_fields: Sequence[str] = ( "url", "response_type", @@ -214,7 +215,14 @@ def serialize(self) -> tuple[str, dict[str, Any]]: ) def get_conn(self) -> RequestAdapter: - return self.hook.get_conn() + connection = self.cached_connections.get(self.conn_id) + if not connection: + self.log.info("No cached connection found for '%s'", self.conn_id) + connection = self.hook.get_conn() + self.cached_connections[self.conn_id] = connection + else: + self.log.info("Cached connection found for '%s'", self.conn_id) + return connection @property def conn_id(self) -> str: From 0f998613e8bb83ad5b9b66f14df7bc75700ac9d1 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 11:39:24 +0200 Subject: [PATCH 098/105] docs: Added example for the MSGraphSensor and additional examples on how you can use the operator for PowerBI --- .../operators/msgraph.rst | 17 ++++ .../sensors/msgraph.rst | 42 ++++++++++ .../microsoft/azure/example_powerbi.py | 80 +++++++++++++++++++ 3 files changed, 139 insertions(+) create mode 100644 docs/apache-airflow-providers-microsoft-azure/sensors/msgraph.rst create mode 100644 tests/system/providers/microsoft/azure/example_powerbi.py diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst index 8ef239924acd0..817b14f783142 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst @@ -48,6 +48,22 @@ Below is an example of using this operator to get a Sharepoint site pages. :start-after: [START howto_operator_graph_site_pages] :end-before: [END howto_operator_graph_site_pages] +Below is an example of using this operator to get PowerBI workspaces. + +.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_powerbi.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_powerbi_workspaces] + :end-before: [END howto_operator_powerbi_workspaces] + +Below is an example of using this operator to get PowerBI workspaces info. + +.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_powerbi.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_powerbi_workspaces_info] + :end-before: [END howto_operator_powerbi_workspaces_info] + Reference --------- @@ -55,3 +71,4 @@ Reference For further information, look at: * `Use the Microsoft Graph API `__ +* `Using the Power BI REST APIs `__ diff --git a/docs/apache-airflow-providers-microsoft-azure/sensors/msgraph.rst b/docs/apache-airflow-providers-microsoft-azure/sensors/msgraph.rst new file mode 100644 index 0000000000000..4ddad88f19fa1 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/sensors/msgraph.rst @@ -0,0 +1,42 @@ + .. 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 Graph API Sensors +============================= + +MSGraphSensor +------------- +Use the +:class:`~airflow.providers.microsoft.azure.sensors.msgraph.MSGraphSensor` to poll a Power BI API. + + +Below is an example of using this sensor to poll the status of a PowerBI workspace. + +.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_powerbi.py + :language: python + :dedent: 0 + :start-after: [START howto_sensor_powerbi_scan_status] + :end-before: [END howto_sensor_powerbi_scan_status] + + +Reference +--------- + +For further information, look at: + +* `Using the Power BI REST APIs `__ diff --git a/tests/system/providers/microsoft/azure/example_powerbi.py b/tests/system/providers/microsoft/azure/example_powerbi.py new file mode 100644 index 0000000000000..cbee9a62af0c4 --- /dev/null +++ b/tests/system/providers/microsoft/azure/example_powerbi.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 + +from datetime import datetime + +from airflow import models +from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator +from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor + +DAG_ID = "example_powerbi" + +with models.DAG( + DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + tags=["example"], +) as dag: + # [START howto_operator_powerbi_workspaces] + workspaces_task = MSGraphAsyncOperator( + task_id="workspaces", + conn_id="powerbi", + url="myorg/admin/workspaces/modified", + result_processor=lambda context, response: list(map(lambda workspace: workspace["id"], response)), + ) + # [END howto_operator_powerbi_workspaces] + + # [START howto_operator_powerbi_workspaces_info] + workspaces_info_task = MSGraphAsyncOperator( + task_id="get_workspace_info", + conn_id="powerbi", + url="myorg/admin/workspaces/getInfo", + method="POST", + query_parameters={ + "lineage": True, + "datasourceDetails": True, + "datasetSchema": True, + "datasetExpressions": True, + "getArtifactUsers": True, + }, + data={"workspaces": workspaces_task.output}, + result_processor=lambda context, response: {"scanId": response["id"]}, + ) + # [END howto_operator_powerbi_workspaces_info] + + # [START howto_sensor_powerbi_scan_status] + check_workspace_status_task = MSGraphSensor.partial( + task_id="check_workspaces_status", + conn_id="powerbi_api", + url="myorg/admin/workspaces/scanStatus/{scanId}", + timeout=350.0, + ).expand(path_parameters=workspaces_info_task.output) + # [END howto_sensor_powerbi_scan_status] + + workspaces_task >> workspaces_info_task >> check_workspace_status_task + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) From 78494c7df20bfe4da3896332521111896a86f32d Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 13:02:02 +0200 Subject: [PATCH 099/105] Revert "refactor: Try implementing cached connections on MSGraphTrigger" This reverts commit 693975eb8dbf8a2982f3b3d05a4d385b312009f9. --- airflow/providers/microsoft/azure/triggers/msgraph.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index 61cf842a30056..e48e4d7e9cd2a 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -135,7 +135,6 @@ class MSGraphTrigger(BaseTrigger): """ DEFAULT_HEADERS = {"Accept": "application/json;q=1"} - cached_connections: dict[str:RequestAdapter] = {} template_fields: Sequence[str] = ( "url", "response_type", @@ -215,14 +214,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: ) def get_conn(self) -> RequestAdapter: - connection = self.cached_connections.get(self.conn_id) - if not connection: - self.log.info("No cached connection found for '%s'", self.conn_id) - connection = self.hook.get_conn() - self.cached_connections[self.conn_id] = connection - else: - self.log.info("Cached connection found for '%s'", self.conn_id) - return connection + return self.hook.get_conn() @property def conn_id(self) -> str: From e571db66ec8e4b600f63f8d44f055148a204266d Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 13:33:12 +0200 Subject: [PATCH 100/105] fix: Fixed serialization of event payload as xcom_value for the MSGraphSensor --- airflow/providers/microsoft/azure/sensors/msgraph.py | 10 +++++++--- .../providers/microsoft/azure/sensors/test_msgraph.py | 11 ++++------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index 2836333bba375..b05cf8e6861a8 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -103,7 +103,7 @@ def __init__( self.proxies = proxies self.api_version = api_version self.event_processor = event_processor - self.serializer = serializer + self.serializer = serializer() @property def trigger(self): @@ -121,7 +121,7 @@ def trigger(self): timeout=self.timeout, proxies=self.proxies, api_version=self.api_version, - serializer=self.serializer, + serializer=type(self.serializer), ) async def async_poke(self, context: Context) -> bool | PokeReturnValue: @@ -134,7 +134,11 @@ async def async_poke(self, context: Context) -> bool | PokeReturnValue: self.log.debug("is_done: %s", is_done) - return PokeReturnValue(is_done=is_done, xcom_value=response) + value = self.serializer.deserialize(response.payload["response"]) + + self.log.debug("value: %s", value) + + return PokeReturnValue(is_done=is_done, xcom_value=value) return PokeReturnValue(is_done=True) def poke(self, context) -> bool | PokeReturnValue: diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index df55cb1d18f0c..8c03c6fb3b853 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -16,10 +16,7 @@ # under the License. from __future__ import annotations -import json - from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor -from airflow.triggers.base import TriggerEvent from tests.providers.microsoft.azure.base import Base from tests.providers.microsoft.conftest import load_json, mock_context, mock_json_response @@ -39,7 +36,7 @@ def test_execute(self): ) actual = sensor.execute(context=mock_context(task=sensor)) - assert isinstance(actual, TriggerEvent) - assert actual.payload["status"] == "success" - assert actual.payload["type"] == "builtins.dict" - assert actual.payload["response"] == json.dumps(status) + assert isinstance(actual, dict) + assert actual["id"] == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef" + assert actual["createdDateTime"] == "2024-04-10T15:05:17.357" + assert actual["status"] == "Succeeded" From c7a06dbab1c516e9ffb67ad279f425c640d7a851 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 14:02:04 +0200 Subject: [PATCH 101/105] refactor: TestMSGraphAsyncOperator should be allowed to run as a db test --- tests/providers/microsoft/azure/operators/test_msgraph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 1520fe603e7fb..984b4e978f74e 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -30,6 +30,7 @@ class TestMSGraphAsyncOperator(Base): + @pytest.mark.db_test def test_execute(self): users = load_json("resources", "users.json") next_users = load_json("resources", "next_users.json") @@ -57,6 +58,7 @@ def test_execute(self): assert events[1].payload["type"] == "builtins.dict" assert events[1].payload["response"] == json.dumps(next_users) + @pytest.mark.db_test def test_execute_when_do_xcom_push_is_false(self): users = load_json("resources", "users.json") users.pop("@odata.nextLink") @@ -79,6 +81,7 @@ def test_execute_when_do_xcom_push_is_false(self): assert events[0].payload["type"] == "builtins.dict" assert events[0].payload["response"] == json.dumps(users) + @pytest.mark.db_test def test_execute_when_an_exception_occurs(self): with self.patch_hook_and_request_adapter(AirflowException()): operator = MSGraphAsyncOperator( @@ -91,6 +94,7 @@ def test_execute_when_an_exception_occurs(self): with pytest.raises(AirflowException): self.execute_operator(operator) + @pytest.mark.db_test def test_execute_when_response_is_bytes(self): content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) From 3e508b0bdb66dc60b047f09bf92c30b376112858 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 14:36:52 +0200 Subject: [PATCH 102/105] Revert "refactor: TestMSGraphAsyncOperator should be allowed to run as a db test" This reverts commit c7a06dbab1c516e9ffb67ad279f425c640d7a851. --- tests/providers/microsoft/azure/operators/test_msgraph.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 984b4e978f74e..1520fe603e7fb 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -30,7 +30,6 @@ class TestMSGraphAsyncOperator(Base): - @pytest.mark.db_test def test_execute(self): users = load_json("resources", "users.json") next_users = load_json("resources", "next_users.json") @@ -58,7 +57,6 @@ def test_execute(self): assert events[1].payload["type"] == "builtins.dict" assert events[1].payload["response"] == json.dumps(next_users) - @pytest.mark.db_test def test_execute_when_do_xcom_push_is_false(self): users = load_json("resources", "users.json") users.pop("@odata.nextLink") @@ -81,7 +79,6 @@ def test_execute_when_do_xcom_push_is_false(self): assert events[0].payload["type"] == "builtins.dict" assert events[0].payload["response"] == json.dumps(users) - @pytest.mark.db_test def test_execute_when_an_exception_occurs(self): with self.patch_hook_and_request_adapter(AirflowException()): operator = MSGraphAsyncOperator( @@ -94,7 +91,6 @@ def test_execute_when_an_exception_occurs(self): with pytest.raises(AirflowException): self.execute_operator(operator) - @pytest.mark.db_test def test_execute_when_response_is_bytes(self): content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) From e318581abfbecb96e8a763204fd245bf2bfe4f67 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 14:02:04 +0200 Subject: [PATCH 103/105] refactor: TestMSGraphAsyncOperator should be allowed to run as a db test --- tests/providers/microsoft/azure/operators/test_msgraph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 1520fe603e7fb..984b4e978f74e 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -30,6 +30,7 @@ class TestMSGraphAsyncOperator(Base): + @pytest.mark.db_test def test_execute(self): users = load_json("resources", "users.json") next_users = load_json("resources", "next_users.json") @@ -57,6 +58,7 @@ def test_execute(self): assert events[1].payload["type"] == "builtins.dict" assert events[1].payload["response"] == json.dumps(next_users) + @pytest.mark.db_test def test_execute_when_do_xcom_push_is_false(self): users = load_json("resources", "users.json") users.pop("@odata.nextLink") @@ -79,6 +81,7 @@ def test_execute_when_do_xcom_push_is_false(self): assert events[0].payload["type"] == "builtins.dict" assert events[0].payload["response"] == json.dumps(users) + @pytest.mark.db_test def test_execute_when_an_exception_occurs(self): with self.patch_hook_and_request_adapter(AirflowException()): operator = MSGraphAsyncOperator( @@ -91,6 +94,7 @@ def test_execute_when_an_exception_occurs(self): with pytest.raises(AirflowException): self.execute_operator(operator) + @pytest.mark.db_test def test_execute_when_response_is_bytes(self): content = load_file("resources", "dummy.pdf", mode="rb", encoding=None) base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding()) From ea3957eb4237f5a0aab786b893e6357654eedc21 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 18:25:52 +0200 Subject: [PATCH 104/105] refactor: Also added result_processor to MSGraphSensor --- .../microsoft/azure/sensors/msgraph.py | 21 +++++++++++++------ .../microsoft/azure/sensors/test_msgraph.py | 7 +++---- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index b05cf8e6861a8..a3bfe4b62b0c6 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -63,6 +63,9 @@ class MSGraphSensor(BaseSensorOperator): :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. """ @@ -86,6 +89,7 @@ def __init__( proxies: dict | None = None, api_version: APIVersion | None = None, event_processor: Callable[[Context, TriggerEvent], bool] = default_event_processor, + result_processor: Callable[[Context, Any], Any] = lambda context, result: result, serializer: type[ResponseSerializer] = ResponseSerializer, **kwargs, ): @@ -103,6 +107,7 @@ def __init__( self.proxies = proxies self.api_version = api_version self.event_processor = event_processor + self.result_processor = result_processor self.serializer = serializer() @property @@ -127,18 +132,22 @@ def trigger(self): async def async_poke(self, context: Context) -> bool | PokeReturnValue: self.log.info("Sensor triggered") - async for response in self.trigger.run(): - self.log.debug("response: %s", response) + async for event in self.trigger.run(): + self.log.debug("event: %s", event) - is_done = self.event_processor(context, response) + is_done = self.event_processor(context, event) self.log.debug("is_done: %s", is_done) - value = self.serializer.deserialize(response.payload["response"]) + response = self.serializer.deserialize(event.payload["response"]) - self.log.debug("value: %s", value) + self.log.debug("deserialize event: %s", response) - return PokeReturnValue(is_done=is_done, xcom_value=value) + result = self.result_processor(context, response) + + self.log.debug("result: %s", result) + + return PokeReturnValue(is_done=is_done, xcom_value=result) return PokeReturnValue(is_done=True) def poke(self, context) -> bool | PokeReturnValue: diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index 8c03c6fb3b853..73052045a4578 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -32,11 +32,10 @@ def test_execute(self): conn_id="powerbi", url="myorg/admin/workspaces/scanStatus/{scanId}", path_parameters={"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}, + result_processor=lambda context, result: result["id"], timeout=350.0, ) actual = sensor.execute(context=mock_context(task=sensor)) - assert isinstance(actual, dict) - assert actual["id"] == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef" - assert actual["createdDateTime"] == "2024-04-10T15:05:17.357" - assert actual["status"] == "Succeeded" + assert isinstance(actual, str) + assert actual == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef" From 1479af78519aaa022dd29535eb6a3b4a063fd03d Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Apr 2024 20:20:51 +0200 Subject: [PATCH 105/105] refactor: Fixed template_fields in operator, trigger and sensor --- .../providers/microsoft/azure/operators/msgraph.py | 11 ++++++++++- airflow/providers/microsoft/azure/sensors/msgraph.py | 11 ++++++++++- airflow/providers/microsoft/azure/triggers/msgraph.py | 2 +- .../microsoft/azure/operators/test_msgraph.py | 10 ++++++++++ .../providers/microsoft/azure/sensors/test_msgraph.py | 10 ++++++++++ .../microsoft/azure/triggers/test_msgraph.py | 6 ++++++ 6 files changed, 47 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/msgraph.py b/airflow/providers/microsoft/azure/operators/msgraph.py index 59e5b809d09e8..6411f9cc4ac2d 100644 --- a/airflow/providers/microsoft/azure/operators/msgraph.py +++ b/airflow/providers/microsoft/azure/operators/msgraph.py @@ -78,7 +78,16 @@ class MSGraphAsyncOperator(BaseOperator): Bytes will be base64 encoded into a string, so it can be stored as an XCom. """ - template_fields: Sequence[str] = ("url", "conn_id") + template_fields: Sequence[str] = ( + "url", + "response_type", + "path_parameters", + "url_template", + "query_parameters", + "headers", + "data", + "conn_id", + ) def __init__( self, diff --git a/airflow/providers/microsoft/azure/sensors/msgraph.py b/airflow/providers/microsoft/azure/sensors/msgraph.py index a3bfe4b62b0c6..ffbf244dbe88c 100644 --- a/airflow/providers/microsoft/azure/sensors/msgraph.py +++ b/airflow/providers/microsoft/azure/sensors/msgraph.py @@ -70,7 +70,16 @@ class MSGraphSensor(BaseSensorOperator): Bytes will be base64 encoded into a string, so it can be stored as an XCom. """ - template_fields: Sequence[str] = ("url", "conn_id") + template_fields: Sequence[str] = ( + "url", + "response_type", + "path_parameters", + "url_template", + "query_parameters", + "headers", + "data", + "conn_id", + ) def __init__( self, diff --git a/airflow/providers/microsoft/azure/triggers/msgraph.py b/airflow/providers/microsoft/azure/triggers/msgraph.py index e48e4d7e9cd2a..c0e5ee85a0c4c 100644 --- a/airflow/providers/microsoft/azure/triggers/msgraph.py +++ b/airflow/providers/microsoft/azure/triggers/msgraph.py @@ -142,7 +142,7 @@ class MSGraphTrigger(BaseTrigger): "url_template", "query_parameters", "headers", - "content", + "data", "conn_id", ) diff --git a/tests/providers/microsoft/azure/operators/test_msgraph.py b/tests/providers/microsoft/azure/operators/test_msgraph.py index 984b4e978f74e..b7520d731544c 100644 --- a/tests/providers/microsoft/azure/operators/test_msgraph.py +++ b/tests/providers/microsoft/azure/operators/test_msgraph.py @@ -117,3 +117,13 @@ def test_execute_when_response_is_bytes(self): assert events[0].payload["status"] == "success" assert events[0].payload["type"] == "builtins.bytes" assert events[0].payload["response"] == base64_encoded_content + + def test_template_fields(self): + operator = MSGraphAsyncOperator( + task_id="drive_item_content", + conn_id="msgraph_api", + url="users/delta", + ) + + for template_field in MSGraphAsyncOperator.template_fields: + getattr(operator, template_field) diff --git a/tests/providers/microsoft/azure/sensors/test_msgraph.py b/tests/providers/microsoft/azure/sensors/test_msgraph.py index 73052045a4578..50fd2474ab454 100644 --- a/tests/providers/microsoft/azure/sensors/test_msgraph.py +++ b/tests/providers/microsoft/azure/sensors/test_msgraph.py @@ -39,3 +39,13 @@ def test_execute(self): assert isinstance(actual, str) assert actual == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef" + + def test_template_fields(self): + sensor = MSGraphSensor( + task_id="check_workspaces_status", + conn_id="powerbi", + url="myorg/admin/workspaces/scanStatus/{scanId}", + ) + + for template_field in MSGraphSensor.template_fields: + getattr(sensor, template_field) diff --git a/tests/providers/microsoft/azure/triggers/test_msgraph.py b/tests/providers/microsoft/azure/triggers/test_msgraph.py index 107a5d993fcef..900d0875cd0f2 100644 --- a/tests/providers/microsoft/azure/triggers/test_msgraph.py +++ b/tests/providers/microsoft/azure/triggers/test_msgraph.py @@ -126,6 +126,12 @@ def test_serialize(self): "serializer": "airflow.providers.microsoft.azure.triggers.msgraph.ResponseSerializer", } + def test_template_fields(self): + trigger = MSGraphTrigger("users/delta", response_type="bytes", conn_id="msgraph_api") + + for template_field in MSGraphTrigger.template_fields: + getattr(trigger, template_field) + class TestResponseHandler: def test_handle_response_async(self):