From 6ee920c78de51a0a126a3024170d4dcce38c62a2 Mon Sep 17 00:00:00 2001 From: Meredith Heller Date: Mon, 4 Nov 2024 17:22:30 -0800 Subject: [PATCH] dont hard code formatter --- snuba/cli/lw_deletions_consumer.py | 5 +++-- snuba/lw_deletions/formatters.py | 8 +++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/snuba/cli/lw_deletions_consumer.py b/snuba/cli/lw_deletions_consumer.py index 04fce1beee9..a509d514a59 100644 --- a/snuba/cli/lw_deletions_consumer.py +++ b/snuba/cli/lw_deletions_consumer.py @@ -18,7 +18,7 @@ from snuba.datasets.storages.factory import get_writable_storage from snuba.datasets.storages.storage_key import StorageKey from snuba.environment import setup_logging, setup_sentry -from snuba.lw_deletions.formatters import SearchIssuesFormatter +from snuba.lw_deletions.formatters import STORAGE_FORMATTER from snuba.lw_deletions.strategy import ConsumerStrategyFactory from snuba.utils.metrics.wrapper import MetricsWrapper from snuba.utils.streams.metrics_adapter import StreamMetricsAdapter @@ -143,11 +143,12 @@ def handler(signum: int, frame: Any) -> None: ) storage = get_writable_storage(StorageKey(storage_name)) + formatter = STORAGE_FORMATTER[storage_name]() strategy_factory = ConsumerStrategyFactory( max_batch_size=max_batch_size, max_batch_time_ms=max_batch_time_ms, storage=storage, - formatter=SearchIssuesFormatter(), + formatter=formatter, ) consumer = consumer_builder.build_lw_deletions_consumer(strategy_factory) diff --git a/snuba/lw_deletions/formatters.py b/snuba/lw_deletions/formatters.py index 6dcfe5f0150..3132779646a 100644 --- a/snuba/lw_deletions/formatters.py +++ b/snuba/lw_deletions/formatters.py @@ -1,8 +1,9 @@ from abc import ABC, abstractmethod -from typing import MutableMapping, Sequence +from typing import Mapping, MutableMapping, Sequence, Type from attr import dataclass +from snuba.datasets.storages.storage_key import StorageKey from snuba.web.bulk_delete_query import ConditionsType, DeleteQueryMessage @@ -51,3 +52,8 @@ def format( } ) return and_conditions + + +STORAGE_FORMATTER: Mapping[str, Type[Formatter]] = { + StorageKey.SEARCH_ISSUES.value: SearchIssuesFormatter +}