This is an automated email from the ASF dual-hosted git repository.
Miretpl 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 68e1eb4557f Parse additional Kafka, Redis & SQS options in
`broker_transport_options` (#70228)
68e1eb4557f is described below
commit 68e1eb4557f670c78ef2b5b9abb34abe2691182b
Author: stephen-bracken <[email protected]>
AuthorDate: Wed Aug 19 16:50:06 2026 +0100
Parse additional Kafka, Redis & SQS options in `broker_transport_options`
(#70228)
* parse additional dict options in broker_transport_options
* Apply suggestions from code review
* remove formatting from `fetch_message_attributes` example
* change `AirflowException` to `ValueError`
Co-authored-by: Przemysław Mirowski
<[email protected]>
* update type hinting
* add predefined_queues test and config doc
* fix mypy and linting
* Update exception messages
Co-authored-by: Przemysław Mirowski
<[email protected]>
* fix linting
---------
Co-authored-by: Stephen Bracken <email-protected>
Co-authored-by: Przemysław Mirowski
<[email protected]>
---
generated/known_airflow_exceptions.txt | 2 +-
providers/celery/provider.yaml | 71 ++++++++++
.../providers/celery/executors/default_celery.py | 95 ++++++++++----
.../airflow/providers/celery/get_provider_info.py | 70 ++++++++++
.../unit/celery/executors/test_celery_executor.py | 144 ++++++++++++++++++++-
5 files changed, 351 insertions(+), 31 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index 6866f875992..acfd9ae3eb8 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -148,7 +148,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
diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml
index dfe7e2e4ed9..bdde5389f8d 100644
--- a/providers/celery/provider.yaml
+++ b/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: ~
diff --git
a/providers/celery/src/airflow/providers/celery/executors/default_celery.py
b/providers/celery/src/airflow/providers/celery/executors/default_celery.py
index 3af3aa7744e..c5cd8773d5f 100644
--- a/providers/celery/src/airflow/providers/celery/executors/default_celery.py
+++ b/providers/celery/src/airflow/providers/celery/executors/default_celery.py
@@ -23,12 +23,17 @@ import json
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 @@ except (ImportError, ModuleNotFoundError):
_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:
diff --git a/providers/celery/src/airflow/providers/celery/get_provider_info.py
b/providers/celery/src/airflow/providers/celery/get_provider_info.py
index 4a79cc58821..02c71ab36ac 100644
--- a/providers/celery/src/airflow/providers/celery/get_provider_info.py
+++ b/providers/celery/src/airflow/providers/celery/get_provider_info.py
@@ -292,6 +292,48 @@ 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,
@@ -299,6 +341,27 @@ def get_provider_info():
"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/backend [...]
"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,
+ },
},
},
},
diff --git
a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
index 28b08405902..5d030d269d0 100644
--- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
+++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
@@ -724,15 +724,147 @@ def
test_celery_executor_with_no_recommended_result_backend(caplog):
) in caplog.text
-@conf_vars({("celery_broker_transport_options", "sentinel_kwargs"):
'{"service_name": "mymaster"}'})
-def test_sentinel_kwargs_loaded_from_string():
+_dict_options_test_cases: list[tuple[str, str, dict]] = [
+ (
+ "client-config",
+ '{"connect_timeout": 5}',
+ {
+ "connect_timeout": 5,
+ },
+ ),
+ (
+ "fetch_message_attributes",
+ """
+ {
+ "MessageSystemAttributeNames": ["SenderId", "SentTimestamp"],
+ "MessageAttributeNames": ["S3MessageBodyKey"]
+ }
+ """,
+ {
+ "MessageSystemAttributeNames": ["SenderId", "SentTimestamp"],
+ "MessageAttributeNames": ["S3MessageBodyKey"],
+ },
+ ),
+ (
+ "predefined_exchanges",
+ """
+ {
+ "exchange-1": {
+ "arn": "arn:aws:sns:us-east-1:xxx:exchange-1",
+ "access_key_id": "a",
+ "secret_access_key": "b"
+ },
+ "exchange-2.fifo": {
+ "arn": "arn:aws:sns:us-east-1:xxx:exchange-2",
+ "access_key_id": "c",
+ "secret_access_key": "d"
+ }
+ }
+ """,
+ {
+ "exchange-1": {
+ "arn": "arn:aws:sns:us-east-1:xxx:exchange-1",
+ "access_key_id": "a",
+ "secret_access_key": "b",
+ },
+ "exchange-2.fifo": {
+ "arn": "arn:aws:sns:us-east-1:xxx:exchange-2",
+ "access_key_id": "c",
+ "secret_access_key": "d",
+ },
+ },
+ ),
+ (
+ "predefined_queues",
+ """
+ {
+ "queue-1": {
+ "url": "https://sqs.us-east-1.amazonaws.com/xxx/aaa";,
+ "access_key_id": "a",
+ "secret_access_key": "b",
+ "backoff_tasks": ["svc.tasks.tasks.task1"]
+ },
+ "queue-2.fifo": {
+ "url": "https://sqs.us-east-1.amazonaws.com/xxx/bbb.fifo";,
+ "access_key_id": "c",
+ "secret_access_key": "d"
+ }
+ }
+ """,
+ {
+ "queue-1": {
+ "url": "https://sqs.us-east-1.amazonaws.com/xxx/aaa";,
+ "access_key_id": "a",
+ "secret_access_key": "b",
+ "backoff_tasks": ["svc.tasks.tasks.task1"],
+ },
+ "queue-2.fifo": {
+ "url": "https://sqs.us-east-1.amazonaws.com/xxx/bbb.fifo";,
+ "access_key_id": "c",
+ "secret_access_key": "d",
+ },
+ },
+ ),
+ (
+ "queue_tags",
+ """
+ {
+ "Environment": "production",
+ "Team": "backend"
+ }
+ """,
+ {
+ "Environment": "production",
+ "Team": "backend",
+ },
+ ),
+ (
+ "sqs-creation-attributes",
+ """
+ {
+ "KmsMasterKeyId": "alias/aws/sqs"
+ }
+ """,
+ {
+ "KmsMasterKeyId": "alias/aws/sqs",
+ },
+ ),
+ (
+ "kafka_admin_config",
+ '{"sasl.username": "foo", "sasl.password": "bar"}',
+ {"sasl.username": "foo", "sasl.password": "bar"},
+ ),
+ ("kafka_common_config", '{"compression.type": "zstd"}',
{"compression.type": "zstd"}),
+ (
+ "kafka_consumer_config",
+ '{"group.id": "myconsumer"}',
+ {"group.id": "myconsumer"},
+ ),
+ (
+ "kafka_producer_config",
+ '{"ssl.certificate.location": "/foo/bar"}',
+ {"ssl.certificate.location": "/foo/bar"},
+ ),
+ ("sentinel_kwargs", '{"service_name": "mymaster"}', {"service_name":
"mymaster"}),
+]
+
+
[email protected](
+ (
+ "option",
+ "value",
+ "expected",
+ ),
+ _dict_options_test_cases,
+ ids=[t[0] for t in _dict_options_test_cases],
+)
+def test_dict_options_loaded_from_string(option, value, expected):
import importlib
# Reload celery conf to apply the new config.
- importlib.reload(default_celery)
- assert
default_celery.DEFAULT_CELERY_CONFIG["broker_transport_options"]["sentinel_kwargs"]
== {
- "service_name": "mymaster"
- }
+ with conf_vars({("celery_broker_transport_options", option): value}):
+ importlib.reload(default_celery)
+ assert
default_celery.DEFAULT_CELERY_CONFIG["broker_transport_options"][option] ==
expected
@conf_vars({("celery", "task_acks_late"): "False"})
