Source code for satorbis_kit.pgstac.models.config

"""Ingestion configuration model for STAC workflows."""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field




class IngestionConfig(BaseModel):
    """Configuration for STAC ingestion workflow.

    This model represents the configuration sent to Airflow DAG for STAC ingestion.
    It includes all necessary parameters for raster stacking and STAC catalog ingestion.

    Attributes:
        raster_s3_urls: List of S3 URLs to raster files
        collection: STAC collection name
        costing_labels: Costing labels for tracking (always included)
        ingestion_batch_size: Batch size for ingestion (default: 100)
        parameters: Optional COG conversion parameters (default: {})

    Example:
        >>> config = IngestionConfig(
        ...     raster_s3_urls=["s3://bucket/TEST_COL/path/20240401.tif"],
        ...     collection="TEST_COL",
        ...     costing_labels=DEFAULT_COSTING_LABELS,
        ... )
        >>> payload = config.to_payload()
    """

    raster_s3_urls: List[str]
    collection: str
    costing_labels: Dict[str, str]
    ingestion_batch_size: Optional[int] = Field(default=100)
    parameters: Optional[Dict[str, Any]] = Field(default_factory=dict)



    def to_payload(self) -> Dict[str, Any]:
        """Convert to Airflow DAG payload format.

        Transforms the configuration into the exact format expected by
        the Airflow stac_item_batch_creation_dag. The configuration is
        wrapped in a "conf" key as per Airflow DAG trigger requirements.

        Returns:
            Dictionary payload ready for Airflow API with structure:
            {
                "conf": {
                    "costing_labels": {...},
                    "raster_s3_urls": [...],
                    "collection": "...",
                    "ingestion_batch_size": 100,
                    "parameters": {...}
                }
            }

        Example:
            >>> config = IngestionConfig(...)
            >>> payload = config.to_payload()
            >>> # payload is ready to send to Airflow
        """
        # Use defaults if None
        batch_size = (
            self.ingestion_batch_size if self.ingestion_batch_size is not None else 100
        )
        params = self.parameters if self.parameters is not None else {}

        payload = {
            "conf": {
                "costing_labels": self.costing_labels,
                "raster_s3_urls": self.raster_s3_urls,
                "collection": self.collection,
                "ingestion_batch_size": batch_size,
                "parameters": params,
            }
        }

        return payload