Parse additional Kafka, Redis & SQS options in broker_transport_options

generated/known_airflow_exceptions.txt

@@ -149,7 +149,7 @@ providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py:
 providers/arangodb/src/airflow/providers/arangodb/hooks/arangodb.py::9
 providers/atlassian/jira/src/airflow/providers/atlassian/jira/hooks/jira.py::1
 providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py::2
-providers/celery/src/airflow/providers/celery/executors/default_celery.py::2
+providers/celery/src/airflow/providers/celery/executors/default_celery.py::1
 providers/celery/tests/integration/celery/test_celery_executor.py::2
 providers/cloudant/src/airflow/providers/cloudant/hooks/cloudant.py::2
 providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/backcompat/backwards_compat_converters.py::3

providers/celery/provider.yaml

@@ -460,6 +460,49 @@ config:
       Redis Sentinel as the result backend. See:
       https://docs.celeryq.dev/en/latest/userguide/configuration.html#std:setting-result_backend_transport_options
     options:
+      client-config:
+        description: |
+          SQS botocore client config.
+        version_added: "3.24.0"
+        type: string
+        example: '{"connect_timeout": 5}'
+        default: ~
+      fetch_message_attributes:
+        description: |
+          SQS fetch message attributes.
+        version_added: "3.24.0"
+        type: string
+        example: '{"MessageSystemAttributeNames": ["SenderId", "SentTimestamp"]}'
+        default: ~
+      kafka_admin_config:
+        description: |
+          Kafka admin only config, see ``kafka_common_config``.
+        version_added: "3.24.0"
+        type: string
+        example: ~
+        default: ~
+      kafka_common_config:
+        description: |
+          Kafka config common to producers and consumers,
+          see https://docs.celeryq.dev/en/stable/getting-started/backends-and-brokers/kafka.html.
+        version_added: "3.24.0"
+        type: string
+        example: '{"bootstrap.servers": "broker:9094"}'
+        default: ~
+      kafka_consumer_config:
+        description: |
+          Kafka consumer specific config, see ``kafka_common_config``.
+        version_added: "3.24.0"
+        type: string
+        example: ~
+        default: ~
+      kafka_producer_config:
+        description: |
+          Kafka producer specific config, see ``kafka_common_config``.
+        version_added: "3.24.0"
+        type: string
+        example: ~
+        default: ~
       master_name:
         description: |
           The name of the Redis Sentinel primary node to connect to.
@@ -468,6 +511,27 @@ config:
         type: string
         example: "mymaster"
         default: ~
+      predefined_exchanges:
+        description: |
+          SQS Predefined SNS topics.
+        version_added: "3.24.0"
+        type: string
+        example: '{"exchange-1": {"arn": "arn:aws:sns:us-east-1:xxx:exchange-1"}}'
+        default: ~
+      predefined_queues:
+        description:
+          SQS Predefined Queues.
+        version_added: "3.24.0"
+        type: string
+        example: '{"queue-1": {"url": "https://sqs.us-east-1.amazonaws.com/xxx/aaa"}}'
+        default: ~
+      queue_tags:
+        description: |
+          SQS Queue tags to apply during queue creation.
+        version_added: "3.24.0"
+        type: string
+        example: '{"Environment": "production", "Team": "backend"}'
+        default: ~
       sentinel_kwargs:
         description: |
           The sentinel_kwargs parameter allows passing additional options to the Sentinel client
@@ -482,3 +546,10 @@ config:
         sensitive: true
         example: '{"password": "password_for_redis_server"}'
         default: ~
+      sqs-creation-attributes:
+        description: |
+          Additional SQS queue attributes to apply as key/value pairs.
+        version_added: "3.24.0"
+        type: string
+        example: '{"KmsMasterKeyId": "alias/aws/sqs"}'
+        default: ~

providers/celery/src/airflow/providers/celery/executors/default_celery.py

@@ -23,12 +23,17 @@
 import logging
 import re
 import ssl
-from typing import Any
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowConfigException
 from airflow.providers.celery.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.providers.common.compat.sdk import AirflowException, conf
 
+if TYPE_CHECKING:
+    from typing import Any
+
+    from airflow.sdk.configuration import AirflowSDKConfigParser
+
 log = logging.getLogger(__name__)
 
 _USE_PSYCOPG3: bool
@@ -44,11 +49,70 @@
     _USE_PSYCOPG3 = False
 
 
