diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index 8be3fafad4ac0..5716b3f767d3a 100644 --- a/generated/known_airflow_exceptions.txt +++ b/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 diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml index dfe7e2e4ed9e6..bdde5389f8dd4 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 3af3aa7744e66..c5cd8773d5f01 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 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 = " + ) + + # 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 = " - ) - - 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 4a79cc58821d5..02c71ab36ac47 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/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, + }, }, }, }, 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 28b0840590235..5d030d269d0ad 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"}), +] + + +@pytest.mark.parametrize( + ( + "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"})