ingestion.py

"""Low-level wrapper for the IngestionAPI.

This module provides thin wrappers around the autogenerated bindings for the IngestionAPI.
It handles common concerns like error handling and retries.

It provides an asynchronous client for the IngestionAPI.
"""

from __future__ import annotations

import asyncio
import atexit
import hashlib
import logging
import threading
import time
from collections import namedtuple
from queue import Queue
from typing import TYPE_CHECKING, Any, cast

from sift.ingestion_configs.v2.ingestion_configs_pb2 import (
    GetIngestionConfigRequest,
    ListIngestionConfigFlowsResponse,
    ListIngestionConfigsRequest,
    ListIngestionConfigsResponse,
)
from sift.ingestion_configs.v2.ingestion_configs_pb2_grpc import (
    IngestionConfigServiceStub,
)

from sift_client._internal.low_level_wrappers.base import (
    LowLevelClientBase,
)
from sift_client._internal.util.timestamp import to_rust_py_timestamp
from sift_client.sift_types.ingestion import Flow, IngestionConfig, _to_rust_value
from sift_client.transport import GrpcClient, WithGrpcClient
from sift_client.util import cel_utils as cel

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from datetime import datetime

    from sift_stream_bindings import (
        IngestionConfigFormPy,
        IngestWithConfigDataStreamRequestPy,
        SiftStreamBuilderPy,
    )


class IngestionThread(threading.Thread):
    """Manages ingestion for a single ingestion config."""

    IDLE_LOOP_PERIOD = 0.1  # Time of intervals loop will sleep while waiting for data.
    SIFT_STREAM_FINISH_TIMEOUT = 0.06  # Measured ~0.05s to finish stream.
    CLEANUP_TIMEOUT = IDLE_LOOP_PERIOD + SIFT_STREAM_FINISH_TIMEOUT

    def __init__(
        self,
        sift_stream_builder: SiftStreamBuilderPy,
        data_queue: Queue,
        ingestion_config: IngestionConfigFormPy,
        no_data_timeout: int = 1,
        metric_interval: float = 0.5,
    ):
        """Initialize the IngestionThread.

        Args:
            sift_stream_builder: The sift stream builder to build a new stream.
            data_queue: The queue to put IngestWithConfigDataStreamRequestPy requests into for ingestion.
            ingestion_config: The ingestion config to use for ingestion.
            no_data_timeout: The number of (whole number) seconds to wait for data before stopping the thread (Saves minorly on startup resources. Ingesting new data will always restart the thread if it is stopped).
            metric_interval: Time (seconds) to wait between logging metrics.
        """
        super().__init__(daemon=True)
        self.data_queue = data_queue
        self._stop_event = threading.Event()
        self.sift_stream_builder = sift_stream_builder
        self.ingestion_config = ingestion_config
        self.no_data_timeout = no_data_timeout
        self.metric_interval = metric_interval
        self.initialized = False

    def stop(self):
        self._stop_event.set()
        # Give a brief chance to finish the stream (should take < 50ms).
        time.sleep(self.CLEANUP_TIMEOUT)
        self.task.cancel()

    async def await_stream_build(self):
        while not self.initialized:
            await asyncio.sleep(0.01)

    async def main(self):
        logger.debug("Ingestion thread started")
        self.sift_stream_builder.ingestion_config = self.ingestion_config
        sift_stream = await self.sift_stream_builder.build()
        time_of_last_metric = time.time()
        time_of_last_data = time.time()
        count = 0
        self.initialized = True
        try:
            while True:
                while not self.data_queue.empty():
                    if self._stop_event.is_set():
                        # Being forced to stop. Try to finish the stream.
                        logger.info(
                            f"Ingestion thread received stop signal. Exiting. Sent {count} requests. {self.data_queue.qsize()} requests remaining."
                        )
                        await sift_stream.finish()
                        return
                    time_of_last_metric = time.time()
                    item = self.data_queue.get()
                    sift_stream = await sift_stream.send_requests(item)
                    count += 1
                    time_since_last_metric = time.time() - time_of_last_metric
                    if time_since_last_metric > self.metric_interval:
                        logger.debug(
                            f"Ingestion thread sent {count} requests, remaining: {self.data_queue.qsize()}"
                        )
                        time_of_last_metric = time.time()

                # Queue empty, check if we should stop.
                time_since_last_data = time.time() - time_of_last_data
                if self._stop_event.is_set() or time_since_last_data > self.no_data_timeout:
                    logger.debug(
                        f"No more requests. Stopping. Sent {count} requests. {self.data_queue.qsize()} requests remaining."
                    )
                    await sift_stream.finish()
                    return
                else:
                    await asyncio.sleep(self.IDLE_LOOP_PERIOD)

        except asyncio.CancelledError:
            # It's possible the thread was joined while sleeping waiting for data. Only note error if we have data left.
            if self.data_queue.qsize() > 0:
                logger.error(
                    f"Ingestion thread cancelled without finishing stream. {self.data_queue.qsize()} requests were not sent."
                )

    async def _run(self):
        self.task = asyncio.create_task(self.main())
        await self.task

    def run(self):
        """This thread will handle sending data to Sift."""
        # Even thought this is a thread, we need to run this async task to await send_requests otherwise we get sift_stream consumed errors.
        asyncio.run(self._run())


