airflow_scheduler/.venv/Lib/site-packages/airflow/providers/http/triggers/http.py

# 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 base64
import pickle
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any

import aiohttp
import requests
from requests.cookies import RequestsCookieJar
from requests.structures import CaseInsensitiveDict

from airflow.exceptions import AirflowException
from airflow.providers.http.hooks.http import HttpAsyncHook
from airflow.triggers.base import BaseTrigger, TriggerEvent

if TYPE_CHECKING:
    from aiohttp.client_reqrep import ClientResponse


class HttpTrigger(BaseTrigger):
    """
    HttpTrigger run on the trigger worker.

    :param http_conn_id: http connection id that has the base
        API url i.e https://www.google.com/ and optional authentication credentials. Default
        headers can also be specified in the Extra field in json format.
    :param auth_type: The auth type for the service
    :param method: the API method to be called
    :param endpoint: Endpoint to be called, i.e. ``resource/v1/query?``.
    :param headers: Additional headers to be passed through as a dict.
    :param data: Payload to be uploaded or request parameters.
    :param extra_options: Additional kwargs to pass when creating a request.
        For example, ``run(json=obj)`` is passed as
        ``aiohttp.ClientSession().get(json=obj)``.
        2XX or 3XX status codes
    """

    def __init__(
        self,
        http_conn_id: str = "http_default",
        auth_type: Any = None,
        method: str = "POST",
        endpoint: str | None = None,
        headers: dict[str, str] | None = None,
        data: dict[str, Any] | str | None = None,
        extra_options: dict[str, Any] | None = None,
    ):
        super().__init__()
        self.http_conn_id = http_conn_id
        self.method = method
        self.auth_type = auth_type
        self.endpoint = endpoint
        self.headers = headers
        self.data = data
        self.extra_options = extra_options

    def serialize(self) -> tuple[str, dict[str, Any]]:
        """Serialize HttpTrigger arguments and classpath."""
        return (
            "airflow.providers.http.triggers.http.HttpTrigger",
            {
                "http_conn_id": self.http_conn_id,
                "method": self.method,
                "auth_type": self.auth_type,
                "endpoint": self.endpoint,
                "headers": self.headers,
                "data": self.data,
                "extra_options": self.extra_options,
            },
        )

    async def run(self) -> AsyncIterator[TriggerEvent]:
        """Make a series of asynchronous http calls via a http hook."""
        hook = HttpAsyncHook(
            method=self.method,
            http_conn_id=self.http_conn_id,
            auth_type=self.auth_type,
        )
        try:
            async with aiohttp.ClientSession() as session:
                client_response = await hook.run(
                    session=session,
                    endpoint=self.endpoint,
                    data=self.data,
                    headers=self.headers,
                    extra_options=self.extra_options,
                )
                response = await self._convert_response(client_response)
            yield TriggerEvent(
                {
                    "status": "success",
                    "response": base64.standard_b64encode(pickle.dumps(response)).decode("ascii"),
                }
            )
        except Exception as e:
            yield TriggerEvent({"status": "error", "message": str(e)})

    @staticmethod
    async def _convert_response(client_response: ClientResponse) -> requests.Response:
        """Convert aiohttp.client_reqrep.ClientResponse to requests.Response."""
        response = requests.Response()
        response._content = await client_response.read()
        response.status_code = client_response.status
        response.headers = CaseInsensitiveDict(client_response.headers)
        response.url = str(client_response.url)
        response.history = [await HttpTrigger._convert_response(h) for h in client_response.history]
        response.encoding = client_response.get_encoding()
        response.reason = str(client_response.reason)
        cookies = RequestsCookieJar()
        for k, v in client_response.cookies.items():
            cookies.set(k, v)
        response.cookies = cookies
        return response


class HttpSensorTrigger(BaseTrigger):
    """
    A trigger that fires when the request to a URL returns a non-404 status code.

    :param endpoint: The relative part of the full url
    :param http_conn_id: The HTTP Connection ID to run the sensor against
    :param method: The HTTP request method to use
    :param data: payload to be uploaded or aiohttp parameters
    :param headers: The HTTP headers to be added to the GET request
    :param extra_options: Additional kwargs to pass when creating a request.
        For example, ``run(json=obj)`` is passed as ``aiohttp.ClientSession().get(json=obj)``
    :param poke_interval: Time to sleep using asyncio
    """

    def __init__(
        self,
        endpoint: str | None = None,
        http_conn_id: str = "http_default",
        method: str = "GET",
        data: dict[str, Any] | str | None = None,
        headers: dict[str, str] | None = None,
        extra_options: dict[str, Any] | None = None,
        poke_interval: float = 5.0,
    ):
        super().__init__()
        self.endpoint = endpoint
        self.method = method
        self.data = data
        self.headers = headers
        self.extra_options = extra_options or {}
        self.http_conn_id = http_conn_id
        self.poke_interval = poke_interval

    def serialize(self) -> tuple[str, dict[str, Any]]:
        """Serialize HttpTrigger arguments and classpath."""
        return (
            "airflow.providers.http.triggers.http.HttpSensorTrigger",
            {
                "endpoint": self.endpoint,
                "data": self.data,
                "method": self.method,
                "headers": self.headers,
                "extra_options": self.extra_options,
                "http_conn_id": self.http_conn_id,
                "poke_interval": self.poke_interval,
            },
        )

    async def run(self) -> AsyncIterator[TriggerEvent]:
        """Make a series of asynchronous http calls via an http hook."""
        hook = self._get_async_hook()
        while True:
            try:
                async with aiohttp.ClientSession() as session:
                    await hook.run(
                        session=session,
                        endpoint=self.endpoint,
                        data=self.data,
                        headers=self.headers,
                        extra_options=self.extra_options,
                    )
                yield TriggerEvent(True)
                return
            except AirflowException as exc:
                if str(exc).startswith("404"):
                    await asyncio.sleep(self.poke_interval)

    def _get_async_hook(self) -> HttpAsyncHook:
        return HttpAsyncHook(
            method=self.method,
            http_conn_id=self.http_conn_id,
        )