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 hashlib
import logging
from collections import namedtuple
from typing import TYPE_CHECKING, Any, Iterable, cast

from sift.ingestion_configs.v2.ingestion_configs_pb2 import (
    ListIngestionConfigFlowsRequest,
    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.sift_types.ingestion import FlowConfig, IngestionConfig
from sift_client.transport import GrpcClient, WithGrpcClient
from sift_client.util import cel_utils as cel

logger = logging.getLogger(__name__)

DEFAULT_INGESTION_CONFIG_PAGE_SIZE = 100
"""Default page size for ingestion config and flow list calls (flow configs can be large)."""

if TYPE_CHECKING:
    from datetime import datetime

    from sift_stream_bindings import (
        DiskBackupPolicyPy,
        FlowConfigPy,
        FlowDescriptorPy,
        FlowPy,
        IngestionConfigFormPy,
        IngestWithConfigDataStreamRequestPy,
        IngestWithConfigDataStreamRequestWrapperPy,
        MetadataPy,
        RetryPolicyPy,
        RunFormPy,
        RunSelectorPy,
        SiftStreamBuilderPy,
        SiftStreamMetricsSnapshotPy,
        SiftStreamPy,
        TimeValuePy,
    )

    from sift_client.resources.ingestion import StreamingMode, TracingConfig


def _to_rust_py_timestamp(time: datetime) -> TimeValuePy:
    """Convert a Python datetime to a Rust TimeValuePy.

    Args:
        time: The datetime to convert

    Returns:
        A TimeValuePy representation
    """
    # Importing here to allow sift_stream_bindings to be an optional dependancy for non-ingestion users
    from sift_stream_bindings import TimeValuePy

    ts = time.timestamp()
    secs = int(ts)
    nsecs = int((ts - secs) * 1_000_000_000)
    return TimeValuePy.from_timestamp(secs, nsecs)


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 | None
    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.
        """
        super().__init__(grpc_client=grpc_client)

    async def list_ingestion_configs(
        self,
        filter_query: str | None = None,
        page_size: int | None = DEFAULT_INGESTION_CONFIG_PAGE_SIZE,
        page_token: str | None = None,
        order_by: str | None = None,
    ) -> tuple[list[IngestionConfig], str]:
        """List ingestion configs (single page).

        Args:
            filter_query: The CEL filter query.
            page_size: Number of results per page.
            page_token: Token for the next page.
            order_by: Unused; accepted for _handle_pagination compatibility.

        Returns:
            A tuple of (list of IngestionConfig, next_page_token).
        """
        request_kwargs: dict[str, Any] = {}

        if page_size is not None:
            request_kwargs["page_size"] = page_size
        if page_token is not None:
            request_kwargs["page_token"] = page_token
        if filter_query is not None:
            request_kwargs["filter"] = filter_query
        if order_by is not None:
            request_kwargs["order_by"] = order_by

        request = ListIngestionConfigsRequest(**request_kwargs)
        res = await self._grpc_client.get_stub(IngestionConfigServiceStub).ListIngestionConfigs(
            request
        )
        res = cast("ListIngestionConfigsResponse", res)
        configs = [IngestionConfig._from_proto(config) for config in res.ingestion_configs]
        return configs, res.next_page_token

    async def list_all_ingestion_configs(
        self,
        filter_query: str | None = None,
        page_size: int | None = DEFAULT_INGESTION_CONFIG_PAGE_SIZE,
        max_results: int | None = None,
    ) -> list[IngestionConfig]:
        """List all ingestion configs matching the filter, using pagination.

        Args:
            filter_query: The CEL filter query.
            page_size: Number of results per page.
            max_results: Maximum total results to return; None for no limit.

        Returns:
            A list of all matching IngestionConfigs.
        """
        return await self._handle_pagination(
            self.list_ingestion_configs,
            kwargs={"filter_query": filter_query},
            page_size=page_size,
            max_results=max_results,
        )

    async def list_ingestion_config_flows(
        self,
        ingestion_config_id: str,
        page_size: int | None = DEFAULT_INGESTION_CONFIG_PAGE_SIZE,
        page_token: str | None = None,
        order_by: str | None = None,
        filter_query: str | None = None,
    ) -> tuple[list[FlowConfig], str]:
        """List ingestion config flows (single page).

        Args:
            ingestion_config_id: The ingestion config ID.
            page_size: Number of results per page.
            page_token: Token for the next page.
            order_by: Unused; accepted for _handle_pagination compatibility.
            filter_query: Optional CEL filter for flows.

        Returns:
            A tuple of (list of FlowConfig, next_page_token).
        """
        request_kwargs: dict[str, Any] = {"ingestion_config_id": ingestion_config_id}

        if page_size is not None:
            request_kwargs["page_size"] = page_size
        if page_token is not None:
            request_kwargs["page_token"] = page_token
        if filter_query is not None:
            request_kwargs["filter"] = filter_query
        if order_by is not None:
            request_kwargs["order_by"] = order_by

        request = ListIngestionConfigFlowsRequest(**request_kwargs)
        res = await self._grpc_client.get_stub(IngestionConfigServiceStub).ListIngestionConfigFlows(
            request
        )
        res = cast("ListIngestionConfigFlowsResponse", res)
        flows = [FlowConfig._from_proto(flow) for flow in res.flows]
        return flows, res.next_page_token

    async def get_ingestion_config_flows(
        self,
        ingestion_config_id: str,
        page_size: int | None = DEFAULT_INGESTION_CONFIG_PAGE_SIZE,
        max_results: int | None = None,
    ) -> list[FlowConfig]:
        """Get all flows for an ingestion config, using pagination.

        Args:
            ingestion_config_id: The ingestion config ID.
            page_size: Number of results per page.
            max_results: Maximum total results to return; None for no limit.

        Returns:
            A list of all FlowConfigs for the ingestion config.
        """
        return await self._handle_pagination(
            self.list_ingestion_config_flows,
            kwargs={"ingestion_config_id": ingestion_config_id},
            page_size=page_size,
            max_results=max_results,
        )

    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_all_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 _hash_flows(self, asset_name: str, flows: list[FlowConfig]) -> str:
        """Generate a client key that should be unique but deterministic for the given asset and flow configuration."""
        return _hash_flows(asset_name=asset_name, flows=flows)


class IngestionConfigStreamingLowLevelClient(LowLevelClientBase):
    DEFAULT_MAX_LOG_FILES = 7  # Equal to 1 week of logs
    DEFAULT_LOGFILE_PREFIX = "sift_stream_bindings.log"
    _sift_stream_instance: SiftStreamPy

    def __init__(self, sift_stream_instance: SiftStreamPy):
        super().__init__()
        self._sift_stream_instance = sift_stream_instance

    @classmethod
    async def create_sift_stream_instance(
        cls,
        api_key: str,
        grpc_uri: str,
        ingestion_config_form: IngestionConfigFormPy,
        run_form: RunFormPy | None = None,
        run_id: str | None = None,
        asset_tags: list[str] | None = None,
        asset_metadata: list[MetadataPy] | None = None,
        streaming_mode: StreamingMode = ...,  # type: ignore[assignment]
        retry_policy: RetryPolicyPy | None = None,
        disk_backup_policy: DiskBackupPolicyPy | None = None,
        checkpoint_interval_seconds: int | None = None,
        enable_tls: bool = True,
        tracing_config: TracingConfig | None = None,
    ) -> IngestionConfigStreamingLowLevelClient:
        # Importing here to allow sift_stream_bindings to be an optional dependancy for non-ingestion users
        # TODO(nathan): Fix bindings to fix mypy issues with tracing functions
        from sift_stream_bindings import (  # type: ignore[attr-defined]
            DurationPy,
            SiftStreamBuilderPy,
            init_tracing,  # type: ignore[attr-defined]
            init_tracing_with_file,  # type: ignore[attr-defined]
            is_tracing_initialized,  # type: ignore[attr-defined]
        )

        from sift_client.resources.ingestion import StreamingMode, TracingConfig

        if streaming_mode is ...:  # type: ignore[comparison-overlap]
            streaming_mode = StreamingMode.LIVE_WITH_BACKUPS

        if not is_tracing_initialized():
            if tracing_config is None:
                tracing_config = TracingConfig.with_file()

            if tracing_config.log_dir is not None:
                # Use file logging
                init_tracing_with_file(
                    tracing_config.level,
                    tracing_config.log_dir,
                    tracing_config.filename_prefix or cls.DEFAULT_LOGFILE_PREFIX,
                    tracing_config.max_log_files or cls.DEFAULT_MAX_LOG_FILES,
                )
            else:
                # Use stdout/stderr only
                init_tracing(tracing_config.level)

        sift_builder = SiftStreamBuilderPy(uri=grpc_uri, apikey=api_key)
        sift_builder.enable_tls = enable_tls

        config_builder = sift_builder.ingestion_config(ingestion_config_form)
        config_builder.run = run_form
        config_builder.run_id = run_id
        config_builder.asset_tags = asset_tags
        config_builder.metadata = asset_metadata

        if streaming_mode == StreamingMode.LIVE_ONLY:
            sift_stream_instance = await config_builder.live_only().build()

        elif streaming_mode == StreamingMode.FILE_BACKUP:
            fb_builder = config_builder.file_backup()
            if disk_backup_policy is not None:
                fb_builder.disk_backup_policy = disk_backup_policy
            sift_stream_instance = await fb_builder.build()

        else:  # LIVE_WITH_BACKUPS (default)
            lwb_builder = config_builder.live_with_backups()
            if retry_policy is not None:
                lwb_builder.retry_policy = retry_policy
            if disk_backup_policy is not None:
                lwb_builder.disk_backup_policy = disk_backup_policy
            if checkpoint_interval_seconds is not None:
                lwb_builder.checkpoint_interval = DurationPy(
                    secs=checkpoint_interval_seconds, nanos=0
                )
            sift_stream_instance = await lwb_builder.build()

        return cls(sift_stream_instance)

    async def send(self, flow: FlowPy):
        await self._sift_stream_instance.send(flow)

    async def batch_send(self, flows: Iterable[FlowPy]):
        await self._sift_stream_instance.batch_send(flows)

    async def send_requests(self, requests: list[IngestWithConfigDataStreamRequestPy]):
        await self._sift_stream_instance.send_requests(requests)

    def try_send_requests(
        self, requests: Iterable[IngestWithConfigDataStreamRequestWrapperPy]
    ) -> None:
        self._sift_stream_instance.try_send_requests(requests)

    def try_send(self, flow: FlowPy) -> None:
        self._sift_stream_instance.try_send(flow)

    def get_flow_descriptor(self, flow_name: str) -> FlowDescriptorPy:
        return self._sift_stream_instance.get_flow_descriptor(flow_name)

    async def add_new_flows(self, flow_configs: list[FlowConfigPy]):
        await self._sift_stream_instance.add_new_flows(flow_configs)

    async def attach_run(self, run_selector: RunSelectorPy):
        await self._sift_stream_instance.attach_run(run_selector)

    def detach_run(self):
        self._sift_stream_instance.detach_run()

    def get_run_id(self) -> str | None:
        return self._sift_stream_instance.run()

    async def finish(self):
        await self._sift_stream_instance.finish()

    def get_metrics_snapshot(self) -> SiftStreamMetricsSnapshotPy:
        return self._sift_stream_instance.get_metrics_snapshot()


def _hash_flows(asset_name: str, flows: list[FlowConfig]) -> 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()