uranusjr commented on code in PR #70517:
URL: https://github.com/apache/airflow/pull/70517#discussion_r3800276430


##########
task-sdk/src/airflow/sdk/observability/metrics/statsd_logger.py:
##########
@@ -23,32 +23,49 @@
 from airflow.sdk._shared.observability.metrics import statsd_logger
 from airflow.sdk.configuration import conf
 
+log = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from airflow.sdk._shared.observability.metrics.statsd_logger import 
SafeStatsdLogger
 
-log = logging.getLogger(__name__)
-
 
 def get_statsd_logger() -> SafeStatsdLogger:
-    stats_class = conf.getimport("metrics", "statsd_custom_client_path", 
fallback=None)
+    # Local import to avoid requiring statsd when other backends are used 
(e.g. Datadog)
+    from statsd import StatsClient, UnixSocketStatsClient
 
-    # no need to check for the scheduler/statsd_on -> this method is only 
called when it is set
-    # and previously it would crash with None is callable if it was called 
without it.
-    from statsd import StatsClient
+    socket_path = conf.get("metrics", "statsd_socket_path", fallback=None)
+    custom_class = conf.getimport("metrics", "statsd_custom_client_path", 
fallback=None)
 
-    if stats_class:
-        if not issubclass(stats_class, StatsClient):
+    if custom_class is not None:
+        if socket_path is not None:
+            if not (isinstance(custom_class, type) and 
issubclass(custom_class, UnixSocketStatsClient)):
+                raise AirflowConfigException(
+                    "Your custom StatsD client must extend the 
statsd.UnixSocketStatsClient "
+                    "when using a socket path in order to ensure backwards 
compatibility."
+                )
+        elif not (
+            isinstance(custom_class, type) and issubclass(custom_class, 
(StatsClient, UnixSocketStatsClient))
+        ):
             raise AirflowConfigException(
-                "Your custom StatsD client must extend the statsd.StatsClient 
in order to ensure "
-                "backwards compatibility."
+                "Your custom StatsD client must extend the statsd.StatsClient 
or "
+                "statsd.UnixSocketStatsClient in order to ensure backwards 
compatibility."
             )
         log.info("Successfully loaded custom StatsD client")
 
-    else:
-        stats_class = StatsClient
-
-    return statsd_logger.get_statsd_logger(
-        stats_class=stats_class,
+    if socket_path is not None:
+        return statsd_logger.get_socket_statsd_logger(
+            stats_class=custom_class,
+            socket_path=socket_path,
+            prefix=conf.get("metrics", "statsd_prefix"),
+            influxdb_tags_enabled=conf.getboolean("metrics", 
"statsd_influxdb_enabled", fallback=False),
+            statsd_disabled_tags=conf.get("metrics", "statsd_disabled_tags", 
fallback=None),
+            metrics_allow_list=conf.get("metrics", "metrics_allow_list", 
fallback=None),
+            metrics_block_list=conf.get("metrics", "metrics_block_list", 
fallback=None),
+            stat_name_handler=conf.getimport("metrics", "stat_name_handler"),
+            statsd_influxdb_enabled=conf.getboolean("metrics", 
"statsd_influxdb_enabled", fallback=False),

Review Comment:
   Using the same relative import technique, we should be able to move all 
these into `get_socket_statsd_logger` instead so they are not repeated.
   
   Alternatively, at least do
   
   ```python
   logger_kwargs = {
       "prefix": conf.get("metrics", "statsd_prefix"),
       ...
   }
   if socket_path is None:
       return statsd_logger.get_udp_statsd_logger(**kwargs)
   return statsd_logger.get_socket_statsd_logger(socket_path=socket_path, 
**kwargs)
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to