This is an automated email from the ASF dual-hosted git repository.
jasonliu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1790ec41eea Refactor Common Queue Interface (#54651)
1790ec41eea is described below
commit 1790ec41eeaff6e33f78b39ad50b2bcd98107e67
Author: LIU ZHE YOU <[email protected]>
AuthorDate: Sun Aug 31 00:05:12 2025 +0800
Refactor Common Queue Interface (#54651)
* Replace queue URI with scheme for MessageQueueTrigger
* Refactor KafkaMessageQueueProvider
* Refactor SqsMessageQueueProvider
* Refactor RedisPubSubMessageQueueProvider
* Fix tests
- Add test_utils/common_msg_queue
- Fix unit tests for sqs, kafka, common messaging
- Fix integration test for redis
- Fix system test for redis
* Fix fixture import for unit tests
* fixup! Fix test_sqs
* Remove docs for deprecated URI
* Fix KafkaMessageQueueTrigger with scheme based usage
* Update example dags
* Fix doc nits
- remove "New scheme-based approach" comment
- remove redundent "the"
* Fix Kafka TestMessageQueueTrigger
---
.../tests_common/test_utils/common_msg_queue.py | 39 ++++++
providers/amazon/docs/message-queues/index.rst | 23 +---
.../src/airflow/providers/amazon/aws/queues/sqs.py | 30 ++++-
.../tests/unit/amazon/aws/queues/test_sqs.py | 17 +++
.../tests/unit/amazon/aws/triggers/test_sqs.py | 23 +++-
.../apache/kafka/docs/message-queues/index.rst | 42 +-----
.../airflow/providers/apache/kafka/queues/kafka.py | 38 ++----
.../providers/apache/kafka/triggers/msg_queue.py | 38 +-----
.../kafka/example_dag_message_queue_trigger.py | 3 +-
.../tests/unit/apache/kafka/queues/test_kafka.py | 13 ++
.../apache/kafka/triggers/test_await_message.py | 24 +++-
.../unit/apache/kafka/triggers/test_msg_queue.py | 28 ----
.../common/messaging/providers/base_provider.py | 13 ++
.../common/messaging/triggers/msg_queue.py | 76 +++++++++--
.../messaging/example_message_queue_trigger.py | 3 +-
.../common/messaging/triggers/test_msg_queue.py | 143 +++++++++++++++++++--
providers/redis/docs/message-queues.rst | 74 ++++++-----
.../src/airflow/providers/redis/queues/redis.py | 58 ++-------
.../queues/test_redis_pubsub_message_queue.py | 7 +-
.../redis/example_dag_message_queue_trigger.py | 2 +-
.../redis/tests/unit/redis/queues/test_redis.py | 64 +++------
21 files changed, 450 insertions(+), 308 deletions(-)
diff --git a/devel-common/src/tests_common/test_utils/common_msg_queue.py
b/devel-common/src/tests_common/test_utils/common_msg_queue.py
new file mode 100644
index 00000000000..c0b19eec7d1
--- /dev/null
+++ b/devel-common/src/tests_common/test_utils/common_msg_queue.py
@@ -0,0 +1,39 @@
+# 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 pytest
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
+
+from tests_common.test_utils.version_compat import
get_base_airflow_version_tuple
+
+
[email protected]
+def collect_queue_param_deprecation_warning():
+ """Collect deprecation warnings for queue parameter."""
+ with pytest.warns(
+ AirflowProviderDeprecationWarning,
+ match="The `queue` parameter is deprecated and will be removed in
future versions. Use the `scheme` parameter instead and pass configuration as
keyword arguments to `MessageQueueTrigger`.",
+ ):
+ yield
+
+
+mark_common_msg_queue_test = pytest.mark.skipif(
+ get_base_airflow_version_tuple() < (3, 0, 1),
reason="CommonMessageQueueTrigger Requires Airflow 3.0.1+"
+)
diff --git a/providers/amazon/docs/message-queues/index.rst
b/providers/amazon/docs/message-queues/index.rst
index 388be374c66..2d07a423afa 100644
--- a/providers/amazon/docs/message-queues/index.rst
+++ b/providers/amazon/docs/message-queues/index.rst
@@ -1,4 +1,4 @@
- .. Licensed to the Apache Software Foundation (ASF) under one
+.. 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
@@ -23,21 +23,12 @@ Amazon SQS Queue Provider
Implemented by
:class:`~airflow.providers.amazon.aws.queues.sqs.SqsMessageQueueProvider`
-The Amazon SQS Queue Provider is a message queue provider that uses
-Amazon Simple Queue Service (SQS) as the underlying message queue system.
-It allows you to send and receive messages using SQS queues in your Airflow
workflows.
-The provider supports both standard and FIFO queues, and it provides features
-such as message visibility timeout, message retention period, and dead-letter
queues.
-
-The queue must be matching this regex:
-.. exampleinclude:: /../src/airflow/providers/amazon/aws/queues/sqs.py
- :language: python
- :dedent: 0
- :start-after: [START queue_regexp]
- :end-before: [END queue_regexp]
+The Amazon SQS Queue Provider is a
:class:`~airflow.providers.common.messaging.providers.base_provider.BaseMessageQueueProvider`
that uses
+Amazon Simple Queue Service (SQS) as the underlying message queue system.
+It allows you to send and receive messages using SQS queues in your Airflow
workflows with
:class:`~airflow.providers.common.messaging.triggers.msg_queue.MessageQueueTrigger`
common message queue interface.
-The queue parameter is passed directly to ``sqs_queue`` parameter of the
underlying
-:class:`~airflow.providers.amazon.aws.triggers.sqs.SqsSensorTrigger` class,
and passes
-all the kwargs directly to the trigger constructor if added.
+.. include:: /../src/airflow/providers/amazon/aws/queues/sqs.py
+ :start-after: [START sqs_message_queue_provider_description]
+ :end-before: [END sqs_message_queue_provider_description]
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/queues/sqs.py
b/providers/amazon/src/airflow/providers/amazon/aws/queues/sqs.py
index aff5b462b7f..e7c0086a0a7 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/queues/sqs.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/queues/sqs.py
@@ -38,7 +38,35 @@ QUEUE_REGEXP =
r"^https://sqs\.[^.]+\.amazonaws\.com/[0-9]+/.+";
class SqsMessageQueueProvider(BaseMessageQueueProvider):
- """Configuration for SQS integration with common-messaging."""
+ """
+ Configuration for SQS integration with common-messaging.
+
+ [START sqs_message_queue_provider_description]
+
+ * It uses ``sqs`` as scheme for identifying SQS queues.
+ * For parameter definitions take a look at
:class:`~airflow.providers.amazon.aws.triggers.sqs.SqsSensorTrigger`.
+
+ .. code-block:: python
+
+ from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+ from airflow.sdk import Asset, AssetWatcher
+
+ trigger = MessageQueueTrigger(
+ scheme="sqs",
+ # Additional AWS SqsSensorTrigger parameters as needed
+
sqs_queue="https://sqs.us-east-1.amazonaws.com/123456789012/my-queue";,
+ aws_conn_id="aws_default",
+ )
+
+ asset = Asset("sqs_queue_asset",
watchers=[AssetWatcher(name="sqs_watcher", trigger=trigger)])
+
+ For a complete example, see:
+ :mod:`tests.system.amazon.aws.example_dag_sqs_message_queue_trigger`
+
+ [END sqs_message_queue_provider_description]
+ """
+
+ scheme = "sqs"
def queue_matches(self, queue: str) -> bool:
return bool(re.match(QUEUE_REGEXP, queue))
diff --git a/providers/amazon/tests/unit/amazon/aws/queues/test_sqs.py
b/providers/amazon/tests/unit/amazon/aws/queues/test_sqs.py
index 457c61c7d87..103a38682f7 100644
--- a/providers/amazon/tests/unit/amazon/aws/queues/test_sqs.py
+++ b/providers/amazon/tests/unit/amazon/aws/queues/test_sqs.py
@@ -41,6 +41,23 @@ def test_message_sqs_queue_matches():
assert not provider.queue_matches("https://sqs.us-east-1.amazonaws.com/";)
[email protected](
+ "scheme, expected_result",
+ [
+ pytest.param("sqs", True, id="sqs_scheme"),
+ pytest.param("kafka", False, id="kafka_scheme"),
+ pytest.param("redis+pubsub", False, id="redis_scheme"),
+ pytest.param("unknown", False, id="unknown_scheme"),
+ ],
+)
+def test_message_sqs_scheme_matches(scheme, expected_result):
+ """Test the scheme_matches method with various schemes."""
+ from airflow.providers.amazon.aws.queues.sqs import SqsMessageQueueProvider
+
+ provider = SqsMessageQueueProvider()
+ assert provider.scheme_matches(scheme) == expected_result
+
+
def test_message_sqs_queue_trigger_class():
from airflow.providers.amazon.aws.queues.sqs import SqsMessageQueueProvider
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_sqs.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_sqs.py
index 11ea1c9ca0c..4522ecd05e3 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_sqs.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_sqs.py
@@ -20,7 +20,12 @@ from unittest.mock import AsyncMock
import pytest
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
get_base_airflow_version_tuple
+from tests_common.test_utils.common_msg_queue import (
+ collect_queue_param_deprecation_warning,
+ mark_common_msg_queue_test,
+)
+
+USED_FIXTURES = [collect_queue_param_deprecation_warning]
TEST_SQS_QUEUE = "test-sqs-queue"
TEST_AWS_CONN_ID = "test-aws-conn-id"
@@ -96,14 +101,22 @@ class TestSqsTriggers:
assert len(messages) == 0
[email protected](not AIRFLOW_V_3_0_PLUS, reason="Requires Airflow 3.0.+")
+@mark_common_msg_queue_test
class TestMessageQueueTrigger:
- def test_provider_integrations(self, cleanup_providers_manager):
- if get_base_airflow_version_tuple() < (3, 0, 1):
- pytest.skip("This test is only for Airflow 3.0.1+")
+ @pytest.mark.usefixtures("collect_queue_param_deprecation_warning")
+ def test_provider_integrations_with_queue_param(self,
cleanup_providers_manager):
queue = "https://sqs.us-east-1.amazonaws.com/0123456789/Test";
from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
trigger = MessageQueueTrigger(queue=queue)
assert isinstance(trigger.trigger, SqsSensorTrigger)
+
+ def test_provider_integrations_with_scheme_param(self,
cleanup_providers_manager):
+ from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
+ from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+
+ trigger = MessageQueueTrigger(
+ scheme="sqs",
sqs_queue="https://sqs.us-east-1.amazonaws.com/0123456789/Test";
+ )
+ assert isinstance(trigger.trigger, SqsSensorTrigger)
diff --git a/providers/apache/kafka/docs/message-queues/index.rst
b/providers/apache/kafka/docs/message-queues/index.rst
index 34785e34977..a676f2b18af 100644
--- a/providers/apache/kafka/docs/message-queues/index.rst
+++ b/providers/apache/kafka/docs/message-queues/index.rst
@@ -24,51 +24,20 @@
Apache Kafka Message Queue
==========================
-
-
Apache Kafka Queue Provider
---------------------------
Implemented by
:class:`~airflow.providers.apache.kafka.queues.kafka.KafkaMessageQueueProvider`
-The Apache Kafka Queue Provider is a message queue provider that uses
+The Apache Kafka Queue Provider is a
:class:`~airflow.providers.common.messaging.providers.base_provider.BaseMessageQueueProvider`
that uses
Apache Kafka as the underlying message queue system.
-It allows you to send and receive messages using Kafka topics in your Airflow
workflows.
-The provider supports Kafka topics and provides features for consuming and
processing
-messages from Kafka brokers.
-
-The queue must be matching this regex:
-
-.. exampleinclude:: /../src/airflow/providers/apache/kafka/queues/kafka.py
- :language: python
- :dedent: 0
- :start-after: [START queue_regexp]
- :end-before: [END queue_regexp]
-
-Queue URI Format:
-
-.. code-block:: text
+It allows you to send and receive messages using Kafka topics in your Airflow
workflows with
:class:`~airflow.providers.common.messaging.triggers.msg_queue.MessageQueueTrigger`
common message queue interface.
- kafka://<broker>/<topic_list>
-Where:
-
-- ``broker``: Kafka brokers (hostname:port)
-- ``topic_list``: Comma-separated list of Kafka topics to consume messages from
-
-The ``queue`` parameter is used to configure the underlying
-:class:`~airflow.providers.apache.kafka.triggers.await_message.AwaitMessageTrigger`
class and
-passes all kwargs directly to the trigger constructor, if provided.
-The ``apply_function`` kwarg is **required**.
-
-Topics can also be specified via the Queue URI instead of the ``topics``
kwarg. The provider will extract topics from the URI as follows:
-
-.. exampleinclude:: /../src/airflow/providers/apache/kafka/queues/kafka.py
- :language: python
- :dedent: 0
- :start-after: [START extract_topics]
- :end-before: [END extract_topics]
+.. include:: /../src/airflow/providers/apache/kafka/queues/kafka.py
+ :start-after: [START kafka_message_queue_provider_description]
+ :end-before: [END kafka_message_queue_provider_description]
.. _howto/triggers:KafkaMessageQueueTrigger:
@@ -93,6 +62,7 @@ Below is an example of how you can configure an Airflow DAG
to be triggered by a
How it works
------------
+
1. **Kafka Message Queue Trigger**: The ``KafkaMessageQueueTrigger`` listens
for messages from Apache Kafka Topic(s).
2. **Asset and Watcher**: The ``Asset`` abstracts the external entity, the
Kafka queue in this example.
diff --git
a/providers/apache/kafka/src/airflow/providers/apache/kafka/queues/kafka.py
b/providers/apache/kafka/src/airflow/providers/apache/kafka/queues/kafka.py
index 45617544d69..c5f7cebaec5 100644
--- a/providers/apache/kafka/src/airflow/providers/apache/kafka/queues/kafka.py
+++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/queues/kafka.py
@@ -35,44 +35,34 @@ class KafkaMessageQueueProvider(BaseMessageQueueProvider):
"""
Configuration for Apache Kafka integration with common-messaging.
- It uses the ``kafka://`` URI scheme for identifying Kafka queues.
+ [START kafka_message_queue_provider_description]
- **URI Format**:
-
- .. code-block:: text
-
- kafka://<broker>/<topic_list>
-
- Where:
-
- * ``broker``: Kafka brokers (hostname:port)
- * ``topic_list``: Comma-separated list of Kafka topics to consume messages
from
-
- **Examples**:
-
- .. code-block:: text
-
- kafka://localhost:9092/my_topic
-
- **Required kwargs**:
-
- * ``apply_function``: Function to process each Kafka message
-
- You can also provide ``topics`` directly in kwargs instead of in the URI.
+ * It uses ``kafka`` as scheme for identifying Kafka queues.
+ * For parameter definitions take a look at
:class:`~airflow.providers.apache.kafka.triggers.await_message.AwaitMessageTrigger`.
.. code-block:: python
from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+ from airflow.sdk import Asset, AssetWatcher
trigger = MessageQueueTrigger(
- queue="kafka://localhost:9092/test",
+ scheme="kafka",
+ # Additional Kafka AwaitMessageTrigger parameters as needed
+ topics=["my_topic"],
apply_function="module.apply_function",
+ bootstrap_servers="localhost:9092",
)
+ asset = Asset("kafka_queue_asset",
watchers=[AssetWatcher(name="kafka_watcher", trigger=trigger)])
+
For a complete example, see:
:mod:`tests.system.common.messaging.kafka_message_queue_trigger`
+
+ [END kafka_message_queue_provider_description]
"""
+ scheme = "kafka"
+
def queue_matches(self, queue: str) -> bool:
return bool(re.match(QUEUE_REGEXP, queue))
diff --git
a/providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/msg_queue.py
b/providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/msg_queue.py
index 4a36fa2e448..c28b5335f8c 100644
---
a/providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/msg_queue.py
+++
b/providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/msg_queue.py
@@ -31,7 +31,7 @@ class KafkaMessageQueueTrigger(MessageQueueTrigger):
for Kafka message queue operations while leveraging the unified messaging
framework.
:param topics: The topic (or topic regex) that should be searched for
messages
- :param kafka_config_id: The Kafka queue identifier in the format
kafka://<broker>/<topic_list>
+ :param kafka_config_id: The Kafka connection configuration ID, defaults to
"kafka_default"
:param apply_function: the location of the function to apply to messages
for determination of matching
criteria. (In python dot notation as a string)
:param apply_function_args: A set of arguments to apply to the callable,
defaults to None
@@ -56,10 +56,10 @@ class KafkaMessageQueueTrigger(MessageQueueTrigger):
poll_interval: float = 5,
**kwargs: Any,
) -> None:
- queue =
self.__class__.get_kafka_queue_uri(kafka_config_id=kafka_config_id,
topics=topics)
# Pass all required parameters to MessageQueueTrigger
super().__init__(
- queue=queue,
+ scheme="kafka",
+ topics=topics,
apply_function=apply_function,
kafka_config_id=kafka_config_id,
apply_function_args=apply_function_args or [],
@@ -68,35 +68,3 @@ class KafkaMessageQueueTrigger(MessageQueueTrigger):
poll_interval=poll_interval,
**kwargs,
)
-
- @classmethod
- def _get_brokers_from_connection(cls, kafka_config_id: str =
"kafka_default") -> str:
- """
- Get the brokers from the Kafka connection.
-
- :param kafka_config_id: The Kafka connection ID, defaults to
"kafka_default"
- :return: A string of brokers
- """
- from airflow.models.connection import Connection
-
- conn = Connection.get_connection_from_secrets(kafka_config_id)
- if not (brokers := conn.extra_dejson.get("bootstrap.servers", None)):
- raise ValueError("config['bootstrap.servers'] must be provided.")
- if isinstance(brokers, list):
- brokers = ",".join(brokers)
- return brokers.strip().replace(" ", "")
-
- @classmethod
- def get_kafka_queue_uri(cls, kafka_config_id: str, topics: Sequence[str])
-> str:
- """
- Generate a Kafka queue URI string from a Kafka configuration ID and a
list of topics.
-
- :param kafka_config_id: The Kafka connection configuration ID.
- :param topics: A sequence of topic names to include in the URI.
- :return: A formatted Kafka URI string in the format
"kafka://brokers/topics".
- """
- queue_topics = ",".join(topics).strip().replace(" ", "")
- if not queue_topics:
- raise ValueError("At least one valid topic must be provided.")
- queue_brokers = cls._get_brokers_from_connection(kafka_config_id)
- return f"kafka://{queue_brokers}/{queue_topics}"
diff --git
a/providers/apache/kafka/tests/system/apache/kafka/example_dag_message_queue_trigger.py
b/providers/apache/kafka/tests/system/apache/kafka/example_dag_message_queue_trigger.py
index 289ecb0961b..e07911b7298 100644
---
a/providers/apache/kafka/tests/system/apache/kafka/example_dag_message_queue_trigger.py
+++
b/providers/apache/kafka/tests/system/apache/kafka/example_dag_message_queue_trigger.py
@@ -32,7 +32,8 @@ def apply_function(message):
# Define a trigger that listens to an external message queue (Apache Kafka in
this case)
trigger = MessageQueueTrigger(
- queue="kafka://localhost:9092/test",
+ scheme="kafka",
+ topics=["test1, test2"],
apply_function="example_dag_message_queue_trigger.apply_function",
)
diff --git
a/providers/apache/kafka/tests/unit/apache/kafka/queues/test_kafka.py
b/providers/apache/kafka/tests/unit/apache/kafka/queues/test_kafka.py
index 0edd823d741..9c7ff192760 100644
--- a/providers/apache/kafka/tests/unit/apache/kafka/queues/test_kafka.py
+++ b/providers/apache/kafka/tests/unit/apache/kafka/queues/test_kafka.py
@@ -55,6 +55,19 @@ class TestKafkaMessageQueueProvider:
"""Test the queue_matches method with various URLs."""
assert self.provider.queue_matches(queue_uri) == expected_result
+ @pytest.mark.parametrize(
+ "scheme, expected_result",
+ [
+ pytest.param("kafka", True, id="kafka_scheme"),
+ pytest.param("redis+pubsub", False, id="redis_scheme"),
+ pytest.param("sqs", False, id="sqs_scheme"),
+ pytest.param("unknown", False, id="unknown_scheme"),
+ ],
+ )
+ def test_scheme_matches(self, scheme, expected_result):
+ """Test the scheme_matches method with various schemes."""
+ assert self.provider.scheme_matches(scheme) == expected_result
+
def test_trigger_class(self):
"""Test the trigger_class method."""
assert self.provider.trigger_class() == AwaitMessageTrigger
diff --git
a/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
b/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
index cd9380e5308..982089bb7c0 100644
---
a/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
+++
b/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
@@ -25,7 +25,12 @@ from airflow.models import Connection
from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
from airflow.providers.apache.kafka.triggers.await_message import
AwaitMessageTrigger
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
get_base_airflow_version_tuple
+from tests_common.test_utils.common_msg_queue import (
+ collect_queue_param_deprecation_warning,
+ mark_common_msg_queue_test,
+)
+
+USED_FIXTURES = [collect_queue_param_deprecation_warning]
def apply_function_false(message):
@@ -129,14 +134,21 @@ class TestTrigger:
asyncio.get_event_loop().stop()
[email protected](not AIRFLOW_V_3_0_PLUS, reason="Requires Airflow 3.0.+")
+@mark_common_msg_queue_test
class TestMessageQueueTrigger:
- def test_provider_integrations(self, cleanup_providers_manager):
- if get_base_airflow_version_tuple() < (3, 0, 1):
- pytest.skip("This test is only for Airflow 3.0.1+")
-
+ @pytest.mark.usefixtures("collect_queue_param_deprecation_warning")
+ def test_provider_integrations_with_queue_param(self,
cleanup_providers_manager):
queue = "kafka://localhost:9092/topic1"
+
from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
trigger = MessageQueueTrigger(queue=queue,
apply_function="mock_kafka_trigger_apply_function")
assert isinstance(trigger.trigger, AwaitMessageTrigger)
+
+ def test_provider_integrations_with_scheme_param(self,
cleanup_providers_manager):
+ from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+
+ trigger = MessageQueueTrigger(
+ scheme="kafka",
apply_function="mock_kafka_trigger_apply_function", topics=["topic1"]
+ )
+ assert isinstance(trigger.trigger, AwaitMessageTrigger)
diff --git
a/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_msg_queue.py
b/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_msg_queue.py
index 8b3ae123446..151b5f48d32 100644
--- a/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_msg_queue.py
+++ b/providers/apache/kafka/tests/unit/apache/kafka/triggers/test_msg_queue.py
@@ -84,34 +84,6 @@ class TestMessageQueueTrigger:
)
)
- @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Requires Airflow
3.0.+")
- @pytest.mark.parametrize(
- "kafka_config_id, topics, expected_uri",
- [
- ("kafka_d", ["test"], "kafka://localhost:9092/test"),
- ("kafka_d", ["test1", "test2"],
"kafka://localhost:9092/test1,test2"),
- (
- "kafka_multi_bootstraps",
- ["test1", "test2", "test3"],
-
"kafka://localhost:9091,localhost:9092,localhost:9093/test1,test2,test3",
- ),
- ],
- )
- def test_queue_uri(self, kafka_config_id, topics, expected_uri):
- """
- Test the queue URI generation for KafkaMessageQueueTrigger.
- """
- trigger = KafkaMessageQueueTrigger(
- kafka_config_id=kafka_config_id,
- topics=topics,
- apply_function="test.noop",
- apply_function_args=[],
- apply_function_kwargs={},
- poll_timeout=1,
- poll_interval=5,
- )
- assert trigger.queue == expected_uri
-
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Requires Airflow
3.0.+")
def test_trigger_serialization(self):
if get_base_airflow_version_tuple() < (3, 0, 1):
diff --git
a/providers/common/messaging/src/airflow/providers/common/messaging/providers/base_provider.py
b/providers/common/messaging/src/airflow/providers/common/messaging/providers/base_provider.py
index 5546e789ae8..505e1b647ec 100644
---
a/providers/common/messaging/src/airflow/providers/common/messaging/providers/base_provider.py
+++
b/providers/common/messaging/src/airflow/providers/common/messaging/providers/base_provider.py
@@ -31,6 +31,19 @@ class BaseMessageQueueProvider:
to ``MESSAGE_QUEUE_PROVIDERS``.
"""
+ scheme: str | None = None
+
+ def scheme_matches(self, scheme: str) -> bool:
+ """
+ Return whether a given scheme (string) matches a specific provider's
pattern.
+
+ This function must be as specific as possible to avoid collision with
other providers.
+ Functions in this provider should NOT overlap with each other in their
matching criteria.
+
+ :param scheme: The scheme identifier
+ """
+ return self.scheme == scheme
+
@abstractmethod
def queue_matches(self, queue: str) -> bool:
"""
diff --git
a/providers/common/messaging/src/airflow/providers/common/messaging/triggers/msg_queue.py
b/providers/common/messaging/src/airflow/providers/common/messaging/triggers/msg_queue.py
index 7dcfcc19bcd..5dd56333d25 100644
---
a/providers/common/messaging/src/airflow/providers/common/messaging/triggers/msg_queue.py
+++
b/providers/common/messaging/src/airflow/providers/common/messaging/triggers/msg_queue.py
@@ -18,10 +18,12 @@
from __future__ import annotations
import importlib
+import warnings
from collections.abc import AsyncIterator
from functools import cached_property
from typing import Any
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers_manager import ProvidersManager
from airflow.triggers.base import BaseEventTrigger, TriggerEvent
@@ -47,15 +49,36 @@ class MessageQueueTrigger(BaseEventTrigger):
This makes it easy to switch providers without modifying the trigger.
- :param queue: The queue identifier
+ :param scheme: The queue scheme (e.g., 'kafka', 'redis+pubsub', 'sqs').
Used for provider matching.
+ :param queue: **Deprecated** The queue identifier (URI format). If
provided, this takes precedence over scheme parameter.
+ This parameter is deprecated and will be removed in future versions.
Use the 'scheme' parameter instead.
.. seealso::
For more information on how to use this trigger, take a look at the
guide:
:ref:`howto/trigger:MessageQueueTrigger`
"""
- def __init__(self, *, queue: str, **kwargs: Any) -> None:
- self.queue = queue
+ queue: str | None = None
+ scheme: str | None = None
+
+ def __init__(self, *, queue: str | None = None, scheme: str | None = None,
**kwargs: Any) -> None:
+ if queue is None and scheme is None:
+ raise ValueError("Either `queue` or `scheme` parameter must be
provided.")
+
+ # For backward compatibility, queue takes precedence
+ if queue is not None:
+ warnings.warn(
+ "The `queue` parameter is deprecated and will be removed in
future versions. "
+ "Use the `scheme` parameter instead and pass configuration as
keyword arguments to `MessageQueueTrigger`.",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ self.queue = queue
+ self.scheme = None
+ else:
+ self.queue = None
+ self.scheme = scheme
+
self.kwargs = kwargs
@cached_property
@@ -66,27 +89,54 @@ class MessageQueueTrigger(BaseEventTrigger):
"Please ensure that you have the necessary providers
installed."
)
raise ValueError("No message queue providers are available. ")
- providers = [provider for provider in MESSAGE_QUEUE_PROVIDERS if
provider.queue_matches(self.queue)]
+
+ # Find matching providers based on queue URI or scheme
+ if self.queue is not None:
+ # Use existing queue-based matching for backward compatibility
+ providers = [
+ provider for provider in MESSAGE_QUEUE_PROVIDERS if
provider.queue_matches(self.queue)
+ ]
+ identifier = self.queue
+ match_by = "queue"
+ elif self.scheme is not None:
+ # Use new scheme-based matching
+ providers = [
+ provider for provider in MESSAGE_QUEUE_PROVIDERS if
provider.scheme_matches(self.scheme)
+ ]
+ identifier = self.scheme
+ match_by = "scheme"
+
if len(providers) == 0:
self.log.error(
- "The queue '%s' is not recognized by any of the registered
providers. "
+ "The %s '%s' is not recognized by any of the registered
providers. "
"The available providers are: '%s'.",
- self.queue,
+ match_by,
+ identifier,
", ".join([type(provider).__name__ for provider in
MESSAGE_QUEUE_PROVIDERS]),
)
- raise ValueError("The queue is not recognized by any of the
registered providers.")
+ raise ValueError(
+ f"The {match_by} '{identifier}' is not recognized by any of
the registered providers."
+ )
+
if len(providers) > 1:
self.log.error(
- "The queue '%s' is recognized by more than one provider. "
+ "The %s '%s' is recognized by more than one provider. "
"At least two providers in ``MESSAGE_QUEUE_PROVIDERS`` are
colliding with each "
"other: '%s'",
- self.queue,
+ match_by,
+ identifier,
", ".join([type(provider).__name__ for provider in providers]),
)
- raise ValueError(f"The queue '{self.queue}' is recognized by more
than one provider.")
- return providers[0].trigger_class()(
- **providers[0].trigger_kwargs(self.queue, **self.kwargs),
**self.kwargs
- )
+ raise ValueError(f"The {match_by} '{identifier}' is recognized by
more than one provider.")
+
+ # Create trigger instance
+ selected_provider = providers[0]
+ if self.queue is not None:
+ # Pass queue to trigger_kwargs for backward compatibility
+ trigger_kwargs = selected_provider.trigger_kwargs(self.queue,
**self.kwargs)
+ return selected_provider.trigger_class()(**trigger_kwargs,
**self.kwargs)
+ # For scheme-based matching, we need to pass all current kwargs to the
trigger
+ return selected_provider.trigger_class()(**self.kwargs)
def serialize(self) -> tuple[str, dict[str, Any]]:
return self.trigger.serialize()
diff --git
a/providers/common/messaging/tests/system/common/messaging/example_message_queue_trigger.py
b/providers/common/messaging/tests/system/common/messaging/example_message_queue_trigger.py
index 7eba4dc35aa..6e2ec6e7aef 100644
---
a/providers/common/messaging/tests/system/common/messaging/example_message_queue_trigger.py
+++
b/providers/common/messaging/tests/system/common/messaging/example_message_queue_trigger.py
@@ -22,8 +22,7 @@ from airflow.providers.standard.operators.empty import
EmptyOperator
from airflow.sdk import DAG, Asset, AssetWatcher
# Define a trigger that listens to an external message queue (AWS SQS in this
case)
-trigger =
MessageQueueTrigger(queue="https://sqs.us-east-1.amazonaws.com/0123456789/my-queue";)
-
+trigger = MessageQueueTrigger(scheme="sqs",
sqs_queue="https://sqs.us-east-1.amazonaws.com/0123456789/Test";)
# Define an asset that watches for messages on the queue
asset = Asset("sqs_queue_asset", watchers=[AssetWatcher(name="sqs_watcher",
trigger=trigger)])
diff --git
a/providers/common/messaging/tests/unit/common/messaging/triggers/test_msg_queue.py
b/providers/common/messaging/tests/unit/common/messaging/triggers/test_msg_queue.py
index 7fd4394c942..4ee43f027d2 100644
---
a/providers/common/messaging/tests/unit/common/messaging/triggers/test_msg_queue.py
+++
b/providers/common/messaging/tests/unit/common/messaging/triggers/test_msg_queue.py
@@ -21,21 +21,32 @@ from unittest.mock import MagicMock
import pytest
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
from airflow.triggers.base import BaseEventTrigger
+from tests_common.test_utils.common_msg_queue import (
+ collect_queue_param_deprecation_warning,
+)
+
+USED_FIXTURES = [collect_queue_param_deprecation_warning]
+
class MockProvider:
"""Mock provider for testing."""
- def __init__(self, name: str, pattern: str):
+ def __init__(self, name: str, pattern: str, scheme: str | None = None):
self.name = name
self.pattern = pattern
+ self.scheme = scheme
self.mock_trigger = MagicMock(spec=BaseEventTrigger)
def queue_matches(self, queue: str) -> bool:
return queue.startswith(self.pattern)
+ def scheme_matches(self, scheme: str) -> bool:
+ return scheme == self.scheme
+
def trigger_class(self):
return type(self.mock_trigger)
@@ -45,10 +56,12 @@ class MockProvider:
PROVIDER_1_NAME = "SQSMessageQueueProvider"
PROVIDER_1_PATTERN = "sqs://"
+PROVIDER_1_SCHEME = "sqs"
PROVIDER_1_QUEUE = "sqs://test-queue"
PROVIDER_2_NAME = "KafkaMessageQueueProvider"
PROVIDER_2_PATTERN = "kafka://"
PROVIDER_2_QUEUE = "kafka://test-queue"
+PROVIDER_2_SCHEME = "kafka"
UNKNOWN_QUEUE = "unknown://queue"
UNSUPPORTED_QUEUE = "unsupported://queue"
@@ -56,13 +69,14 @@ TEST_QUEUE = "test://queue"
NO_PROVIDERS_ERROR = "No message queue providers are available"
INSTALL_PROVIDERS_MESSAGE = "Please ensure that you have the necessary
providers installed"
-NOT_RECOGNIZED_ERROR = "The queue is not recognized by any of the registered
providers"
-MULTIPLE_PROVIDERS_ERROR = "The queue '{queue}' is recognized by more than one
provider"
+NOT_RECOGNIZED_ERROR = "The {match_type} '{queue_name}' is not recognized by
any of the registered providers"
+MULTIPLE_PROVIDERS_ERROR = "The {match_type} '{queue}' is recognized by more
than one provider"
MESSAGE_QUEUE_PROVIDERS_PATH =
"airflow.providers.common.messaging.triggers.msg_queue.MESSAGE_QUEUE_PROVIDERS"
[email protected]("collect_queue_param_deprecation_warning")
class TestMessageQueueTrigger:
"""Test cases for MessageQueueTrigger error handling and provider
matching."""
@@ -82,7 +96,9 @@ class TestMessageQueueTrigger:
with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider1, provider2]):
trigger = MessageQueueTrigger(queue=UNKNOWN_QUEUE)
- with pytest.raises(ValueError, match=NOT_RECOGNIZED_ERROR):
+ with pytest.raises(
+ ValueError,
match=NOT_RECOGNIZED_ERROR.format(queue_name=UNKNOWN_QUEUE, match_type="queue")
+ ):
_ = trigger.trigger
def test_queue_recognized_by_multiple_providers(self):
@@ -94,7 +110,9 @@ class TestMessageQueueTrigger:
with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider1, provider2]):
trigger = MessageQueueTrigger(queue=PROVIDER_1_QUEUE)
- with pytest.raises(ValueError,
match=MULTIPLE_PROVIDERS_ERROR.format(queue=PROVIDER_1_QUEUE)):
+ with pytest.raises(
+ ValueError,
match=MULTIPLE_PROVIDERS_ERROR.format(queue=PROVIDER_1_QUEUE,
match_type="queue")
+ ):
_ = trigger.trigger
def test_successful_provider_matching(self):
@@ -191,10 +209,111 @@ class TestMessageQueueTrigger:
assert event is not None
[email protected](
-
"airflow.providers.common.messaging.triggers.msg_queue.MESSAGE_QUEUE_PROVIDERS",
- new_callable=mock.PropertyMock,
-)
-def test_provider_integrations(_):
- trigger = MessageQueueTrigger(queue="any queue")
- assert trigger is not None
+class TestMessageQueueTriggerScheme:
+ """Test cases for the new scheme parameter functionality."""
+
+ @pytest.mark.usefixtures("collect_queue_param_deprecation_warning")
+ def test_queue_takes_precedence_over_scheme(self):
+ """Test that queue parameter takes precedence when both are
provided."""
+ trigger = MessageQueueTrigger(queue=PROVIDER_1_QUEUE,
scheme=PROVIDER_2_SCHEME)
+ assert trigger.queue == PROVIDER_1_QUEUE
+ assert trigger.scheme is None
+
+ def test_scheme_only_initialization(self):
+ """Test initialization with scheme parameter only."""
+ trigger = MessageQueueTrigger(scheme=PROVIDER_2_SCHEME)
+ assert trigger.queue is None
+ assert trigger.scheme == PROVIDER_2_SCHEME
+
+ def test_scheme_provider_matching(self):
+ """Test that scheme matching works correctly."""
+ provider1 = MockProvider(PROVIDER_1_NAME, PROVIDER_1_PATTERN,
scheme=PROVIDER_1_SCHEME)
+ provider2 = MockProvider(PROVIDER_2_NAME, PROVIDER_2_PATTERN,
scheme=PROVIDER_2_SCHEME)
+
+ with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider1, provider2]):
+ msg_queue = MessageQueueTrigger(scheme=PROVIDER_2_SCHEME)
+
+ result_trigger = msg_queue.trigger
+ assert result_trigger is not None
+
+ def test_scheme_not_recognized_by_any_provider(self):
+ """Test error when scheme is not recognized by any provider."""
+ provider1 = MockProvider(PROVIDER_1_NAME, PROVIDER_1_PATTERN,
scheme=PROVIDER_1_SCHEME)
+ provider2 = MockProvider(PROVIDER_2_NAME, PROVIDER_2_PATTERN,
scheme=PROVIDER_2_SCHEME)
+
+ with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider1, provider2]):
+ msg_queue = MessageQueueTrigger(scheme="unknown")
+
+ with pytest.raises(
+ ValueError,
match=NOT_RECOGNIZED_ERROR.format(queue_name="unknown", match_type="scheme")
+ ):
+ _ = msg_queue.trigger
+
+ def test_scheme_recognized_by_multiple_providers(self):
+ """Test error when scheme is recognized by multiple providers."""
+ provider1 = MockProvider(PROVIDER_1_NAME, PROVIDER_1_PATTERN,
scheme=PROVIDER_2_SCHEME)
+ provider2 = MockProvider(PROVIDER_2_NAME, PROVIDER_2_PATTERN,
scheme=PROVIDER_2_SCHEME)
+
+ with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider1, provider2]):
+ trigger = MessageQueueTrigger(scheme=PROVIDER_2_SCHEME)
+
+ with pytest.raises(
+ ValueError,
+ match=MULTIPLE_PROVIDERS_ERROR.format(queue=PROVIDER_2_SCHEME,
match_type="scheme"),
+ ):
+ _ = trigger.trigger
+
+ def test_scheme_kwargs_passed_correctly(self):
+ """Test that kwargs are passed correctly with scheme parameter."""
+ provider = MockProvider(PROVIDER_1_NAME, PROVIDER_1_PATTERN,
scheme=PROVIDER_1_SCHEME)
+
+ mock_trigger_class = MagicMock()
+ mock_trigger_instance = MagicMock(spec=BaseEventTrigger)
+ mock_trigger_class.return_value = mock_trigger_instance
+
+ provider.trigger_class = MagicMock(return_value=mock_trigger_class)
+ provider.trigger_kwargs = MagicMock()
+
+ with mock.patch(MESSAGE_QUEUE_PROVIDERS_PATH, [provider]):
+ msg_queue = MessageQueueTrigger(scheme=PROVIDER_1_SCHEME,
param1="value1", param2="value2")
+
+ result = msg_queue.trigger
+
+ provider.trigger_kwargs.assert_not_called()
+
+ # Verify trigger class was instantiated with combined kwargs
+ mock_trigger_class.assert_called_once_with(param1="value1",
param2="value2")
+
+ assert result == mock_trigger_instance
+
+
+class TestMessageQueueTriggerIntegration:
+ @pytest.mark.usefixtures("collect_queue_param_deprecation_warning")
+ @mock.patch(
+ MESSAGE_QUEUE_PROVIDERS_PATH,
+ new_callable=mock.PropertyMock,
+ )
+ def test_provider_integrations_with_queue(self, _):
+ trigger_by_queue = MessageQueueTrigger(queue="any queue")
+ assert trigger_by_queue is not None
+
+ @mock.patch(
+ MESSAGE_QUEUE_PROVIDERS_PATH,
+ new_callable=mock.PropertyMock,
+ )
+ def test_provider_integrations_with_scheme(self, _):
+ trigger_by_scheme = MessageQueueTrigger(scheme="any scheme")
+ assert trigger_by_scheme is not None
+
+ def test_provider_integrations_validation(self):
+ """Test that either queue or scheme parameter must be provided."""
+ with pytest.raises(ValueError, match="Either `queue` or `scheme`
parameter must be provided"):
+ MessageQueueTrigger()
+
+ def test_provider_integrations_deprecation_warning(self):
+ """Test that a deprecation warning is raised when using the 'queue'
parameter."""
+ with pytest.warns(
+ AirflowProviderDeprecationWarning,
+ match="The `queue` parameter is deprecated and will be removed in
future versions. Use the `scheme` parameter instead and pass configuration as
keyword arguments to `MessageQueueTrigger`.",
+ ):
+ MessageQueueTrigger(queue="any queue")
diff --git a/providers/redis/docs/message-queues.rst
b/providers/redis/docs/message-queues.rst
index 84614e9a1f4..c948fe0b343 100644
--- a/providers/redis/docs/message-queues.rst
+++ b/providers/redis/docs/message-queues.rst
@@ -1,4 +1,4 @@
- .. Licensed to the Apache Software Foundation (ASF) under one
+.. 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
@@ -15,9 +15,18 @@
specific language governing permissions and limitations
under the License.
+.. NOTE TO CONTRIBUTORS:
+ Please, only add notes to the Changelog just below the "Changelog" header
when there are some breaking changes
+ and you want to add an explanation to the users on how they are supposed to
deal with them.
+ The changelog is updated and maintained semi-automatically by release
manager.
+
+
Redis Message Queue
-===================
+====================
+.. contents::
+ :local:
+ :depth: 2
Redis Queue Provider
--------------------
@@ -25,48 +34,47 @@ Redis Queue Provider
Implemented by
:class:`~airflow.providers.redis.queues.redis.RedisPubSubMessageQueueProvider`
-The Redis Queue Provider is a message queue provider that uses
+The Redis Queue Provider is a
:class:`~airflow.providers.common.messaging.providers.base_provider.BaseMessageQueueProvider`
that uses
Redis as the underlying message queue system.
-It allows you to send and receive messages using Redis in your Airflow
workflows.
-The provider supports Redis channels.
+It allows you to send and receive messages using Redis channels in your
Airflow workflows with
:class:`~airflow.providers.common.messaging.triggers.msg_queue.MessageQueueTrigger`
common message queue interface.
-The queue must be matching this regex:
-.. exampleinclude::/../src/airflow/providers/redis/queues/redis.py
- :language: python
- :dedent: 0
- :start-after: [START queue_regexp]
- :end-before: [END queue_regexp]
+.. include:: /../src/airflow/providers/redis/queues/redis.py
+ :start-after: [START redis_message_queue_provider_description]
+ :end-before: [END redis_message_queue_provider_description]
-Queue URI Format:
-.. code-block:: text
+.. _howto/triggers:RedisMessageQueueTrigger:
- redis://<host>:<port>/<channel_list>
+Redis Message Queue Trigger
+---------------------------
-Where:
+Implemented by
:class:`~airflow.providers.redis.triggers.redis_await_message.AwaitMessageTrigger`
-- ``host``: Redis server hostname
-- ``port``: Redis server port
-- ``channel_list``: Comma-separated list of Redis channels to subscribe to
+Inherited from
:class:`~airflow.providers.common.messaging.triggers.msg_queue.MessageQueueTrigger`
-The ``queue`` parameter is used to configure the underlying
-:class:`~airflow.providers.redis.triggers.redis_await_message.AwaitMessageTrigger`
class and
-passes all kwargs directly to the trigger constructor, if provided.
+Wait for a message in a queue
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Channels can also be specified via the Queue URI instead of the ``channels``
kwarg. The provider will extract channels from the URI as follows:
+Below is an example of how you can configure an Airflow DAG to be triggered by
a message in Redis.
-.. exampleinclude:: /../src/airflow/providers/redis/queues/redis.py
- :language: python
- :dedent: 0
- :start-after: [START extract_channels]
- :end-before: [END extract_channels]
+.. exampleinclude:: /../tests/system/redis/example_dag_message_queue_trigger.py
+ :language: python
+ :start-after: [START howto_trigger_message_queue]
+ :end-before: [END howto_trigger_message_queue]
-Below is an example of how you can configure an Airflow DAG to be triggered by
a message in Redis.
+How it works
+------------
-.. exampleinclude:: /../tests/system/redis/example_dag_message_queue_trigger.py
- :language: python
- :dedent: 0
- :start-after: [START howto_trigger_message_queue]
- :end-before: [END howto_trigger_message_queue]
+1. **Redis Message Queue Trigger**: The ``AwaitMessageTrigger`` listens for
messages from Redis channel(s).
+
+2. **Asset and Watcher**: The ``Asset`` abstracts the external entity, the
Redis queue in this example.
+The ``AssetWatcher`` associate a trigger with a name. This name helps you
identify which trigger is associated to which
+asset.
+
+3. **Event-Driven DAG**: Instead of running on a fixed schedule, the DAG
executes when the asset receives an update
+(e.g., a new message in the queue).
+
+For how to use the trigger, refer to the documentation of the
+:ref:`Messaging Trigger <howto/trigger:MessageQueueTrigger>`
diff --git a/providers/redis/src/airflow/providers/redis/queues/redis.py
b/providers/redis/src/airflow/providers/redis/queues/redis.py
index b36cdc1ad5b..c6ec22b98bb 100644
--- a/providers/redis/src/airflow/providers/redis/queues/redis.py
+++ b/providers/redis/src/airflow/providers/redis/queues/redis.py
@@ -16,9 +16,7 @@
# under the License.
from __future__ import annotations
-import re
from typing import TYPE_CHECKING
-from urllib.parse import urlparse
from airflow.providers.common.messaging.providers.base_provider import
BaseMessageQueueProvider
from airflow.providers.redis.triggers.redis_await_message import
AwaitMessageTrigger
@@ -35,57 +33,29 @@ class
RedisPubSubMessageQueueProvider(BaseMessageQueueProvider):
"""
Configuration for Redis integration with common-messaging.
- It uses the ``redis+pubsub://`` URI scheme for identifying Redis queues.
+ [START redis_message_queue_provider_description]
- **URI Format**:
-
- .. code-block:: text
-
- redis+pubsub://<host>:<port>/<channel_list>
-
- Where:
-
- * ``host``: Redis server hostname
- * ``port``: Redis server port
- * ``channel_list``: Comma-separated list of Redis channels to subscribe to
-
- **Examples**:
-
- .. code-block:: text
-
- redis+pubsub://localhost:6379/my_channel
-
- You can also provide ``channels`` directly in kwargs instead of in the URI.
+ * It uses ``redis+pubsub`` as scheme for identifying Redis queues.
+ * For parameter definitions take a look at
:class:`~airflow.providers.redis.triggers.redis_await_message.AwaitMessageTrigger`.
.. code-block:: python
from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+ from airflow.sdk import Asset, AssetWatcher
+
+ trigger = MessageQueueTrigger(
+ scheme="redis+pubsub",
+ # Additional Redis AwaitMessageTrigger parameters as needed
+ channels=["my_channel"],
+ redis_conn_id="redis_default",
+ )
- trigger =
MessageQueueTrigger(queue="redis+pubsub://localhost:6379/test")
+ asset = Asset("redis_queue_asset",
watchers=[AssetWatcher(name="redis_watcher", trigger=trigger)])
- For a complete example, see:
- :mod:`tests.system.redis.example_dag_message_queue_trigger`
+ [END redis_message_queue_provider_description]
"""
- def queue_matches(self, queue: str) -> bool:
- return bool(re.match(QUEUE_REGEXP, queue))
+ scheme = "redis+pubsub"
def trigger_class(self) -> type[BaseEventTrigger]:
return AwaitMessageTrigger # type: ignore[return-value]
-
- def trigger_kwargs(self, queue: str, **kwargs) -> dict:
- # [START extract_channels]
- # Parse the queue URI
- parsed = urlparse(queue)
- # Extract channels (after host and port)
- # parsed.path starts with a '/', so strip it
- raw_channels = parsed.path.lstrip("/")
- channels = raw_channels.split(",") if raw_channels else []
- # [END extract_channels]
-
- if not channels and "channels" not in kwargs:
- raise ValueError(
- "channels is required in RedisPubSubMessageQueueProvider
kwargs or provide them in the queue URI"
- )
-
- return {} if "channels" in kwargs else {"channels": channels}
diff --git
a/providers/redis/tests/integration/redis/queues/test_redis_pubsub_message_queue.py
b/providers/redis/tests/integration/redis/queues/test_redis_pubsub_message_queue.py
index ae5902b50f6..1b77bf57ad3 100644
---
a/providers/redis/tests/integration/redis/queues/test_redis_pubsub_message_queue.py
+++
b/providers/redis/tests/integration/redis/queues/test_redis_pubsub_message_queue.py
@@ -52,9 +52,4 @@ class TestRedisPubSubMessageQueueProviderIntegration:
pubsub.unsubscribe(self.channel)
def test_queue_matches(self):
- assert
self.provider.queue_matches(f"redis+pubsub://localhost:6379/{self.channel}")
-
- def test_trigger_kwargs(self):
- uri = f"redis+pubsub://localhost:6379/{self.channel}"
- kwargs = self.provider.trigger_kwargs(uri)
- assert kwargs == {"channels": [self.channel]}
+ assert self.provider.scheme_matches("redis+pubsub")
diff --git
a/providers/redis/tests/system/redis/example_dag_message_queue_trigger.py
b/providers/redis/tests/system/redis/example_dag_message_queue_trigger.py
index b7089e0d4dc..9ef4bec0af1 100644
--- a/providers/redis/tests/system/redis/example_dag_message_queue_trigger.py
+++ b/providers/redis/tests/system/redis/example_dag_message_queue_trigger.py
@@ -22,7 +22,7 @@ from airflow.providers.standard.operators.empty import
EmptyOperator
from airflow.sdk import DAG, Asset, AssetWatcher
# Define a trigger that listens to an external message queue (Redis in this
case)
-trigger = MessageQueueTrigger(queue="redis+pubsub://localhost:6379/test")
+trigger = MessageQueueTrigger(scheme="redis+pubsub", channels=["test"])
# Define an asset that watches for messages on the queue
asset = Asset("redis_queue_asset_1",
watchers=[AssetWatcher(name="redis_watcher_1", trigger=trigger)])
diff --git a/providers/redis/tests/unit/redis/queues/test_redis.py
b/providers/redis/tests/unit/redis/queues/test_redis.py
index bf91fc35c4c..7cf8c6a92f8 100644
--- a/providers/redis/tests/unit/redis/queues/test_redis.py
+++ b/providers/redis/tests/unit/redis/queues/test_redis.py
@@ -20,6 +20,8 @@ import pytest
from airflow.providers.redis.triggers.redis_await_message import
AwaitMessageTrigger
+from tests_common.test_utils.common_msg_queue import mark_common_msg_queue_test
+
pytest.importorskip("airflow.providers.common.messaging.providers.base_provider")
@@ -39,57 +41,29 @@ class TestRedisPubSubMessageQueueProvider:
assert isinstance(self.provider, BaseMessageQueueProvider)
@pytest.mark.parametrize(
- "queue_uri, expected_result",
+ "scheme, expected_result",
[
- pytest.param("redis+pubsub://localhost:6379/channel1", True,
id="single_channel"),
- pytest.param("redis+pubsub://localhost:6379/channel1,channel2",
True, id="multiple_channels"),
- pytest.param("http://example.com";, False, id="http_url"),
- pytest.param("not-a-url", False, id="invalid_url"),
+ pytest.param("redis+pubsub", True, id="redis_pubsub_scheme"),
+ pytest.param("kafka", False, id="kafka_scheme"),
+ pytest.param("sqs", False, id="sqs_scheme"),
+ pytest.param("unknown", False, id="unknown_scheme"),
],
)
- def test_queue_matches(self, queue_uri, expected_result):
- """Test the queue_matches method with various URLs."""
- assert self.provider.queue_matches(queue_uri) == expected_result
+ def test_scheme_matches(self, scheme, expected_result):
+ """Test the scheme_matches method with various schemes."""
+ assert self.provider.scheme_matches(scheme) == expected_result
def test_trigger_class(self):
"""Test the trigger_class method."""
assert self.provider.trigger_class() == AwaitMessageTrigger
- @pytest.mark.parametrize(
- "queue_uri, extra_kwargs, expected_result",
- [
- pytest.param(
- "redis+pubsub://localhost:6379/channel1,channel2",
- {"channels": ["channel1", "channel2"]},
- {},
- id="channels_from_uri",
- ),
- pytest.param(
- "redis+pubsub://localhost:6379/",
- {"channels": ["channel1", "channel2"]},
- {},
- id="channels_from_kwargs",
- ),
- ],
- )
- def test_trigger_kwargs_valid_cases(self, queue_uri, extra_kwargs,
expected_result):
- """Test the trigger_kwargs method with valid parameters."""
- kwargs = self.provider.trigger_kwargs(queue_uri, **extra_kwargs)
- assert kwargs == expected_result
- @pytest.mark.parametrize(
- "queue_uri, extra_kwargs, expected_error, error_match",
- [
- pytest.param(
- "redis+pubsub://localhost:6379/",
- {},
- ValueError,
- "channels is required in RedisPubSubMessageQueueProvider
kwargs or provide them in the queue URI",
- id="missing_channels",
- ),
- ],
- )
- def test_trigger_kwargs_error_cases(self, queue_uri, extra_kwargs,
expected_error, error_match):
- """Test that trigger_kwargs raises appropriate errors with invalid
parameters."""
- with pytest.raises(expected_error, match=error_match):
- self.provider.trigger_kwargs(queue_uri, **extra_kwargs)
+@mark_common_msg_queue_test
+class TestMessageQueueTrigger:
+ @pytest.mark.usefixtures("cleanup_providers_manager")
+ def test_provider_integrations_with_scheme_param(self):
+ from airflow.providers.common.messaging.triggers.msg_queue import
MessageQueueTrigger
+ from airflow.providers.redis.triggers.redis_await_message import
AwaitMessageTrigger
+
+ trigger = MessageQueueTrigger(scheme="redis+pubsub",
channels="test_channel")
+ assert isinstance(trigger.trigger, AwaitMessageTrigger)
