Adding MSGraphSDKAsyncOperator

[dabla]

Hello, I'm building a custom async MSGraphSDKAsyncOperator which uses the official python msgraph-sdk-python library from Microsoft. I've implemented a hook, a trigger and an operator. I want the operator to work asynchronously because the msgraph-sdk-python is also implemented in an async way, which is in my opinion a good thing but makes things a bit harder.

I've succeeded in the implementation, and I hope once finished I will be able to contribute this operator to Airflow, but I have one issue which I don't understand. When doing calls with the client, it is possible that you receive a paged response, so my idea was that when it is the case, the trigger would automatically call the next pages until it's done. The trigger is being deferred from the operator.
Below a snippet of the trigger run method how I implemented it:

async def run(self) -> AsyncIterator[TriggerEvent]:
    """Makes a series of asynchronous http calls via a MSGraphSDKHook."""
    try:
        response = await self.hook.run(
            expression=self.expression,
        )

        self.log.info("response: %s", response)

        if response:
            async for index, value in enumerate(self.to_odata_iterator(response)):  # async enumerate from aiostream.stream
                self.log.info("value[%s]: %s", index, value)
                value_type = type(value)
                yield TriggerEvent(
                    {
                        "status": "success",
                        "type": f"{value_type.__module__}.{value_type.__name__}",
                        "response": ResponseSerializer.serialize(value),
                    }
                )
    except Exception as e:
        self.log.error("%s", e)
        yield TriggerEvent({"status": "failure", "message": str(e)})

As you can see I implemented an AsyncIterator named ODataIterator which handles the paged response and calls the next one on each async iteration. I've tested the ODataIterator and it works fine, the problem is that after the second iteration the trigger is being stopped due to an asyncio.CancelledError, and I don't understand why as no exception is being throw and no timeout has been specified so that could also not be the case. When I run the same method shown above outside Airflow, it completes succesfully and goes through all the paged responses (it takes a while to complete as there are a lot of paged responses). This is the code in Airflow where the exception happens in the triggerer_job_runner.py module:

async def run_trigger(self, trigger_id, trigger):
    """Run a trigger (they are async generators) and push their events into our outbound event deque."""
    name = self.triggers[trigger_id]["name"]
    self.log.info("trigger %s starting", name)
    try:
        self.set_individual_trigger_logging(trigger)
        async for event in trigger.run():
            self.log.info("Trigger %s fired: %s", self.triggers[trigger_id]["name"], event)
            self.triggers[trigger_id]["events"] += 1
            self.events.append((trigger_id, event))
    except asyncio.CancelledError as err:
        if timeout := trigger.task_instance.trigger_timeout:
            timeout = timeout.replace(tzinfo=timezone.utc) if not timeout.tzinfo else timeout
            if timeout < timezone.utcnow():
                self.log.error("Trigger cancelled due to timeout")
        self.log.error("Trigger cancelled; message=%s", err)
        raise
    finally:
        # CancelledError will get injected when we're stopped - which is
        # fine, the cleanup process will understand that, but we want to
        # allow triggers a chance to cleanup, either in that case or if
        # they exit cleanly. Exception from cleanup methods are ignored.
        with suppress(Exception):
            await trigger.cleanup()
        if SEND_TRIGGER_END_MARKER:
            self.mark_trigger_end(trigger)

        # unsetting ctx_indiv_trigger var restores stdout logging
        ctx_indiv_trigger.set(None)
        self.log.info("trigger %s completed", name)

As you can see above, when an asyncio.CancelledError occurs, the exception is being raised and of course the triggers is stopped. I've skimmed through the source code of the operators provided with Airflow in hope to see if there where other examples where multiple events are being fired through an AsyncIterator but I couldn't find any. Maybe what I'm trying here isn't also the good way to do it in Airflow. I've also tried an alternative approach by deferring each succesive paged call which actually worked but then I wonder what the purpose of the AsyncIterator is for events if you're only allowed to trigger at most 2 events?

To write my async operator, I got inspired by the following article on betterprogramming.

Below the integration tests which runs the code and succesfully goes through all 11 pages:

from typing import List
from unittest.mock import patch

from airflow import configuration
from airflow.providers.microsoft.msgraph.triggers.msgraph import MSGraphSDKTrigger
from airflow.triggers.base import TriggerEvent
from assertpy import assert_that

