sift_py.ingestion.manager

 1from __future__ import annotations
 2
 3from contextlib import contextmanager
 4from typing import Callable, Dict, Iterator, Optional, TypedDict
 5
 6from typing_extensions import Self, TypeAlias
 7
 8from sift_py.grpc.transport import SiftChannel
 9from sift_py.ingestion.config.telemetry import TelemetryConfig
10from sift_py.ingestion.service import IngestionService
11
12IngestionServiceBuilder: TypeAlias = Callable[[SiftChannel], IngestionService]
13
14
15class IngestionServicesManager:
16    """
17    Allows for the initialization of multiple instances of `sift_py.ingestion.service.IngestionService` from
18    either telemetry configs or builders under a single wrapper class that assists in managing data-ingestion
19    for multiple telemetry configs.
20
21    The initializer of this class can be used directly, but prefer to use either `from_builders` or `from_telemetry_configs`.
22    Prefer to use `from_builders` if you have custom options that you want to provide to `sift_py.ingestion.service.IngestionService.__init__`.
23
24    Example usage:
25
26    ```python
27    manager = IngestionServicesManager.from_telementry_configs(grpc_channel, {
28        "config_a": config_a,
29        "config_b": config_b,
30    })
31
32    with manager.ingestion_service("config_a") as config_a:
33        config_a.try_ingest_flow(...)
34
35    with manager.ingestion_service("config_b") as config_b:
36        config_b.try_ingest_flow(...)
37    ```
38    """
39
40    _transport_channel: SiftChannel
41    _ingestion_services: Dict[str, IngestionService]
42
43    def __init__(
44        self, transport_channel: SiftChannel, ingestion_services: Dict[str, IngestionService]
45    ):
46        self._transport_channel = transport_channel
47        self._ingestion_services = ingestion_services
48
49    @classmethod
50    def from_builders(
51        cls, channel: SiftChannel, builders: Dict[str, IngestionServiceBuilder]
52    ) -> Self:
53        return cls(
54            transport_channel=channel,
55            ingestion_services={key: builder(channel) for key, builder in builders.items()},
56        )
57
58    @classmethod
59    def from_telemetry_configs(
60        cls, channel: SiftChannel, telemetry_configs: Dict[str, TelemetryConfig]
61    ) -> Self:
62        return cls(
63            transport_channel=channel,
64            ingestion_services={
65                key: IngestionService(channel, config) for key, config in telemetry_configs.items()
66            },
67        )
68
69    def get_ingestion_service_by_identifier(self, identifier: str) -> Optional[IngestionService]:
70        return self._ingestion_services.get(identifier)
71
72    def __getitem__(self, identifier: str) -> Optional[IngestionService]:
73        return self.get_ingestion_service_by_identifier(identifier)
74
75    @contextmanager
76    def ingestion_service(self, identifier: str) -> Iterator[IngestionService]:
77        ingestion_service = self[identifier]
78
79        if ingestion_service is None:
80            raise IngestionServiceManagerError(
81                f"An ingestion service is not configured for the identifier '{identifier}'."
82            )
83
84        yield ingestion_service
85
86
87class IngestionServiceManagerError(Exception):
88    def __init__(self, msg: str):
89        return super().__init__(msg)
90
91
92class TelemetryConfigByIdentifierMap(TypedDict):
93    identifier: str
94    telemetry_config: TelemetryConfig
95
96
97class IngestionConfigServiceBuilderIdentifierMap(TypedDict):
98    identifier: str
99    builder: IngestionServiceBuilder