python(feat): lazy flow creation (#296)

Author: nathan-sift

python/lib/sift_py/error.py

@@ -1,5 +1,7 @@
 import warnings
 
+import google.protobuf.message
+
 
 class SiftError(Exception):
     """
@@ -28,3 +30,20 @@ def _component_deprecation_warning():
         "See docs for more details: https://docs.siftstack.com/docs/glossary#component",
         SiftAPIDeprecationWarning,
     )
+
+
+# The default max message size for the Sift gRPC server.
+GRPC_MAX_MESSAGE_SIZE = 4_194_304
+
+
+class ProtobufMaxSizeExceededError(Exception):
+    """
+    The library limits the size of certain protobufs to prevent gRPC messages from being too big.
+    """
+
+
+def raise_if_too_large(pb: google.protobuf.message.Message):
+    size = len(pb.SerializeToString())
+    name = getattr(pb, "name", pb.__class__.__name__)
+    if size > GRPC_MAX_MESSAGE_SIZE:
+        raise ProtobufMaxSizeExceededError(f"{name} too large: {size}")

python/lib/sift_py/ingestion/_internal/ingest.py

@@ -3,9 +3,11 @@
 import logging
 from collections.abc import Callable
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Set, Union, cast
 
+import grpc
 from google.protobuf.timestamp_pb2 import Timestamp
+from grpc import StatusCode
 from sift.ingest.v1.ingest_pb2 import (
     IngestWithConfigDataChannelValue,
     IngestWithConfigDataStreamRequest,
@@ -14,6 +16,7 @@
 from sift.ingestion_configs.v2.ingestion_configs_pb2 import ChannelConfig as ChannelConfigPb
 from sift.ingestion_configs.v2.ingestion_configs_pb2 import IngestionConfig
 
+from sift_py.error import ProtobufMaxSizeExceededError
 from sift_py.grpc.transport import SiftChannel
 from sift_py.ingestion._internal.error import IngestionValidationError
 from sift_py.ingestion._internal.ingestion_config import (
@@ -43,11 +46,13 @@ class _IngestionServiceImpl:
     ingestion_config: IngestionConfig
     asset_name: str
     flow_configs_by_name: Dict[str, FlowConfig]
+    flow_configs_created: Set[str]
     rules: List[RuleConfig]
     run_id: Optional[str]
     organization_id: Optional[str]
     end_stream_on_error: bool
     config: TelemetryConfig
+    use_lazy_flow_creation: bool
 
     ingest_service_stub: IngestServiceStub
     rule_service: RuleService
@@ -58,21 +63,49 @@ def __init__(
         config: TelemetryConfig,
         run_id: Optional[str] = None,
         end_stream_on_error: bool = False,
+        force_lazy_flow_creation: bool = False,
     ):
-        ingestion_config = self.__class__._get_or_create_ingestion_config(channel, config)
+        try:
+            ingestion_config = self.__class__._get_or_create_ingestion_config(
+                channel, config, lazy_flows=force_lazy_flow_creation
+            )
+        except ProtobufMaxSizeExceededError:
+            ingestion_config = self.__class__._get_or_create_ingestion_config(
+                channel, config, lazy_flows=True
+            )
+            self.use_lazy_flow_creation = True
+        else:
+            self.use_lazy_flow_creation = force_lazy_flow_creation
+
         self.ingestion_config = ingestion_config
 
-        if config._ingestion_client_key_is_generated:
-            # If this is a generated key, use the local telemetry config since it is static.
+        # `flow_configs_by_name` will include all flows in the config, and anything already created
+        # `flow_configs_created` only includes those which are already registered
+        if config._ingestion_client_key_is_generated and not self.use_lazy_flow_creation:
+            # If this is a generated key, use the local telemetry config since it is static
+            # All flows have also already been created since we aren't lazy
             self.flow_configs_by_name = {flow.name: flow for flow in config.flows}
+            self.flow_configs_created = {flow.name for flow in config.flows}
         else:
-            # If the user specified a client key, use the configuration from Sift since it
-            # may have been updated.
-            flows = [
-                FlowConfig.from_pb(f)
-                for f in get_ingestion_config_flows(channel, ingestion_config.ingestion_config_id)
-            ]
-            self.flow_configs_by_name = {flow.name: flow for flow in flows}
+            # If using lazy flow creation, assume the list of flows is large enough that
+            # `get_ingestion_config_flows` will be very slow. Instead lazily check for
+            # flow existance during streaming
+            # Otherwise, since the user specified a client key, use the configuration from Sift since it
+            # may have been updated, and we've already registered any new flows
+
+            if self.use_lazy_flow_creation:
+                self.flow_configs_by_name = {flow.name: flow for flow in config.flows}
+                self.flow_configs_created: Set[str] = set()
+            else:
+                flows = [
+                    FlowConfig.from_pb(f)
+                    for f in get_ingestion_config_flows(
+                        channel,
+                        ingestion_config.ingestion_config_id,
+                    )
+                ]
+                self.flow_configs_by_name = {flow.name: flow for flow in flows}
+                self.flow_configs_created = {flow.name for flow in flows}
 
         self.rule_service = RuleService(channel)
         if config.rules:
@@ -94,6 +127,10 @@ def ingest(self, *requests: IngestWithConfigDataStreamRequest):
         """
         Perform data ingestion.
         """
+        # Perform lazy flow registration if needed
+        if self.use_lazy_flow_creation:
+            self._lazy_flow_creation(*requests)
+
         self.ingest_service_stub.IngestWithConfigDataStream(iter(requests))
 
     def ingest_flows(self, *flows: FlowOrderedChannelValues):
@@ -110,6 +147,10 @@ def ingest_flows(self, *flows: FlowOrderedChannelValues):
             req = self.create_ingestion_request(flow_name, timestamp, channel_values)
             requests.append(req)
 
+        # Perform lazy flow registration if needed
+        if self.use_lazy_flow_creation:
+            self._lazy_flow_creation(*requests)
+
         self.ingest_service_stub.IngestWithConfigDataStream(iter(requests))
 
     def try_ingest_flows(self, *flows: Flow):
@@ -126,6 +167,10 @@ def try_ingest_flows(self, *flows: Flow):
             req = self.try_create_ingestion_request(flow_name, timestamp, channel_values)
             requests.append(req)
 
+        # Perform lazy flow registration if needed
+        if self.use_lazy_flow_creation:
+            self._lazy_flow_creation(*requests)
+
         self.ingest_service_stub.IngestWithConfigDataStream(iter(requests))
 
     def attach_run(
@@ -374,6 +419,7 @@ def try_create_flow(self, *flow_config: FlowConfig):
 
         for fc in flow_config:
             self.flow_configs_by_name[fc.name] = fc
+            self.flow_configs_created.add(fc.name)
 
     def try_create_flows(self, *flow_configs: FlowConfig):
         """
@@ -399,13 +445,60 @@ def create_flow(self, *flow_config: FlowConfig):
         )
         for fc in flow_config:
             self.flow_configs_by_name[fc.name] = fc