class IngestionLowLevelClient(LowLevelClientBase, WithGrpcClient):
    """Low-level client for the IngestionAPI.

    This class provides a thin wrapper around the autogenerated bindings for the IngestionAPI.
    It handles common concerns like error handling and retries.
    """

    CacheEntry = namedtuple("CacheEntry", ["data_queue", "ingestion_config", "thread"])

    sift_stream_builder: SiftStreamBuilderPy
    stream_cache: dict[str, CacheEntry]

    def __init__(self, grpc_client: GrpcClient):
        """Initialize the IngestionLowLevelClient.

        Args:
            grpc_client: The gRPC client to use for making API calls.
        """
        from sift_stream_bindings import (
            RecoveryStrategyPy,
            RetryPolicyPy,
            SiftStreamBuilderPy,
        )

        super().__init__(grpc_client=grpc_client)
        # Rust GRPC client expects URI to have http(s):// prefix.
        uri = grpc_client._config.uri
        if not uri.startswith("http"):
            uri = f"https://{uri}" if grpc_client._config.use_ssl else f"http://{uri}"
        self.sift_stream_builder = SiftStreamBuilderPy(
            uri=uri,
            apikey=grpc_client._config.api_key,
        )
        self.sift_stream_builder.enable_tls = grpc_client._config.use_ssl
        # FD-177: Expose configuration for recovery strategy.
        self.sift_stream_builder.recovery_strategy = RecoveryStrategyPy.retry_only(
            RetryPolicyPy.default()
        )
        self.stream_cache = {}

        atexit.register(self.cleanup, timeout=0.1)

    def cleanup(self, timeout: float | None = None):
        """Cleanup the ingestion threads.

        Args:
            timeout: The timeout in seconds to wait for ingestion to complete. If None, will wait forever.
        """
        for cache_entry in self.stream_cache.values():
            data_queue, ingestion_config, thread = cache_entry
            # "None" value on the queue signals its loop to terminate.
            if thread:
                thread.join(timeout=timeout)
                if thread.is_alive():
                    logger.error(
                        f"Ingestion thread did not finish after {timeout} seconds. Forcing stop."
                    )
                    thread.stop()

    async def get_ingestion_config_flows(self, ingestion_config_id: str) -> list[Flow]:
        """Get the flows for an ingestion config."""
        res = await self._grpc_client.get_stub(IngestionConfigServiceStub).GetIngestionConfig(
            GetIngestionConfigRequest(ingestion_config_id=ingestion_config_id)
        )
        res = cast("ListIngestionConfigFlowsResponse", res)
        return [Flow._from_proto(flow) for flow in res.flows]

    async def list_ingestion_configs(self, filter_query: str) -> list[IngestionConfig]:
        """List ingestion configs."""
        res = await self._grpc_client.get_stub(IngestionConfigServiceStub).ListIngestionConfigs(
            ListIngestionConfigsRequest(filter=filter_query)
        )
        res = cast("ListIngestionConfigsResponse", res)
        return [IngestionConfig._from_proto(config) for config in res.ingestion_configs]

    async def get_ingestion_config_id_from_client_key(self, client_key: str) -> str | None:
        """Get the ingestion config id."""
        filter_query = cel.equals("client_key", client_key)
        ingestion_configs = await self.list_ingestion_configs(filter_query)
        if not ingestion_configs:
            return None
        if len(ingestion_configs) > 1:
            raise ValueError(
                f"Expected 1 ingestion config for client key {client_key}, got {len(ingestion_configs)}"
            )
        return ingestion_configs[0].id_

    def _new_ingestion_thread(
        self,
        ingestion_config_id: str,
        ingestion_config: IngestionConfigFormPy | None = None,
    ):
        """Start a new ingestion thread.
        This allows ingestion to happen in the background regardless of if the user is using the sync or async client
        and without them having to set up threading themselves. We are using a thread vs asyncio since our
        sync wrapper will block on incomlete tasks.

        Args:
            ingestion_config_id: The id of the ingestion config for the flows this stream will ingest. Used to cache the stream.
            ingestion_config: The ingestion config to use for ingestion.
        """
        data_queue: Queue[list[IngestWithConfigDataStreamRequestPy]] = Queue()
        existing = self.stream_cache.get(ingestion_config_id)
        if existing:
            existing_data_queue, existing_ingestion_config, existing_thread = existing
            if existing_thread.is_alive():
                return existing_thread
            else:
                ingestion_config = existing_ingestion_config
                # Re-use existing queue since ingest_flow has already put data on it.
                data_queue = existing_data_queue
        assert ingestion_config is not None  # Appease mypy.
        thread = IngestionThread(self.sift_stream_builder, data_queue, ingestion_config)
        thread.start()

        return self.CacheEntry(data_queue, ingestion_config, thread)

    def _hash_flows(self, asset_name: str, flows: list[Flow]) -> str:
        """Generate a client key that should be unique but deterministic for the given asset and flow configuration."""
        # TODO:  Taken from sift_py/ingestion/config/telemetry.py. Confirm intent from Marc.
        m = hashlib.sha256()
        m.update(asset_name.encode())
        for flow in sorted(flows, key=lambda f: f.name):
            m.update(flow.name.encode())
            # Do not sort channels in alphabetical order since order matters.
            for channel in flow.channels:
                m.update(channel.name.encode())
                # Use api_format for data type since that should be consistent between languages.
                m.update(channel.data_type.hash_str(api_format=True).encode())
                m.update((channel.description or "").encode())
                m.update((channel.unit or "").encode())
                if channel.bit_field_elements:
                    for bfe in sorted(channel.bit_field_elements, key=lambda bfe: bfe.index):
                        m.update(bfe.name.encode())
                        m.update(str(bfe.index).encode())
                        m.update(str(bfe.bit_count).encode())
                if channel.enum_types:
                    for enum_name, enum_key in sorted(
                        channel.enum_types.items(), key=lambda it: it[1]
                    ):
                        m.update(str(enum_key).encode())
                        m.update(enum_name.encode())

        return m.hexdigest()

    async def create_ingestion_config(
        self,
        *,
        asset_name: str,
        flows: list[Flow],
        client_key: str | None = None,
    ) -> str:
        """Create an ingestion config.

        Args:
            asset_name: The name of the asset to ingest to.
            flows: The flows to ingest.
            client_key: The client key to use for ingestion. If not provided, a new one will be generated.
            organization_id: The organization id to use for ingestion. Only needed if the user is part of several organizations.

        Returns:
            The id of the new or found ingestion config.
        """
        from sift_stream_bindings import IngestionConfigFormPy

        ingestion_config_id = None
        if client_key:
            logger.debug(f"Getting ingestion config id for client key {client_key}")
            ingestion_config_id = await self.get_ingestion_config_id_from_client_key(client_key)
            if ingestion_config_id:
                # Perform validation that the flows are valid for the ingestion config.
                existing_flows = await self.get_ingestion_config_flows(ingestion_config_id)
                for flow in flows:
                    if flow.name in {existing_flow.name for existing_flow in existing_flows}:
                        raise ValueError(
                            f"Flow {flow.name} already exists for ingestion client {client_key}"
                        )
        else:
            client_key = self._hash_flows(asset_name, flows)
            try:
                logger.debug(f"Getting ingestion config id from generated client key {client_key}")
                ingestion_config_id = await self.get_ingestion_config_id_from_client_key(client_key)
            except ValueError:
                logger.debug(
                    f"No ingestion config found for client key {client_key}. Creating new one."
                )
                pass

        data_queue, ingestion_config, thread = (
            self.stream_cache.get(ingestion_config_id, (None, None, None))
            if ingestion_config_id
            else (None, None, None)
        )
        if not (thread and thread.is_alive()):
            ingestion_config = IngestionConfigFormPy(
                asset_name=asset_name,
                flows=[flow._to_rust_config() for flow in flows],
                client_key=client_key,
            )

            cache_entry = self._new_ingestion_thread(ingestion_config_id or "", ingestion_config)
            if not ingestion_config_id:
                # No ingestion config ID exists for client key but stream builder in ingestion thread should create it.
                await cache_entry.thread.await_stream_build()
                ingestion_config_id = await self.get_ingestion_config_id_from_client_key(client_key)
                assert ingestion_config_id is not None, (
                    "No ingestion config id found after building new stream. Likely server error."
                )
                logger.debug(f"Built new stream for ingestion config {ingestion_config_id}")
            self.stream_cache[ingestion_config_id] = cache_entry

        for flow in flows:
            flow.ingestion_config_id = ingestion_config_id

        if not ingestion_config_id:
            raise ValueError("No ingestion config id found")
        return ingestion_config_id

    def wait_for_ingestion_to_complete(self, timeout: float | None = None):
        """Blocks until all ingestion to complete.

        Args:
            timeout: The timeout in seconds to wait for ingestion to complete. If None, will wait forever.
        """
        logger.debug("Waiting for ingestion to complete")
        self.cleanup(timeout)

    def ingest_flow(
        self,
        *,
        flow: Flow,
        timestamp: datetime,
        channel_values: dict[str, Any],
        organization_id: str | None = None,
    ):
        """Ingest a flow. This is a synchronous call that queues an ingestion request that will be processed asynchronously on a background thread.

        Args:
            flow: The flow to ingest.
            timestamp: The timestamp of the flow.
            channel_values: The channel values to ingest.
            organization_id: The organization id to use for ingestion. Only relevant if the user is part of several organizations.
        """
        from sift_stream_bindings import IngestWithConfigDataStreamRequestPy

        if not flow.ingestion_config_id:
            raise ValueError(
                "Flow has no ingestion config id -- have you created an ingestion config for this flow?"
            )
        cache_entry = self.stream_cache.get(flow.ingestion_config_id)
        if not cache_entry:
            raise ValueError(
                f"Ingestion config {flow.ingestion_config_id} not found. Have you created an ingestion config for this flow?"
            )
        rust_channel_values = []
        # Iterate through all expected channels for flow and convert to ingestion types (missing channels use a special empty type)
        for channel in flow.channels:
            val = channel_values.get(channel.name)
            rust_channel_values.append(_to_rust_value(channel, val))
        req = IngestWithConfigDataStreamRequestPy(
            ingestion_config_id=flow.ingestion_config_id,
            run_id=flow.run_id or "",
            flow=flow.name,
            timestamp=to_rust_py_timestamp(timestamp),
            channel_values=rust_channel_values,
            end_stream_on_validation_error=False,
            organization_id=organization_id or "",  # This will be filled in by the server
        )
        data_queue, ingestion_config, thread = cache_entry
        assert data_queue is not None
        # Put data on queue before potentially starting a new thread so it doesn't initially sleep waiting for data.
        data_queue.put([req])
        if not (thread and thread.is_alive()):
            # We previously had a thread for this ingestion config but it finished ingestion so create a new one.
            self.stream_cache[flow.ingestion_config_id] = self._new_ingestion_thread(
                flow.ingestion_config_id, ingestion_config
            )