Add dedicated module for the Agent interface (#17667)

Author: ofek

datadog_checks_base/changelog.d/17667.added

@@ -0,0 +1 @@
+Add dedicated module for the Agent interface

datadog_checks_base/datadog_checks/base/agent.py

@@ -0,0 +1,15 @@
+# (C) Datadog, Inc. 2024-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+try:
+    import aggregator
+    import datadog_agent
+
+    AGENT_RUNNING = True
+except ImportError:
+    from .stubs import aggregator, datadog_agent
+
+    AGENT_RUNNING = False
+
+
+__all__ = ['AGENT_RUNNING', 'aggregator', 'datadog_agent']

datadog_checks_base/datadog_checks/base/checks/base.py

@@ -30,6 +30,8 @@
 import yaml
 from six import PY2, binary_type, iteritems, raise_from, text_type
 
+from datadog_checks.base.agent import AGENT_RUNNING, aggregator, datadog_agent
+
 from ..config import is_affirmative
 from ..constants import ServiceCheck
 from ..errors import ConfigurationError
@@ -53,32 +55,34 @@
 from ..utils.tls import TlsContextWrapper
 from ..utils.tracing import traced_class
 
-try:
-    import datadog_agent
-
+if AGENT_RUNNING:
     from ..log import CheckLoggingAdapter, init_logging
 
-    init_logging()
-except ImportError:
-    from ..stubs import datadog_agent
+else:
     from ..stubs.log import CheckLoggingAdapter, init_logging
 
-    init_logging()
+init_logging()
 
-try:
-    import aggregator
+if datadog_agent.get_config('disable_unsafe_yaml'):
+    from ..ddyaml import monkey_patch_pyyaml
 
-    using_stub_aggregator = False
-except ImportError:
-    from ..stubs import aggregator
+    monkey_patch_pyyaml()
 
-    using_stub_aggregator = True
+if datadog_agent.get_config('integration_tracing'):
+    from ddtrace import patch
 
+    # handle thread monitoring as an additional option
+    # See: http://pypi.datadoghq.com/trace/docs/other_integrations.html#futures
+    if datadog_agent.get_config('integration_tracing_futures'):
+        patch(logging=True, requests=True, futures=True)
+    else:
+        patch(logging=True, requests=True)
 
-if datadog_agent.get_config('disable_unsafe_yaml'):
-    from ..ddyaml import monkey_patch_pyyaml
+if is_affirmative(datadog_agent.get_config('integration_profiling')):
+    from ddtrace.profiling import Profiler
 
-    monkey_patch_pyyaml()
+    prof = Profiler(service='datadog-agent-integrations')
+    prof.start()
 
 if not PY2:
     from pydantic import BaseModel, ValidationError
@@ -430,7 +434,7 @@ def metadata_manager(self):
         Used for sending metadata via Go bindings.
         """
         if not hasattr(self, '_metadata_manager'):
-            if not self.check_id and not using_stub_aggregator:
+            if not self.check_id and AGENT_RUNNING:
                 raise RuntimeError('Attribute `check_id` must be set')
 
             self._metadata_manager = MetadataManager(self.name, self.check_id, self.log, self.METADATA_TRANSFORMERS)
@@ -608,7 +612,7 @@ def submit_histogram_bucket(
             err_msg = 'Histogram: {} has non integer value: {}. Only integer are valid bucket values (count).'.format(
                 repr(name), repr(value)
             )
-            if using_stub_aggregator:
+            if not AGENT_RUNNING:
                 raise ValueError(err_msg)
             self.warning(err_msg)
             return
@@ -706,7 +710,7 @@ def _submit_metric(
             err_msg = 'Metric: {} has non float value: {}. Only float values can be submitted as metrics.'.format(
                 repr(name), repr(value)
             )
-            if using_stub_aggregator:
+            if not AGENT_RUNNING:
                 raise ValueError(err_msg)
             self.warning(err_msg)
             return

datadog_checks_base/datadog_checks/base/checks/kube_leader/mixins.py

@@ -1,11 +1,7 @@
 # (C) Datadog, Inc. 2018-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
-
-try:
-    import datadog_agent
-except ImportError:
-    from ...stubs import datadog_agent
+from datadog_checks.base.agent import datadog_agent
 
 from .. import AgentCheck
 from .record import ElectionRecordAnnotation, ElectionRecordLease

datadog_checks_base/datadog_checks/base/checks/openmetrics/mixins.py

@@ -14,18 +14,15 @@
 from prometheus_client.samples import Sample
 from six import PY3, iteritems, string_types
 
+from datadog_checks.base.agent import datadog_agent
+
 from ...config import is_affirmative
 from ...errors import CheckException
 from ...utils.common import to_native_string
 from ...utils.http import RequestsWrapper
 from .. import AgentCheck
 from ..libs.prometheus import text_fd_to_metric_families
 
-try:
-    import datadog_agent
-except ImportError:
-    from datadog_checks.base.stubs import datadog_agent
-
 if PY3:
     long = int
 

datadog_checks_base/datadog_checks/base/checks/openmetrics/v2/scraper.py

@@ -13,6 +13,8 @@
 from prometheus_client.parser import text_fd_to_metric_families as parse_prometheus
 from requests.exceptions import ConnectionError
 
+from datadog_checks.base.agent import datadog_agent
+
 from ....config import is_affirmative
 from ....constants import ServiceCheck
 from ....errors import ConfigurationError
@@ -22,11 +24,6 @@
 from .labels import LabelAggregator, get_label_normalizer
 from .transform import MetricTransformer
 
-try:
-    import datadog_agent
-except ImportError:
-    from datadog_checks.base.stubs import datadog_agent
-
 
 class OpenMetricsScraper:
     """

datadog_checks_base/datadog_checks/base/log.py

@@ -9,14 +9,11 @@
 from six import PY2, text_type
 from urllib3.exceptions import InsecureRequestWarning
 
+from datadog_checks.base.agent import datadog_agent
+
 from .utils.common import to_native_string
 from .utils.tracing import tracing_enabled
 
-try:
-    import datadog_agent
-except ImportError:
-    from .stubs import datadog_agent
-
 # Arbitrary number less than 10 (DEBUG)
 TRACE_LEVEL = 7
 

datadog_checks_base/datadog_checks/base/utils/__init__.py

@@ -1,29 +1,3 @@
 # (C) Datadog, Inc. 2018-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
-import os
-
-from ..config import is_affirmative
-
-try:
-    import datadog_agent
-
-    if datadog_agent.get_config('integration_tracing'):
-        from ddtrace import patch
-
-        # handle thread monitoring as an additional option
-        # See: http://pypi.datadoghq.com/trace/docs/other_integrations.html#futures
-        if datadog_agent.get_config('integration_tracing_futures'):
-            patch(logging=True, requests=True, futures=True)
-        else:
-            patch(logging=True, requests=True)
-
-    if is_affirmative(datadog_agent.get_config('integration_profiling')):
-        from ddtrace.profiling import Profiler
-
-        prof = Profiler(service='datadog-agent-integrations')
-        prof.start()
-
-except ImportError:
-    # Tracing & profiling Integrations is only available with Agent 6
-    pass

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -18,18 +18,14 @@
 from cachetools import TTLCache
 
 from datadog_checks.base import is_affirmative
+from datadog_checks.base.agent import datadog_agent
 from datadog_checks.base.log import get_check_logger
 from datadog_checks.base.utils.db.types import Transformer  # noqa: F401
 from datadog_checks.base.utils.serialization import json
 from datadog_checks.base.utils.tracing import INTEGRATION_TRACING_SERVICE_NAME, tracing_enabled
 
 from ..common import to_native_string
 
-try:
-    import datadog_agent
-except ImportError:
-    from ....stubs import datadog_agent
-
 logger = logging.getLogger(__file__)
 
 # AgentCheck methods to transformer name e.g. set_metadata -> metadata

datadog_checks_base/datadog_checks/base/utils/headers.py

@@ -5,10 +5,7 @@
 
 from six import iteritems
 
-try:
-    import datadog_agent
-except ImportError:
-    from ..stubs import datadog_agent
+from datadog_checks.base.agent import datadog_agent
 
 
 def _get_common_headers():

datadog_checks_base/datadog_checks/base/utils/http.py

@@ -23,6 +23,8 @@
 from six.moves.urllib.parse import quote, urlparse, urlunparse
 from wrapt import ObjectProxy
 
+from datadog_checks.base.agent import datadog_agent
+
 from ..config import is_affirmative
 from ..errors import ConfigurationError
 from .common import ensure_bytes, ensure_unicode
@@ -35,11 +37,6 @@
 except ImportError:
     from contextlib2 import ExitStack
 
-try:
-    import datadog_agent
-except ImportError:
-    from ..stubs import datadog_agent
-
 # Import lazily to reduce memory footprint and ease installation for development
 requests_aws = None
 requests_kerberos = None

datadog_checks_base/datadog_checks/base/utils/metadata/core.py

@@ -5,15 +5,12 @@
 
 from six import iteritems
 
+from datadog_checks.base.agent import datadog_agent
+
 from ..common import to_native_string
 from .utils import is_primitive
 from .version import parse_version
 
-try:
-    import datadog_agent
-except ImportError:
-    from ...stubs import datadog_agent
-
 LOGGER = logging.getLogger(__file__)
 
 

datadog_checks_base/datadog_checks/base/utils/replay/redirect.py

@@ -68,7 +68,7 @@ def log(self, level, *args, **kwargs):
         )
 
 
-base.using_stub_aggregator = False
+base.AGENT_RUNNING = True
 base.aggregator = ReplayAggregator()
 base.datadog_agent = core.datadog_agent = ReplayDatadogAgent()
 

datadog_checks_base/datadog_checks/base/utils/tracing.py

@@ -7,14 +7,11 @@
 
 from six import PY2, PY3
 
+from datadog_checks.base.agent import datadog_agent
+
 from ..config import is_affirmative
 from ..utils.common import to_native_string
 
-try:
-    import datadog_agent
-except ImportError:
-    from ..stubs import datadog_agent
-
 EXCLUDED_MODULES = ['threading']
 
 # During regular continuous tracing we trace only the check's top-level 'run' and

datadog_checks_base/tests/test_metadata.py

@@ -27,7 +27,7 @@ def test_default(self):
     def test_no_check_id_error(self):
         check = AgentCheck('test', {}, [{}])
 
-        with mock.patch('datadog_checks.base.checks.base.using_stub_aggregator', False):
+        with mock.patch('datadog_checks.base.checks.base.AGENT_RUNNING', True):
             with pytest.raises(RuntimeError):
                 check.set_metadata('foo', 'bar')