Skip to content

Commit

Permalink
Refactor SafeDogStatsdLogger to use get_validator to enable patte…
Browse files Browse the repository at this point in the history
…rn matching (#39370)

closes #39368

By reusing the get_validators function, we can incorporate pattern search to include allow/block lists.

(cherry picked from commit 0c4ffb8)
  • Loading branch information
rawwar authored and ephraimbuddy committed Jun 4, 2024
1 parent a74c52d commit 2d88cca
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 16 deletions.
16 changes: 2 additions & 14 deletions airflow/metrics/datadog_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from airflow.metrics.validators import (
AllowListValidator,
BlockListValidator,
get_validator,
validate_stat,
)

Expand Down Expand Up @@ -160,25 +161,12 @@ def get_dogstatsd_logger(cls) -> SafeDogStatsdLogger:
"""Get DataDog StatsD logger."""
from datadog import DogStatsd

metrics_validator: ListValidator

dogstatsd = DogStatsd(
host=conf.get("metrics", "statsd_host"),
port=conf.getint("metrics", "statsd_port"),
namespace=conf.get("metrics", "statsd_prefix"),
constant_tags=cls.get_constant_tags(),
)
if conf.get("metrics", "metrics_allow_list", fallback=None):
metrics_validator = AllowListValidator(conf.get("metrics", "metrics_allow_list"))
if conf.get("metrics", "metrics_block_list", fallback=None):
log.warning(
"Ignoring metrics_block_list as both metrics_allow_list "
"and metrics_block_list have been set"
)
elif conf.get("metrics", "metrics_block_list", fallback=None):
metrics_validator = BlockListValidator(conf.get("metrics", "metrics_block_list"))
else:
metrics_validator = AllowListValidator()
datadog_metrics_tags = conf.getboolean("metrics", "statsd_datadog_metrics_tags", fallback=True)
metric_tags_validator = BlockListValidator(conf.get("metrics", "statsd_disabled_tags", fallback=None))
return SafeDogStatsdLogger(dogstatsd, metrics_validator, datadog_metrics_tags, metric_tags_validator)
return SafeDogStatsdLogger(dogstatsd, get_validator(), datadog_metrics_tags, metric_tags_validator)
14 changes: 12 additions & 2 deletions tests/core/test_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,9 @@ def test_enabled_by_config(self):
"""Test that enabling this sets the right instance properties"""
from datadog import DogStatsd

with conf_vars({("metrics", "statsd_datadog_enabled"): "True"}):
with conf_vars(
{("metrics", "statsd_datadog_enabled"): "True", ("metrics", "metrics_use_pattern_match"): "True"}
):
importlib.reload(airflow.stats)
assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd)
assert not hasattr(airflow.stats.Stats, "statsd")
Expand All @@ -263,7 +265,13 @@ def test_enabled_by_config(self):
def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self):
from datadog import DogStatsd

with conf_vars({("metrics", "statsd_on"): "True", ("metrics", "statsd_datadog_enabled"): "True"}):
with conf_vars(
{
("metrics", "statsd_on"): "True",
("metrics", "statsd_datadog_enabled"): "True",
("metrics", "metrics_use_pattern_match"): "True",
}
):
importlib.reload(airflow.stats)
assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd)
assert not hasattr(airflow.stats.Stats, "statsd")
Expand Down Expand Up @@ -515,6 +523,7 @@ def test_does_not_send_stats_using_statsd_when_the_name_is_not_valid(self, mock_
@conf_vars(
{
("metrics", "statsd_datadog_enabled"): "True",
("metrics", "metrics_use_pattern_match"): "True",
("metrics", "stat_name_handler"): "tests.core.test_stats.always_invalid",
}
)
Expand All @@ -539,6 +548,7 @@ def test_does_send_stats_using_statsd_when_the_name_is_valid(self, mock_statsd):
@conf_vars(
{
("metrics", "statsd_datadog_enabled"): "True",
("metrics", "metrics_use_pattern_match"): "True",
("metrics", "stat_name_handler"): "tests.core.test_stats.always_valid",
}
)
Expand Down

0 comments on commit 2d88cca

Please sign in to comment.