Revert "Explicitly initialize snuba take 2"

.github/workflows/ci.yml

@@ -262,12 +262,6 @@ jobs:
           SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_distributed_migrations TEST_LOCATION=test_distributed_migrations docker-compose --profile multi_node -f docker-compose.gcb.yml run --rm snuba-test
         if: ${{ matrix.snuba_settings == 'test_distributed_migrations' }}
 
-      - name: Docker Snuba Init Tests
-        run: |
-          SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_initialization TEST_LOCATION=test_initialization docker-compose -f docker-compose.gcb.yml run --rm snuba-test
-        if: ${{ matrix.snuba_settings == 'test' }}
-
-
       - name: Upload to codecov
         run: |
           curl -Os https://uploader.codecov.io/latest/linux/codecov && chmod +x codecov && ./codecov -t ${CODECOV_TOKEN}

Makefile

@@ -19,14 +19,6 @@ test-distributed-migrations:
 	SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_distributed_migrations docker-compose --profile multi_node -f docker-compose.gcb.yml up -d
 	SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_distributed_migrations TEST_LOCATION=test_distributed_migrations docker-compose --profile multi_node -f docker-compose.gcb.yml run --rm snuba-test
 
-test-initialization:
-	docker build . -t snuba-test
-	SNUBA_IMAGE=snuba-test docker-compose -f docker-compose.gcb.yml down
-	SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_initialization docker-compose -f docker-compose.gcb.yml up -d
-	SNUBA_IMAGE=snuba-test SNUBA_SETTINGS=test_initialization TEST_LOCATION=test_initialization docker-compose -f docker-compose.gcb.yml run --rm snuba-test
-
-
-
 tests: test
 
 api-tests:

snuba/cli/__init__.py

@@ -1,79 +1,51 @@
-from __future__ import annotations
-
 import logging
-import os
 import time
-from typing import Any
+from pkgutil import walk_packages
 
 import click
 import sentry_sdk
 import structlog
 
-from snuba.core import initialize
 from snuba.environment import metrics as environment_metrics
 from snuba.environment import setup_logging, setup_sentry
 from snuba.utils.metrics.wrapper import MetricsWrapper
 
 setup_sentry()
+setup_logging("DEBUG")
+logger = logging.getLogger("snuba_init")
 
-with sentry_sdk.start_transaction(
-    op="snuba_init", name="Snuba CLI Initialization", sampled=True
-):
-
-    setup_logging("DEBUG")
-    logger = logging.getLogger("snuba_init")
+start = time.perf_counter()
+logger.info("Initializing Snuba...")
+metrics = MetricsWrapper(environment_metrics, "cli")
 
-    start = time.perf_counter()
-    logger.info("Initializing Snuba CLI...")
-    metrics = MetricsWrapper(environment_metrics, "cli")
 
-    plugin_folder = os.path.dirname(__file__)
+@click.group()
+@click.version_option()
+def main() -> None:
+    """\b
+                         o
+                        O
+                        O
+                        o
+    .oOo  'OoOo. O   o  OoOo. .oOoO'
+    `Ooo.  o   O o   O  O   o O   o
+        O  O   o O   o  o   O o   O
+    `OoO'  o   O `OoO'o `OoO' `OoO'o"""
 
-    class SnubaCLI(click.MultiCommand):
-        def list_commands(self, ctx: Any) -> list[str]:
-            rv = []
-            for filename in os.listdir(plugin_folder):
-                if filename.endswith(".py") and filename != "__init__.py":
-                    # IMPORTANT!!!!
-                    # Click replaces underscores with dashes for command names
-                    # by default. To mimic this behavior we have to do that here
-                    # since all of our infrastructure depends on this being the
-                    # case
-                    rv.append(filename[:-3].replace("_", "-"))
-            rv.sort()
-            return rv
 
-        def get_command(self, ctx: Any, name: str) -> click.Command:
-            logger.info(f"Loading command {name}")
-            # IMPORTANT!!!!
-            # Click replaces underscores with dashes for command names
-            # by default. To mimic this behavior we have to do that here
-            # since all of our infrastructure depends on this being the
-            # case
-            actual_command_name = name.replace("-", "_")
-            with sentry_sdk.start_span(op="import", description=name):
-                ns: dict[str, click.Command] = {}
-                fn = os.path.join(plugin_folder, actual_command_name + ".py")
-                with open(fn) as f:
-                    code = compile(f.read(), fn, "exec")
-                    eval(code, ns, ns)
-                initialize.initialize()
-                init_time = time.perf_counter() - start
-                metrics.gauge("snuba_init", init_time)
-                logger.info(f"Snuba initialization took {init_time}s")
-                return ns[actual_command_name]
-
-    @click.command(cls=SnubaCLI)
-    @click.version_option()
-    def main() -> None:
-        """\b
-                             o
-                            O
-                            O
-                            o
-        .oOo  'OoOo. O   o  OoOo. .oOoO'
-        `Ooo.  o   O o   O  O   o O   o
-            O  O   o O   o  o   O o   O
-        `OoO'  o   O `OoO'o `OoO' `OoO'o"""
-
-    structlog.reset_defaults()
+with sentry_sdk.start_transaction(
+    op="snuba_init", name="Snuba CLI Initialization", sampled=True
+):
+    for loader, module_name, is_pkg in walk_packages(__path__, __name__ + "."):
+        logger.info(f"Loading module {module_name}")
+        with sentry_sdk.start_span(op="import", description=module_name):
+            module = __import__(module_name, globals(), locals(), ["__name__"])
+            cmd = getattr(module, module_name.rsplit(".", 1)[-1])
+            if isinstance(cmd, click.Command):
+                main.add_command(cmd)
+
+init_time = time.perf_counter() - start
+metrics.gauge("snuba_init", init_time)
+logger.info(f"Snuba initialization took {init_time}s")
+
+structlog.reset_defaults()

snuba/datasets/factory.py

@@ -110,10 +110,10 @@ class InvalidDatasetError(SerializableException):
 _DS_FACTORY: _DatasetFactory | None = None
 
 
-def _ds_factory(reset: bool = False) -> _DatasetFactory:
+def _ds_factory() -> _DatasetFactory:
     # This function can be acessed by many threads at once. It is okay if more than one thread recreates the same object.
     global _DS_FACTORY
-    if _DS_FACTORY is None or reset:
+    if _DS_FACTORY is None:
         _DS_FACTORY = _DatasetFactory()
     return _DS_FACTORY
 
@@ -137,4 +137,5 @@ def get_config_built_datasets() -> dict[str, Dataset]:
 
 
 def reset_dataset_factory() -> None:
-    _ds_factory(reset=True)
+    global _DS_FACTORY
+    _DS_FACTORY = _DatasetFactory()