Explicitly initialize snuba take 2

.github/workflows/ci.yml

@@ -262,6 +262,12 @@ 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,6 +19,14 @@ 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,51 +1,79 @@
+from __future__ import annotations
+
 import logging
+import os
 import time
-from pkgutil import walk_packages
+from typing import Any
 
 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")
 
-start = time.perf_counter()
-logger.info("Initializing Snuba...")
-metrics = MetricsWrapper(environment_metrics, "cli")
+with sentry_sdk.start_transaction(
+    op="snuba_init", name="Snuba CLI Initialization", sampled=True
+):
 
+    setup_logging("DEBUG")
+    logger = logging.getLogger("snuba_init")
 
-@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"""
+    start = time.perf_counter()
+    logger.info("Initializing Snuba CLI...")
+    metrics = MetricsWrapper(environment_metrics, "cli")
 
+    plugin_folder = os.path.dirname(__file__)
 
-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()
+    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()

snuba/core/initialize.py

@@ -0,0 +1,43 @@
+"""
+This file contains the initialization sequence for snuba. Anything that needs to be created before a snuba command
+is able to run is defined in this file under the `initialize` function
+
+
+All imports are hidden behind their respective functions to avoid running any code at initialization time unless explicitly
+directed
+"""
+import logging
+
+logger = logging.getLogger("snuba_core_init")
+
+
+def _load_datasets() -> None:
+    from snuba.datasets.factory import reset_dataset_factory
+
+    reset_dataset_factory()
+
+
+def _load_storages() -> None:
+    from snuba.datasets.storages.factory import initialize_storage_factory
+
+    initialize_storage_factory()
+
+
+def _load_entities() -> None:
+    from snuba.datasets.entities.factory import initialize_entity_factory
+
+    initialize_entity_factory()
+
+
+def initialize() -> None:
+    logger.info("Initializing snuba")
+
+    # The order of the functions matters The reference direction is
+    #
+    #       datasets -> entities -> storages
+    #
+    # The initialization goes bottom up, starting from the
+    # lowest-level concept (storage) to the highest (dataset).
+    _load_storages()
+    _load_entities()
+    _load_datasets()

snuba/datasets/factory.py

@@ -110,10 +110,10 @@ class InvalidDatasetError(SerializableException):
 _DS_FACTORY: _DatasetFactory | None = None
 
 
-def _ds_factory() -> _DatasetFactory:
+def _ds_factory(reset: bool = False) -> _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:
+    if _DS_FACTORY is None or reset:
         _DS_FACTORY = _DatasetFactory()
     return _DS_FACTORY
 
@@ -137,5 +137,4 @@ def get_config_built_datasets() -> dict[str, Dataset]:
 
 
 def reset_dataset_factory() -> None:
-    global _DS_FACTORY
-    _DS_FACTORY = _DatasetFactory()
+    _ds_factory(reset=True)

test_initialization/test_initialize.py

@@ -0,0 +1,34 @@
+from unittest import mock
+
+from snuba.core.initialize import initialize
+
+
+class TestInitialization:
+    def test_init(
+        self,
+    ):
+        # first make sure all the factories are not initialized
+        # this is accessing private module variables but we don't have a
+        # better way of knowing things are initialized (2022-10-27)
+        from snuba.datasets.entities.factory import _ENT_FACTORY
+        from snuba.datasets.factory import _DS_FACTORY
+        from snuba.datasets.storages.factory import _STORAGE_FACTORY
+
+        for factory in (_DS_FACTORY, _ENT_FACTORY, _STORAGE_FACTORY):
+            assert factory is None
+
+        initialize()
+        from snuba.datasets.entities.factory import get_all_entity_names
+        from snuba.datasets.factory import get_enabled_dataset_names
+        from snuba.datasets.storages.factory import get_all_storage_keys
+
+        # now that we have called the initialize function. We should not have to
+        # load the entities anymore
+        with mock.patch("snuba.datasets.configuration.loader.safe_load") as load_func:
+            for factory_func in (
+                get_enabled_dataset_names,
+                get_all_entity_names,
+                get_all_storage_keys,
+            ):
+                factory_func()
+            load_func.assert_not_called()