Skip to content

Commit

Permalink
Revert "Revert "Explicitly initialize snuba (#3311)" (#3325)"
Browse files Browse the repository at this point in the history
This reverts commit ce0184a.
  • Loading branch information
Vlad Kluev committed Nov 1, 2022
1 parent ce0184a commit de7cb3c
Show file tree
Hide file tree
Showing 9 changed files with 145 additions and 38 deletions.
6 changes: 6 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
8 changes: 8 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
85 changes: 51 additions & 34 deletions snuba/cli/__init__.py
Original file line number Diff line number Diff line change
@@ -1,51 +1,68 @@
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":
rv.append(filename[:-3])
rv.sort()
return rv

def get_command(self, ctx: Any, name: str) -> click.Command:
logger.info(f"Loading command {name}")
with sentry_sdk.start_span(op="import", description=name):
ns: dict[str, click.Command] = {}
fn = os.path.join(plugin_folder, 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[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()
Empty file added snuba/core/__init__.py
Empty file.
43 changes: 43 additions & 0 deletions snuba/core/initialize.py
Original file line number Diff line number Diff line change
@@ -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()
7 changes: 3 additions & 4 deletions snuba/datasets/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
Empty file.
Empty file added test_initialization/__init__.py
Empty file.
34 changes: 34 additions & 0 deletions test_initialization/test_initialize.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit de7cb3c

Please sign in to comment.