From d063b5ed7b65184b3fd1cb50ad0b69724c6ba155 Mon Sep 17 00:00:00 2001 From: agreenburg Date: Thu, 9 Nov 2023 11:09:08 -0500 Subject: [PATCH] Add cosmos/propagate_logs Airflow config support for disabling log propagation if desired (#648) Add Airflow config check for cosmos/propagate_logs to allow override of default propagation behavior. Expose entry-point so that Airflow can theoretically detect configuration default. Closes #639 ## Breaking Change? This is backward-compatible as it falls back to default behavior if the `cosmos` section or `propagate_logs` option don't exist. ## Checklist - [X] I have made corresponding changes to the documentation (if required) - [X] I have added tests that prove my fix is effective or that my feature works --------- Co-authored-by: Andrew Greenburg --- cosmos/__init__.py | 27 +++++++++++++++++++++++++++ cosmos/log.py | 5 +++++ docs/configuration/index.rst | 1 + docs/configuration/logging.rst | 19 +++++++++++++++++++ pyproject.toml | 3 +++ tests/test_log.py | 18 ++++++++++++++++++ 6 files changed, 73 insertions(+) create mode 100644 docs/configuration/logging.rst diff --git a/cosmos/__init__.py b/cosmos/__init__.py index 10fea5a1b..18f675750 100644 --- a/cosmos/__init__.py +++ b/cosmos/__init__.py @@ -116,3 +116,30 @@ "LoadMode", "TestBehavior", ] + +""" +Required provider info for using Airflow config for configuration +""" + + +def get_provider_info(): + return { + "package-name": "astronomer-cosmos", # Required + "name": "Astronomer Cosmos", # Required + "description": "Astronomer Cosmos is a library for rendering dbt workflows in Airflow. Contains dags, task groups, and operators.", # Required + "versions": [__version__], # Required + "config": { + "cosmos": { + "description": None, + "options": { + "propagate_logs": { + "description": "Enable log propagation from Cosmos custom logger\n", + "version_added": "1.3.0a1", + "type": "boolean", + "example": None, + "default": "True", + }, + }, + }, + }, + } diff --git a/cosmos/log.py b/cosmos/log.py index 052762153..e4ad6e2ba 100644 --- a/cosmos/log.py +++ b/cosmos/log.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from airflow.configuration import conf from airflow.utils.log.colored_log import CustomTTYColoredFormatter @@ -23,9 +24,13 @@ def get_logger(name: str | None = None) -> logging.Logger: By using this logger, we introduce a (yellow) astronomer-cosmos string into the project's log messages: [2023-08-09T14:20:55.532+0100] {subprocess.py:94} INFO - (astronomer-cosmos) - 13:20:55 Completed successfully """ + propagateLogs: bool = True + if conf.has_option("cosmos", "propagate_logs"): + propagateLogs = conf.getboolean("cosmos", "propagate_logs") logger = logging.getLogger(name) formatter: logging.Formatter = CustomTTYColoredFormatter(fmt=LOG_FORMAT) # type: ignore handler = logging.StreamHandler() handler.setFormatter(formatter) logger.addHandler(handler) + logger.propagate = propagateLogs return logger diff --git a/docs/configuration/index.rst b/docs/configuration/index.rst index 2c027e32d..8c282be03 100644 --- a/docs/configuration/index.rst +++ b/docs/configuration/index.rst @@ -21,3 +21,4 @@ Cosmos offers a number of configuration options to customize its behavior. For m Selecting & Excluding Operator Args Compiled SQL + Logging diff --git a/docs/configuration/logging.rst b/docs/configuration/logging.rst new file mode 100644 index 000000000..9c27a950c --- /dev/null +++ b/docs/configuration/logging.rst @@ -0,0 +1,19 @@ +.. _logging: + +Logging +==================== + +Cosmos uses a custom logger implementation so that all log messages are clearly tagged with ``(astronomer-cosmos)``. By default this logger has propagation enabled. + +In some environments (for example when running Celery workers) this can cause duplicated log messages to appear in the logs. In this case log propagation can be disabled via airflow configuration using the boolean option ``propagate_logs`` under a ``cosmos`` section. + +.. code-block:: cfg + + [cosmos] + propagate_logs = False + +or + +.. code-block:: python + + AIRFLOW__COSMOS__PROPAGATE_LOGS = "False" diff --git a/pyproject.toml b/pyproject.toml index cd90c8d96..0041d9488 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,9 @@ kubernetes = [ ] +[project.entry-points.cosmos] +provider_info = "cosmos:get_provider_info" + [project.urls] Homepage = "https://github.com/astronomer/astronomer-cosmos" Documentation = "https://astronomer.github.io/astronomer-cosmos" diff --git a/tests/test_log.py b/tests/test_log.py index d94949ed3..d145dbca4 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -1,6 +1,8 @@ import logging +from cosmos import get_provider_info from cosmos.log import get_logger +from airflow.configuration import conf def test_get_logger(): @@ -12,3 +14,19 @@ def test_get_logger(): assert custom_logger.propagate is True assert custom_logger.handlers[0].formatter.__class__.__name__ == "CustomTTYColoredFormatter" assert custom_string in custom_logger.handlers[0].formatter._fmt + + +def test_propagate_logs_conf(): + if not conf.has_section("cosmos"): + conf.add_section("cosmos") + conf.set("cosmos", "propagate_logs", "False") + custom_logger = get_logger("cosmos-log") + assert custom_logger.propagate is False + + +def test_get_provider_info(): + provider_info = get_provider_info() + assert "cosmos" in provider_info.get("config").keys() + assert "options" in provider_info.get("config").get("cosmos").keys() + assert "propagate_logs" in provider_info.get("config").get("cosmos").get("options").keys() + assert provider_info["config"]["cosmos"]["options"]["propagate_logs"]["type"] == "boolean"