from tests.unit.base import BaseTestCase


class MSGraphSDKTriggerIntegrationTestCase(BaseTestCase):
    def setUp(self):
        configuration.load_test_config()

    @staticmethod
    async def run_tigger(trigger: MSGraphSDKTrigger) -> List[TriggerEvent]:
        events = []
        async for event in trigger.run():
            events.append(event)
        return events

    def test_run_when_expression_is_valid(self):
        with (patch(
                "airflow.hooks.base.BaseHook.get_connection",
                side_effect=self.get_airflow_connection,
        )):
            trigger = MSGraphSDKTrigger("users.delta.get()", conn_id="msgraph_api")
            actual = self._loop.run_until_complete(self.run_tigger(trigger))

            assert_that(actual).is_length(11)

[dabla]

I finally managed to make it work, the question now is would this new operator be accepted into Airflow?

The operators/triggers/hook are fully tested, below a screenshot from Sonarcube with the coverage.

I've tried to implement it as clean as possible following the clean code principles, even though I'm sure there is still room for improvement. I've also added docstrings to the classes explaining the parameters.

Bellow a snippit on how the Async MS Graph Operator can be used:

import logging

from airflow import DAG
from airflow.decorators import task
from airflow.operators.python_operator import PythonOperator
from airflow.providers.microsoft.msgraph.operators.msgraph import MSGraphSDKAsyncOperator
from pendulum import datetime


with DAG(
    "process_users_delta",
    start_date=datetime(2022, 12, 20),
    default_args={
        "owner": "dabla",
    },
    schedule_interval=None,
    catchup=False,
) as dag:
    @task(dag=dag)
    def get_response(**context):
        logging.info("context: %s", context)
        params = context.get("params")
        logging.info("params: %s", params)
        response = params.get("response")
        logging.info("response: %s", response)
        logging.info("response_type: %s", type(response))
        return response
        
    get_response()


with DAG(
    "test_msgraph_sdk",
    start_date=datetime(2022, 12, 20),
    default_args={
        "owner": "dabla",
    },
    schedule="@daily",
    catchup=False,
) as dag:
    users_delta_task = MSGraphSDKAsyncOperator(
        task_id="users_delta",
        conn_id="msgraph_api",
        expression="users.delta.get()",
        trigger_dag_id="process_users_delta",
    )

    site_task = MSGraphSDKAsyncOperator(
        task_id="wget_site",
        conn_id="msgraph_api",
        expression="sites.by_site_id('850v1v.sharepoint.com:/sites/wget').get()",
    )

    site_pages_task = MSGraphSDKAsyncOperator(
        task_id="wget_site_pages",
        conn_id="msgraph_api",
        expression=(
            "sites.by_site_id('%s').pages.get()"
            % "{{ ti.xcom_pull(task_ids='wget_site')['id'] }}"
        ),
    )

    @task(dag=dag)
    def show(response):
        logging.info("response: %s", response)


users_delta_task >> site_task >> site_pages_task >> show(site_pages_task.output)

A you can see, the first operator named users_delta_task in the test_msgraph_sdk DAG invokes an expression which returns all the delta user pages. Each paged result is directly transfered to the dag 'process_users_delta' as we specified the trigger_dag_id argument which means events don't get accumulated within the operator until it completes which is memory efficient. If on the other hand you don't want to trigger another DAG to process each paged result, but you want to be able to get the processed events, you can set the keep_events parameter to True (default is False) which will keep track of each processed event (e.g. paged result) and once the operator is finished processing, it will return all those events. The later one will of course consume more memory but allows you to work in the "classic way" and can be handy if you now that there won't be too much data. There is also an example on how to get a sharepoint site via the url and also how to get the pages related to that sharepoint site. Same parameters (expression, conn_id and trigger_dag_id) of the operator are templated which means as you can see in example above that you can use jinja expressions.

The source code can be found here on GitHub.

[dabla]

We now have published the provider on pypi: https://pypi.org/project/apache-airflow-providers-msgraph/

[Taragolis]

You also could add link to your provider in Airflow Ecosystem: Third Party Airflow Plugins and Providers by click on Suggest a change on this page button

[dabla]

Thanks for the suggestion, I've created a pull request for that page: apache/airflow-site#933