+            self.flow_configs_created.add(fc.name)
 
     def create_flows(self, *flow_configs: FlowConfig):
         """
         See `create_flow`.
         """
         return self.create_flow(*flow_configs)
 
+    def _lazy_flow_creation(self, *requests: IngestWithConfigDataStreamRequest):
+        """
+        Used for lazy flow creation, which registers flows with sift as they are seen, instead of during
+        the service initialization
+        """
+        missing_flow_config_names: Set[str] = set()
+        for request in requests:
+            # Skip creation if already registered or an unknown flow name
+            if (
+                request.flow not in self.flow_configs_created
+                and request.flow in self.flow_configs_by_name
+            ):
+                missing_flow_config_names.add(request.flow)
+
+        flow_configs_to_create = [
+            self.flow_configs_by_name[flow_config_name]
+            for flow_config_name in missing_flow_config_names
+        ]
+
+        if flow_configs_to_create:
+            try:
+                create_flow_configs(
+                    self.transport_channel,
+                    self.ingestion_config.ingestion_config_id,
+                    flow_configs_to_create,
+                )
+            except (ProtobufMaxSizeExceededError, grpc.RpcError) as e:
+                # Re-raise gRPC errors unless it just an ALREADY_EXISTS error
+                if isinstance(e, grpc.RpcError):
+                    if e.code() != StatusCode.ALREADY_EXISTS:
+                        raise
+
+                # Try creating them individually instead
+                for flow_config in flow_configs_to_create:
+                    try:
+                        create_flow_configs(
+                            self.transport_channel,
+                            self.ingestion_config.ingestion_config_id,
+                            [flow_config],
+                        )
+                    except grpc.RpcError as e:
+                        if e.code() != StatusCode.ALREADY_EXISTS:
+                            raise
+            for flow_config in flow_configs_to_create:
+                self.flow_configs_created.add(flow_config.name)
+
     @staticmethod
     def _update_flow_configs(
         channel: SiftChannel, ingestion_config_id: str, telemetry_config: TelemetryConfig
@@ -464,27 +557,33 @@ def _update_flow_configs(
 
     @classmethod
     def _get_or_create_ingestion_config(
-        cls, channel: SiftChannel, config: TelemetryConfig
+        cls, channel: SiftChannel, config: TelemetryConfig, lazy_flows: bool = False
     ) -> IngestionConfig:
         """
         Retrieves an existing ingestion config or creates a new one. If an existing ingestion config is fetched,
         then flows may be updated to reflect any changes that may have occured in the telemetry config.
+        May raise `ProtobufMaxSizeExceeded` if a large number of flows needing updates or creation are passed
         """
 
         ingestion_config = get_ingestion_config_by_client_key(channel, config.ingestion_client_key)
 
         # Exiting ingestion config.. update flows if necessary
         if ingestion_config is not None:
-            if config._ingestion_client_key_is_generated:
+            if config._ingestion_client_key_is_generated or lazy_flows:
                 return ingestion_config
             else:
                 cls._update_flow_configs(channel, ingestion_config.ingestion_config_id, config)
                 return ingestion_config
 
+        if lazy_flows:
+            config_flows = []
+        else:
+            config_flows = config.flows
+
         ingestion_config = create_ingestion_config(
             channel,
             config.asset_name,
-            config.flows,
+            config_flows,
             config.ingestion_client_key,
             config.organization_id,
         )

python/lib/sift_py/ingestion/_internal/ingestion_config.py

@@ -19,6 +19,7 @@
     IngestionConfigServiceStub,
 )
 
+from sift_py.error import raise_if_too_large
 from sift_py.grpc.transport import SiftChannel
 from sift_py.ingestion.flow import FlowConfig
 
@@ -54,6 +55,7 @@ def create_ingestion_config(
 ) -> IngestionConfig:
     """
     Creates a new ingestion config
+    Will raise `ProtobufMaxSizeExceeded` if message size is too large to send
     """
 
     svc = IngestionConfigServiceStub(channel)
@@ -63,6 +65,9 @@ def create_ingestion_config(
         organization_id=organization_id or "",
         flows=[flow.as_pb(FlowConfigPb) for flow in flows],
     )
+
+    raise_if_too_large(req)
+
     res = cast(CreateIngestionConfigResponse, svc.CreateIngestionConfig(req))
     return res.ingestion_config
 
@@ -128,10 +133,14 @@ def create_flow_configs(
 ):
     """
     Adds flow configs to an existing ingestion config.
+    Will raise `ProtobufMaxSizeExceeded` if message size is too large to send
     """
     svc = IngestionConfigServiceStub(channel)
     req = CreateIngestionConfigFlowsRequest(
         ingestion_config_id=ingestion_config_id,
         flows=[f.as_pb(FlowConfigPb) for f in flow_configs],
     )
+
+    raise_if_too_large(req)
+
     _ = cast(CreateIngestionConfigFlowsResponse, svc.CreateIngestionConfigFlows(req))