Source code for airflow.providers.amazon.aws.utils.sqs

# 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 logging
from typing import Any

import jsonpath_ng
import jsonpath_ng.ext
from typing_extensions import Literal

[docs]log = logging.getLogger(__name__)


[docs]MessageFilteringType = Literal["literal", "jsonpath", "jsonpath-ext"]


[docs]def process_response(
    response: Any,
    message_filtering: MessageFilteringType | None = None,
    message_filtering_match_values: Any = None,
    message_filtering_config: Any = None,
) -> Any:
    """
    Process the response from SQS.

    :param response: The response from SQS
    :return: The processed response
    """
    if not isinstance(response, dict) or "Messages" not in response:
        return []

    messages = response["Messages"]
    num_messages = len(messages)
    log.info("Received %d messages", num_messages)

    if num_messages and message_filtering:
        messages = filter_messages(
            messages, message_filtering, message_filtering_match_values, message_filtering_config
        )
        num_messages = len(messages)
        log.info("There are %d messages left after filtering", num_messages)
    return messages


[docs]def filter_messages(
    messages, message_filtering, message_filtering_match_values, message_filtering_config
) -> list[Any]:
    if message_filtering == "literal":
        return filter_messages_literal(messages, message_filtering_match_values)
    if message_filtering == "jsonpath":
        return filter_messages_jsonpath(
            messages, message_filtering_match_values, message_filtering_config, jsonpath_ng.parse
        )
    if message_filtering == "jsonpath-ext":
        return filter_messages_jsonpath(
            messages, message_filtering_match_values, message_filtering_config, jsonpath_ng.ext.parse
        )
    else:
        raise NotImplementedError("Override this method to define custom filters")


[docs]def filter_messages_literal(messages, message_filtering_match_values) -> list[Any]:
    return [message for message in messages if message["Body"] in message_filtering_match_values]


[docs]def filter_messages_jsonpath(
    messages, message_filtering_match_values, message_filtering_config, parse
) -> list[Any]:
    jsonpath_expr = parse(message_filtering_config)
    filtered_messages = []
    for message in messages:
        body = message["Body"]
        # Body is a string, deserialize to an object and then parse
        body = json.loads(body)
        results = jsonpath_expr.find(body)
        if results and (
            message_filtering_match_values is None
            or any(result.value in message_filtering_match_values for result in results)
        ):
            filtered_messages.append(message)
    return filtered_messages