From ab56d43258cf2dce4c68dd46a9d3fb4aca7ef5cd Mon Sep 17 00:00:00 2001 From: Julian Rojas Valvo Date: Tue, 5 Dec 2023 15:00:35 -0500 Subject: [PATCH 1/9] feat: added datadog monitoring support --- docker-compose.yml | 20 ++++++ querybook/server/__init__.py | 0 querybook/server/env.py | 10 +-- querybook/server/lib/celery/celery_stats.py | 29 ++++++++ querybook/server/lib/stats_logger/__init__.py | 4 ++ .../lib/stats_logger/all_stats_loggers.py | 12 +++- .../loggers/datadog_stats_logger.py | 68 +++++++++++++++++++ querybook/server/tasks/all_tasks.py | 45 ++++++++++-- requirements/datadog/datadog.txt | 2 + requirements/extra.txt | 3 + requirements/test.txt | 1 + 11 files changed, 184 insertions(+), 10 deletions(-) create mode 100644 querybook/server/__init__.py create mode 100644 querybook/server/lib/celery/celery_stats.py create mode 100644 querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py create mode 100644 requirements/datadog/datadog.txt diff --git a/docker-compose.yml b/docker-compose.yml index 1e36f3d86..af242addf 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -20,6 +20,26 @@ x-querybook-volumes: &querybook-volumes # - file:/opt/store/ services: + datadog: + image: gcr.io/datadoghq/agent:latest + ports: + - 8125:8125/udp + expose: + - 8125/udp + environment: + DD_DOGSTATSD_NON_LOCAL_TRAFFIC: true + DD_USE_DOGSTATSD: true + DD_APM_ENABLED: false + DD_AC_EXCLUDE: '.*' + DD_ENV: 'dev' + DD_SERVICE: 'federated-querybook' + DD_VERSION: 'latest' + env_file: + - ./.env.local + volumes: + - /var/run/docker.sock:/var/run/docker.sock + - /proc/:/host/proc/:ro + - /sys/fs/cgroup:/host/sys/fs/cgroup:ro web: container_name: querybook_web image: *querybook-image diff --git a/querybook/server/__init__.py b/querybook/server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/querybook/server/env.py b/querybook/server/env.py index b634adf15..537e2f983 100644 --- a/querybook/server/env.py +++ b/querybook/server/env.py @@ -135,7 +135,9 @@ class QuerybookSettings(object): AI_ASSISTANT_PROVIDER = get_env_config("AI_ASSISTANT_PROVIDER") AI_ASSISTANT_CONFIG = get_env_config("AI_ASSISTANT_CONFIG") or {} - VECTOR_STORE_PROVIDER = get_env_config("VECTOR_STORE_PROVIDER") - VECTOR_STORE_CONFIG = get_env_config("VECTOR_STORE_CONFIG") or {} - EMBEDDINGS_PROVIDER = get_env_config("EMBEDDINGS_PROVIDER") - EMBEDDINGS_CONFIG = get_env_config("EMBEDDINGS_CONFIG") or {} + # Datadog + DD_AGENT_HOST = get_env_config("DD_AGENT_HOST", optional=True) + DD_DOGSTATSD_PORT = int(get_env_config("DD_DOGSTATSD_PORT", optional=True) or 8125) + DD_PREFIX = get_env_config("DD_PREFIX", optional=True) + DD_SERVICE = get_env_config("DD_SERVICE", optional=True) or "querybook" + DD_TAGS = get_env_config("DD_TAGS", optional=True) or [] diff --git a/querybook/server/lib/celery/celery_stats.py b/querybook/server/lib/celery/celery_stats.py new file mode 100644 index 000000000..1059ac7e8 --- /dev/null +++ b/querybook/server/lib/celery/celery_stats.py @@ -0,0 +1,29 @@ +import time +import threading +from lib.stats_logger import stats_logger, ACTIVE_WORKERS, ACTIVE_TASKS + + +def send_stats_logger_metrics(celery): + while True: + i = celery.control.inspect() + + stats = i.stats() or {} + active = i.active() or {} + + active_workers = list(stats.keys()) + active_tasks = 0 + for worker in active_workers: + if worker in active: + if len(active[worker]) > 0: + active_tasks += 1 + + stats_logger.gauge(ACTIVE_WORKERS, len(stats)) + stats_logger.gauge(ACTIVE_TASKS, active_tasks) + time.sleep(5) + + +def start_stats_logger_monitor(celery): + thread = threading.Thread( + target=send_stats_logger_metrics, args=[celery], daemon=True + ) + thread.start() diff --git a/querybook/server/lib/stats_logger/__init__.py b/querybook/server/lib/stats_logger/__init__.py index 687043bc9..0a184d7f6 100644 --- a/querybook/server/lib/stats_logger/__init__.py +++ b/querybook/server/lib/stats_logger/__init__.py @@ -7,8 +7,12 @@ WS_CONNECTIONS = "ws.connections" SQL_SESSION_FAILURES = "sql_session.failures" TASK_FAILURES = "task.failures" +TASK_SUCCESSES = "task.successes" +TASK_RECEIVED = "task.received" REDIS_OPERATIONS = "redis.operations" QUERY_EXECUTIONS = "query.executions" +ACTIVE_WORKERS = "celery.active_workers" +ACTIVE_TASKS = "celery.active_tasks" logger_name = QuerybookSettings.STATS_LOGGER_NAME diff --git a/querybook/server/lib/stats_logger/all_stats_loggers.py b/querybook/server/lib/stats_logger/all_stats_loggers.py index 3f570786b..0fba53a03 100644 --- a/querybook/server/lib/stats_logger/all_stats_loggers.py +++ b/querybook/server/lib/stats_logger/all_stats_loggers.py @@ -1,4 +1,6 @@ from lib.utils.import_helper import import_module_with_default + +from .loggers.datadog_stats_logger import DatadogStatsLogger from .loggers.null_stats_logger import NullStatsLogger from .loggers.console_stats_logger import ConsoleStatsLogger @@ -8,11 +10,19 @@ default=[], ) -ALL_STATS_LOGGERS = [NullStatsLogger(), ConsoleStatsLogger()] + ALL_PLUGIN_STATS_LOGGERS +ALL_STATS_LOGGERS = [ + NullStatsLogger(), + ConsoleStatsLogger(), + DatadogStatsLogger(), +] + ALL_PLUGIN_STATS_LOGGERS def get_stats_logger_class(name: str): for logger in ALL_STATS_LOGGERS: if logger.logger_name == name: + if hasattr(logger, "initialize") and callable( + getattr(logger, "initialize") + ): + logger.initialize() return logger raise ValueError(f"Unknown stats logger name {name}") diff --git a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py new file mode 100644 index 000000000..add3e1bb4 --- /dev/null +++ b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py @@ -0,0 +1,68 @@ +from lib.stats_logger.base_stats_logger import BaseStatsLogger +from datadog import initialize, statsd +from env import QuerybookSettings +from lib.logger import get_logger + +LOG = get_logger(__file__) + + +class DatadogStatsLogger(BaseStatsLogger): + metric_prefix = "" + dd_tags = [] + + def metric_prefix_helper(self, key): + return self.metric_prefix + "." + key + + def tag_helper(self, tags): + if tags: + return [f"{k}:{v}" for k, v in tags.items()] + return [] + + def initialize(self): + if QuerybookSettings.DD_AGENT_HOST and QuerybookSettings.DD_DOGSTATSD_PORT: + LOG.info("Initializing Datadog") + + self.dd_tags = ( + QuerybookSettings.DD_TAGS.split(",") + if QuerybookSettings.DD_TAGS + else [] + ) + self.metric_prefix = ( + QuerybookSettings.DD_PREFIX or QuerybookSettings.DD_SERVICE + ) + + options = { + "statsd_host": QuerybookSettings.DD_AGENT_HOST, + "statsd_port": QuerybookSettings.DD_DOGSTATSD_PORT, + "statsd_constant_tags": self.dd_tags, + } + + initialize(**options) + else: + LOG.info("Datadog environment variables are not set") + + @property + def logger_name(self) -> str: + return "datadog" + + def incr(self, key: str, tags: dict[str, str] = None) -> None: + statsd.increment( + self.metric_prefix_helper(key), + 1, + tags=self.tag_helper(tags), + ) + + def decr(self, key: str, tags: dict[str, str] = None) -> None: + statsd.decrement( + self.metric_prefix_helper(key), + 1, + tags=self.tag_helper(tags), + ) + + def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None: + statsd.histogram( + self.metric_prefix_helper(key), value, tags=self.tag_helper(tags) + ) + + def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None: + statsd.set(self.metric_prefix_helper(key), value, tags=self.tag_helper(tags)) diff --git a/querybook/server/tasks/all_tasks.py b/querybook/server/tasks/all_tasks.py index 609eac5ea..41cae41e5 100644 --- a/querybook/server/tasks/all_tasks.py +++ b/querybook/server/tasks/all_tasks.py @@ -1,12 +1,25 @@ -from celery.signals import celeryd_init, task_failure +from celery.signals import ( + celeryd_init, + task_failure, + task_success, + task_received, + beat_init, +) from celery.utils.log import get_task_logger from importlib import import_module from app.flask_app import celery from env import QuerybookSettings from lib.logger import get_logger -from lib.stats_logger import TASK_FAILURES, stats_logger -from logic.schedule import get_schedule_task_type +from lib.stats_logger import ( + TASK_FAILURES, + TASK_SUCCESSES, + TASK_RECEIVED, + ACTIVE_WORKERS, + ACTIVE_TASKS, + stats_logger, +) +from lib.celery.celery_stats import start_stats_logger_monitor from .export_query_execution import export_query_execution_task from .run_query import run_query_task @@ -20,6 +33,7 @@ from .presto_hive_function_scrapper import presto_hive_function_scrapper from .db_clean_up_jobs import run_all_db_clean_up_jobs from .disable_scheduled_docs import disable_scheduled_docs +from logic.schedule import get_schedule_task_type LOG = get_logger(__file__) @@ -60,5 +74,26 @@ def configure_workers(sender=None, conf=None, **kwargs): @task_failure.connect def handle_task_failure(sender, signal, *args, **kwargs): - task_type = get_schedule_task_type(sender.name) - stats_logger.incr(TASK_FAILURES, tags={"task_type": task_type}) + tags = {"task_name": sender.name, "task_type": get_schedule_task_type(sender.name)} + stats_logger.incr(TASK_FAILURES, tags=tags) + + +@task_success.connect +def handle_task_success(sender, signal, *args, **kwargs): + tags = {"task_name": sender.name, "task_type": get_schedule_task_type(sender.name)} + stats_logger.incr(TASK_SUCCESSES, tags=tags) + + +@task_received.connect +def handle_task_received(sender, signal, *args, **kwargs): + tags = {"task_name": kwargs["request"].name} + tags["task_type"] = get_schedule_task_type(tags["task_name"]) + stats_logger.incr(TASK_RECEIVED, tags=tags) + + +@beat_init.connect +def start_celery_stats_logging(sender=None, conf=None, **kwargs): + LOG.info("Starting Celery Beat") + start_stats_logger_monitor(celery) + stats_logger.gauge(ACTIVE_WORKERS, 0) + stats_logger.gauge(ACTIVE_TASKS, 0) diff --git a/requirements/datadog/datadog.txt b/requirements/datadog/datadog.txt new file mode 100644 index 000000000..2b9866e0f --- /dev/null +++ b/requirements/datadog/datadog.txt @@ -0,0 +1,2 @@ +datadog==0.47.0 +ddtrace==2.0.2 diff --git a/requirements/extra.txt b/requirements/extra.txt index db5edd23e..e99892e88 100644 --- a/requirements/extra.txt +++ b/requirements/extra.txt @@ -28,3 +28,6 @@ # AI Assistant -r ai/langchain.txt + +# Datadog +-r datadog/datadog.txt diff --git a/requirements/test.txt b/requirements/test.txt index 01c2f4ce0..200423617 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -8,3 +8,4 @@ moto==2.2.4 -r exporter/gspread.txt -r parser/sqlglot.txt -r ai/langchain.txt +-r datadog/datadog.txt From 868eea3059e9176805a24872c2d3b65430abd4ff Mon Sep 17 00:00:00 2001 From: jvalvo Date: Wed, 13 Dec 2023 11:50:40 -0800 Subject: [PATCH 2/9] fix: fixed active_workers monitor --- querybook/server/lib/celery/celery_stats.py | 2 +- .../server/lib/stats_logger/loggers/datadog_stats_logger.py | 2 +- querybook/server/tasks/all_tasks.py | 4 ---- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/querybook/server/lib/celery/celery_stats.py b/querybook/server/lib/celery/celery_stats.py index 1059ac7e8..c5cef597b 100644 --- a/querybook/server/lib/celery/celery_stats.py +++ b/querybook/server/lib/celery/celery_stats.py @@ -17,7 +17,7 @@ def send_stats_logger_metrics(celery): if len(active[worker]) > 0: active_tasks += 1 - stats_logger.gauge(ACTIVE_WORKERS, len(stats)) + stats_logger.gauge(ACTIVE_WORKERS, len(active_workers)) stats_logger.gauge(ACTIVE_TASKS, active_tasks) time.sleep(5) diff --git a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py index add3e1bb4..a71172ab3 100644 --- a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py +++ b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py @@ -65,4 +65,4 @@ def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None: ) def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None: - statsd.set(self.metric_prefix_helper(key), value, tags=self.tag_helper(tags)) + statsd.gauge(self.metric_prefix_helper(key), value, tags=self.tag_helper(tags)) diff --git a/querybook/server/tasks/all_tasks.py b/querybook/server/tasks/all_tasks.py index 41cae41e5..f9afc3535 100644 --- a/querybook/server/tasks/all_tasks.py +++ b/querybook/server/tasks/all_tasks.py @@ -15,8 +15,6 @@ TASK_FAILURES, TASK_SUCCESSES, TASK_RECEIVED, - ACTIVE_WORKERS, - ACTIVE_TASKS, stats_logger, ) from lib.celery.celery_stats import start_stats_logger_monitor @@ -95,5 +93,3 @@ def handle_task_received(sender, signal, *args, **kwargs): def start_celery_stats_logging(sender=None, conf=None, **kwargs): LOG.info("Starting Celery Beat") start_stats_logger_monitor(celery) - stats_logger.gauge(ACTIVE_WORKERS, 0) - stats_logger.gauge(ACTIVE_TASKS, 0) From f23037aaeeed2368fd2d0225e3cd6e55edc0031f Mon Sep 17 00:00:00 2001 From: jvalvo Date: Tue, 19 Dec 2023 10:49:54 -0800 Subject: [PATCH 3/9] fix: fixed active_tasks only displaying 1 tasks per worker --- querybook/server/lib/celery/celery_stats.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/querybook/server/lib/celery/celery_stats.py b/querybook/server/lib/celery/celery_stats.py index c5cef597b..43fce3366 100644 --- a/querybook/server/lib/celery/celery_stats.py +++ b/querybook/server/lib/celery/celery_stats.py @@ -7,15 +7,13 @@ def send_stats_logger_metrics(celery): while True: i = celery.control.inspect() - stats = i.stats() or {} active = i.active() or {} - active_workers = list(stats.keys()) + active_workers = list(active.keys()) active_tasks = 0 for worker in active_workers: if worker in active: - if len(active[worker]) > 0: - active_tasks += 1 + active_tasks += len(active[worker]) stats_logger.gauge(ACTIVE_WORKERS, len(active_workers)) stats_logger.gauge(ACTIVE_TASKS, active_tasks) From 9e3697a1dd1fabf98aff57de5bf0656706f56b2a Mon Sep 17 00:00:00 2001 From: Julian Rojas Valvo Date: Wed, 20 Dec 2023 12:14:40 -0800 Subject: [PATCH 4/9] Delete querybook/server/__init__.py --- querybook/server/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 querybook/server/__init__.py diff --git a/querybook/server/__init__.py b/querybook/server/__init__.py deleted file mode 100644 index e69de29bb..000000000 From 0efaad97757528f5674af1bcdf5a3766204d5a07 Mon Sep 17 00:00:00 2001 From: Julian Rojas Valvo Date: Wed, 20 Dec 2023 12:15:13 -0800 Subject: [PATCH 5/9] Update env.py --- querybook/server/env.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/querybook/server/env.py b/querybook/server/env.py index 537e2f983..99b819ddb 100644 --- a/querybook/server/env.py +++ b/querybook/server/env.py @@ -135,6 +135,11 @@ class QuerybookSettings(object): AI_ASSISTANT_PROVIDER = get_env_config("AI_ASSISTANT_PROVIDER") AI_ASSISTANT_CONFIG = get_env_config("AI_ASSISTANT_CONFIG") or {} + VECTOR_STORE_PROVIDER = get_env_config("VECTOR_STORE_PROVIDER") + VECTOR_STORE_CONFIG = get_env_config("VECTOR_STORE_CONFIG") or {} + EMBEDDINGS_PROVIDER = get_env_config("EMBEDDINGS_PROVIDER") + EMBEDDINGS_CONFIG = get_env_config("EMBEDDINGS_CONFIG") or {} + # Datadog DD_AGENT_HOST = get_env_config("DD_AGENT_HOST", optional=True) DD_DOGSTATSD_PORT = int(get_env_config("DD_DOGSTATSD_PORT", optional=True) or 8125) From 8f4f91c10d2a6119a7b3268158f98decb6489891 Mon Sep 17 00:00:00 2001 From: Dave Bauman Date: Thu, 18 Jan 2024 21:46:59 -0500 Subject: [PATCH 6/9] fix: Make Datadog optional --- docker-compose.yml | 46 +++++++++++-------- plugins/stats_logger_plugin/__init__.py | 9 ++++ .../lib/stats_logger/all_stats_loggers.py | 2 - .../loggers/datadog_stats_logger.py | 14 ++++++ 4 files changed, 49 insertions(+), 22 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index af242addf..d4bdc85ec 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -20,26 +20,6 @@ x-querybook-volumes: &querybook-volumes # - file:/opt/store/ services: - datadog: - image: gcr.io/datadoghq/agent:latest - ports: - - 8125:8125/udp - expose: - - 8125/udp - environment: - DD_DOGSTATSD_NON_LOCAL_TRAFFIC: true - DD_USE_DOGSTATSD: true - DD_APM_ENABLED: false - DD_AC_EXCLUDE: '.*' - DD_ENV: 'dev' - DD_SERVICE: 'federated-querybook' - DD_VERSION: 'latest' - env_file: - - ./.env.local - volumes: - - /var/run/docker.sock:/var/run/docker.sock - - /proc/:/host/proc/:ro - - /sys/fs/cgroup:/host/sys/fs/cgroup:ro web: container_name: querybook_web image: *querybook-image @@ -147,6 +127,32 @@ services: timeout: 30s retries: 3 + # + # OPTIONAL SERVICES + # Start separately if you need them via `docker compose up ` + # + + datadog: + image: gcr.io/datadoghq/agent:latest + ports: + - 8125:8125/udp + expose: + - 8125/udp + environment: + DD_DOGSTATSD_NON_LOCAL_TRAFFIC: true + DD_USE_DOGSTATSD: true + DD_APM_ENABLED: false + DD_AC_EXCLUDE: '.*' + DD_ENV: 'dev' + DD_SERVICE: 'querybook' + DD_VERSION: 'latest' + volumes: + - /var/run/docker.sock:/var/run/docker.sock + - /proc/:/host/proc/:ro + - /sys/fs/cgroup:/host/sys/fs/cgroup:ro + profiles: + - datadog + # EMAIL SERVER EXAMPLE # If you need email to work use this # dockerhostforward: diff --git a/plugins/stats_logger_plugin/__init__.py b/plugins/stats_logger_plugin/__init__.py index cd87502c7..7a4c797c9 100644 --- a/plugins/stats_logger_plugin/__init__.py +++ b/plugins/stats_logger_plugin/__init__.py @@ -1 +1,10 @@ ALL_PLUGIN_STATS_LOGGERS = [] + +# Example to add Datadog stats logger +# Requires requirements/datadog/datadog.txt to be installed +# +# from lib.stats_logger.loggers.datadog_stats_logger import DatadogStatsLogger +# +# ALL_PLUGIN_STATS_LOGGERS = [ +# DatadogStatsLogger(), +# ] diff --git a/querybook/server/lib/stats_logger/all_stats_loggers.py b/querybook/server/lib/stats_logger/all_stats_loggers.py index 0fba53a03..86f6473f2 100644 --- a/querybook/server/lib/stats_logger/all_stats_loggers.py +++ b/querybook/server/lib/stats_logger/all_stats_loggers.py @@ -1,6 +1,5 @@ from lib.utils.import_helper import import_module_with_default -from .loggers.datadog_stats_logger import DatadogStatsLogger from .loggers.null_stats_logger import NullStatsLogger from .loggers.console_stats_logger import ConsoleStatsLogger @@ -13,7 +12,6 @@ ALL_STATS_LOGGERS = [ NullStatsLogger(), ConsoleStatsLogger(), - DatadogStatsLogger(), ] + ALL_PLUGIN_STATS_LOGGERS diff --git a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py index a71172ab3..21034a0eb 100644 --- a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py +++ b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py @@ -7,6 +7,20 @@ class DatadogStatsLogger(BaseStatsLogger): + """ + Stats Logger implemention for Datadog using DogStatsD. + + Required environment variables: + - DD_API_KEY: The API key for Datadog. + - DD_AGENT_HOST: The host of the Datadog agent. + - DD_DOGSTATSD_PORT: The port of the Datadog agent, defaults to 8125. + - DD_SERVICE: The service name. + + Optional environment variables: + - DD_PREFIX: The prefix for all metrics. + - DD_TAGS: Additional tags to be added to all metrics. + """ + metric_prefix = "" dd_tags = [] From 8b25c98476ecb753d5ad0d5e6a2995532f983621 Mon Sep 17 00:00:00 2001 From: Dave Bauman Date: Thu, 18 Jan 2024 22:14:00 -0500 Subject: [PATCH 7/9] fix: Switch to lazy importing to simplify usage --- plugins/stats_logger_plugin/__init__.py | 9 -------- .../lib/stats_logger/all_stats_loggers.py | 2 ++ .../loggers/datadog_stats_logger.py | 21 ++++++++++++++----- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/plugins/stats_logger_plugin/__init__.py b/plugins/stats_logger_plugin/__init__.py index 7a4c797c9..cd87502c7 100644 --- a/plugins/stats_logger_plugin/__init__.py +++ b/plugins/stats_logger_plugin/__init__.py @@ -1,10 +1 @@ ALL_PLUGIN_STATS_LOGGERS = [] - -# Example to add Datadog stats logger -# Requires requirements/datadog/datadog.txt to be installed -# -# from lib.stats_logger.loggers.datadog_stats_logger import DatadogStatsLogger -# -# ALL_PLUGIN_STATS_LOGGERS = [ -# DatadogStatsLogger(), -# ] diff --git a/querybook/server/lib/stats_logger/all_stats_loggers.py b/querybook/server/lib/stats_logger/all_stats_loggers.py index 86f6473f2..0fba53a03 100644 --- a/querybook/server/lib/stats_logger/all_stats_loggers.py +++ b/querybook/server/lib/stats_logger/all_stats_loggers.py @@ -1,5 +1,6 @@ from lib.utils.import_helper import import_module_with_default +from .loggers.datadog_stats_logger import DatadogStatsLogger from .loggers.null_stats_logger import NullStatsLogger from .loggers.console_stats_logger import ConsoleStatsLogger @@ -12,6 +13,7 @@ ALL_STATS_LOGGERS = [ NullStatsLogger(), ConsoleStatsLogger(), + DatadogStatsLogger(), ] + ALL_PLUGIN_STATS_LOGGERS diff --git a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py index 21034a0eb..699fa54a5 100644 --- a/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py +++ b/querybook/server/lib/stats_logger/loggers/datadog_stats_logger.py @@ -1,5 +1,4 @@ from lib.stats_logger.base_stats_logger import BaseStatsLogger -from datadog import initialize, statsd from env import QuerybookSettings from lib.logger import get_logger @@ -23,6 +22,7 @@ class DatadogStatsLogger(BaseStatsLogger): metric_prefix = "" dd_tags = [] + _statsd = None def metric_prefix_helper(self, key): return self.metric_prefix + "." + key @@ -33,6 +33,15 @@ def tag_helper(self, tags): return [] def initialize(self): + try: + from datadog import initialize, statsd + + self._statsd = statsd + except ImportError: + raise ImportError( + "Datadog is not installed. Please install `requirements/datadog/datadog.txt` to use the Datadog stats logger." + ) + if QuerybookSettings.DD_AGENT_HOST and QuerybookSettings.DD_DOGSTATSD_PORT: LOG.info("Initializing Datadog") @@ -60,23 +69,25 @@ def logger_name(self) -> str: return "datadog" def incr(self, key: str, tags: dict[str, str] = None) -> None: - statsd.increment( + self._statsd.increment( self.metric_prefix_helper(key), 1, tags=self.tag_helper(tags), ) def decr(self, key: str, tags: dict[str, str] = None) -> None: - statsd.decrement( + self._statsd.decrement( self.metric_prefix_helper(key), 1, tags=self.tag_helper(tags), ) def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None: - statsd.histogram( + self._statsd.histogram( self.metric_prefix_helper(key), value, tags=self.tag_helper(tags) ) def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None: - statsd.gauge(self.metric_prefix_helper(key), value, tags=self.tag_helper(tags)) + self._statsd.gauge( + self.metric_prefix_helper(key), value, tags=self.tag_helper(tags) + ) From ec69f7d7fa79ce48db3bce414b9ac09f278e05bf Mon Sep 17 00:00:00 2001 From: Julian Rojas Valvo Date: Mon, 22 Jan 2024 08:18:14 -0800 Subject: [PATCH 8/9] Update test.txt --- requirements/test.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements/test.txt b/requirements/test.txt index 200423617..01c2f4ce0 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -8,4 +8,3 @@ moto==2.2.4 -r exporter/gspread.txt -r parser/sqlglot.txt -r ai/langchain.txt --r datadog/datadog.txt From 59baad79c290810949d0d7b611f99f63e9ff52f4 Mon Sep 17 00:00:00 2001 From: Julian Rojas Valvo Date: Mon, 22 Jan 2024 09:17:37 -0800 Subject: [PATCH 9/9] Update extra.txt --- requirements/extra.txt | 3 --- 1 file changed, 3 deletions(-) diff --git a/requirements/extra.txt b/requirements/extra.txt index e99892e88..db5edd23e 100644 --- a/requirements/extra.txt +++ b/requirements/extra.txt @@ -28,6 +28,3 @@ # AI Assistant -r ai/langchain.txt - -# Datadog --r datadog/datadog.txt