+# broker_transport_options accessed as dict
+# e.g. https://github.com/celery/kombu/blob/4281680ef3a275a7d87433a703790251d9805803/kombu/transport/confluentkafka.py#L338
+_BROKER_TRANSPORT_DICT_OPTIONS = [
+    "client-config",
+    "fetch_message_attributes",
+    "kafka_admin_config",
+    "kafka_common_config",
+    "kafka_consumer_config",
+    "kafka_producer_config",
+    "predefined_exchanges",
+    "predefined_queues",
+    "queue_tags",
+    "sentinel_kwargs",
+    "sqs-creation-attributes",
+]
+
+
 def _broker_supports_visibility_timeout(url):
     return url.startswith(("redis://", "rediss://", "sqs://", "sentinel://"))
 
 
-def get_default_celery_config(team_conf) -> dict[str, Any]:
+def _broker_transport_options(broker_url: str, conf: AirflowSDKConfigParser | Any) -> dict[str, Any]:
+    """
+    Parse broker_transport_options including dict options.
+
+    :param broker_url: Celery broker url
+    :param conf: ExecutorConf object
+    :return: broker_transport_options dict
+    """
+    broker_transport_options: dict[str, str | int | float | Any] = (
+        conf.getsection("celery_broker_transport_options") or {}
+    )
+    if "visibility_timeout" not in broker_transport_options:
+        if _broker_supports_visibility_timeout(broker_url):
+            broker_transport_options["visibility_timeout"] = 86400
+            log.warning(
+                "No visibility_timeout configured in [celery_broker_transport_options]. "
+                "Using default of 86400 seconds (24 hours). Celery tasks running longer than this "
+                "will be redelivered by the broker, which terminates the original task. "
+                "If you have long-running tasks, increase this value in your Airflow configuration: "
+                "[celery_broker_transport_options] visibility_timeout = <seconds>"
+            )
+
+    # Parse dict options
+    for option in _BROKER_TRANSPORT_DICT_OPTIONS:
+        if option in broker_transport_options:
+            try:
+                option_value = broker_transport_options[option]
+                if not isinstance(option_value, str):
+                    raise ValueError(f"broker_transport_option {option} is not string: {option_value}")
+                option_json = json.loads(option_value)
+                if not isinstance(option_json, dict):
+                    raise ValueError(
+                        f"broker_transport_option {option} value is not dictionary: {option_json}"
+                    )
+                broker_transport_options[option] = option_json
+            except Exception as exc:
+                raise ValueError(
+                    f"Broker transport option {option} value should be written in the correct JSON format."
+                ) from exc
+    return broker_transport_options
+
+
+def get_default_celery_config(team_conf: AirflowSDKConfigParser | Any) -> dict[str, Any]:
     """
     Build Celery configuration using team-aware config.
 
@@ -65,26 +129,7 @@ def get_default_celery_config(team_conf) -> dict[str, Any]:
 
     broker_url = team_conf.get("celery", "BROKER_URL", fallback="redis://redis:6379/0")
 
-    broker_transport_options: dict = team_conf.getsection("celery_broker_transport_options") or {}
-    if "visibility_timeout" not in broker_transport_options:
-        if _broker_supports_visibility_timeout(broker_url):
-            broker_transport_options["visibility_timeout"] = 86400
-            log.warning(
-                "No visibility_timeout configured in [celery_broker_transport_options]. "
-                "Using default of 86400 seconds (24 hours). Celery tasks running longer than this "
-                "will be redelivered by the broker, which terminates the original task. "
-                "If you have long-running tasks, increase this value in your Airflow configuration: "
-                "[celery_broker_transport_options] visibility_timeout = <seconds>"
-            )
-
-    if "sentinel_kwargs" in broker_transport_options:
-        try:
-            sentinel_kwargs = json.loads(broker_transport_options["sentinel_kwargs"])
-            if not isinstance(sentinel_kwargs, dict):
-                raise ValueError
-            broker_transport_options["sentinel_kwargs"] = sentinel_kwargs
-        except Exception:
-            raise AirflowException("sentinel_kwargs should be written in the correct dictionary format.")
+    broker_transport_options = _broker_transport_options(broker_url, team_conf)
 
     if team_conf.has_option("celery", "RESULT_BACKEND"):
         result_backend = team_conf.get_mandatory_value("celery", "RESULT_BACKEND")
@@ -177,15 +222,17 @@ def get_default_celery_config(team_conf) -> dict[str, Any]:
                     "Set SSL_MUTUAL_TLS=True if you intend to use mutual TLS."
                 )
 
+            broker_use_ssl: dict[str, str | int] = {}
+
             if broker_url and re.search(r"amqps?://", broker_url):
-                broker_use_ssl = {"cert_reqs": ssl.CERT_REQUIRED}
+                broker_use_ssl["cert_reqs"] = ssl.CERT_REQUIRED
                 if ssl_cacert:
                     broker_use_ssl["ca_certs"] = ssl_cacert
                 if ssl_mutual_tls:
                     broker_use_ssl["keyfile"] = ssl_key
                     broker_use_ssl["certfile"] = ssl_cert
             elif broker_url and re.search("rediss?://|sentinel://", broker_url):
-                broker_use_ssl = {"ssl_cert_reqs": ssl.CERT_REQUIRED}
+                broker_use_ssl["ssl_cert_reqs"] = ssl.CERT_REQUIRED
                 if ssl_cacert:
                     broker_use_ssl["ssl_ca_certs"] = ssl_cacert
                 if ssl_mutual_tls:

providers/celery/src/airflow/providers/celery/get_provider_info.py

@@ -292,13 +292,76 @@ def get_provider_info():
             "celery_result_backend_transport_options": {
                 "description": "This section is for specifying options which can be passed to the\nunderlying celery result backend transport. This is particularly useful when using\nRedis Sentinel as the result backend. See:\nhttps://docs.celeryq.dev/en/latest/userguide/configuration.html#std:setting-result_backend_transport_options\n",
                 "options": {
+                    "client-config": {
+                        "description": "SQS botocore client config.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"connect_timeout": 5}',
+                        "default": None,
+                    },
+                    "fetch_message_attributes": {
+                        "description": "SQS fetch message attributes.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"MessageSystemAttributeNames": ["SenderId", "SentTimestamp"]}',
+                        "default": None,
+                    },
+                    "kafka_admin_config": {
+                        "description": "Kafka admin only config, see ``kafka_common_config``.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": None,
+                        "default": None,
+                    },
+                    "kafka_common_config": {
+                        "description": "Kafka config common to producers and consumers,\nsee https://docs.celeryq.dev/en/stable/getting-started/backends-and-brokers/kafka.html.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"bootstrap.servers": "broker:9094"}',
+                        "default": None,
+                    },
+                    "kafka_consumer_config": {
+                        "description": "Kafka consumer specific config, see ``kafka_common_config``.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": None,
+                        "default": None,
+                    },
+                    "kafka_producer_config": {
+                        "description": "Kafka producer specific config, see ``kafka_common_config``.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": None,
+                        "default": None,
+                    },
                     "master_name": {
                         "description": "The name of the Redis Sentinel primary node to connect to.\nRequired when using Redis Sentinel as the result backend.\n",
                         "version_added": None,
                         "type": "string",
                         "example": "mymaster",
                         "default": None,
                     },
+                    "predefined_exchanges": {
+                        "description": "SQS Predefined SNS topics.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"exchange-1": {"arn": "arn:aws:sns:us-east-1:xxx:exchange-1"}}',
+                        "default": None,
+                    },
+                    "predefined_queues": {
+                        "description": "SQS Predefined Queues.",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"queue-1": {"url": "https://sqs.us-east-1.amazonaws.com/xxx/aaa"}}',
+                        "default": None,
+                    },
+                    "queue_tags": {
+                        "description": "SQS Queue tags to apply during queue creation.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"Environment": "production", "Team": "backend"}',
+                        "default": None,
+                    },
                     "sentinel_kwargs": {
                         "description": "The sentinel_kwargs parameter allows passing additional options to the Sentinel client\nfor the result backend. In a typical scenario where Redis Sentinel is used as the\nresult backend and Redis servers are password-protected, the password needs to be\npassed through this parameter. Although its type is string, it is required to pass\na string that conforms to the dictionary format.\nSee:\nhttps://docs.celeryq.dev/en/stable/getting-started/backends-and-brokers/redis.html#configuration\n",
                         "version_added": None,
@@ -307,6 +370,13 @@ def get_provider_info():
                         "example": '{"password": "password_for_redis_server"}',
                         "default": None,
                     },
+                    "sqs-creation-attributes": {
+                        "description": "Additional SQS queue attributes to apply as key/value pairs.\n",
+                        "version_added": "3.24.0",
+                        "type": "string",
+                        "example": '{"KmsMasterKeyId": "alias/aws/sqs"}',
+                        "default": None,
+                    },
                 },
             },
         },