openaleph_procrastinate.settings

class DeferSettings(BaseSettings):
    """
    Adjust the worker queues and tasks for different stages.

    This is useful e.g. for launching a priority queuing setup for a specific dataset:

    Example:
        ```bash
        # ingest service
        export OPENALEPH_INGEST_QUEUE=ingest-prio-dataset
        export OPENALEPH_ANALYZE_QUEUE=analyze-prio-dataset
        ingestors ingest -d prio_dataset ./documents
        procrastinate worker -q ingest-prio-dataset --one-shot  # stop worker after complete

        # analyze service
        procrastinate worker -q analyze-prio-dataset --one-shot  # stop worker after complete
        ```
    """

    model_config = SettingsConfigDict(
        env_prefix="openaleph_",
        env_nested_delimiter="_",
        env_file=".env",
        nested_model_default_partial_update=True,
        extra="ignore",  # other envs in .env file
    )

    ingest: ServiceSettings = ServiceSettings(
        queue="ingest", task="ingestors.tasks.ingest", dedupe=True
    )
    """ingest-file. Deduplicated: the same file is reached through several
    parents, and a failing job replays the fan-out that queued it"""

    analyze: ServiceSettings = ServiceSettings(
        queue="analyze", task="ftm_analyze.tasks.analyze"
    )
    """ftm-analyze"""

    transcribe: ServiceSettings = ServiceSettings(
        queue="transcribe", task="ftm_transcribe.tasks.transcribe", defer=False
    )
    """ftm-transcribe"""

    translate: ServiceSettings = ServiceSettings(
        queue="translate", task="ftm_translate.tasks.translate", defer=False
    )
    """ftm-translate"""

    geocode: ServiceSettings = ServiceSettings(
        queue="geocode", task="ftm_geocode.tasks.geocode", defer=False
    )
    """ftm-geocode"""

    assets: ServiceSettings = ServiceSettings(
        queue="assets", task="ftm_assets.tasks.resolve", defer=False
    )
    """ftm-assets"""

    vectorize: ServiceSettings = ServiceSettings(
        queue="vectorize", task="ftm_vectorize.tasks.vectorize", defer=False
    )
    """ftm-vectorize"""

    # OpenAleph

    index: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.index_entities",
        min_priority=70,
    )
    """openaleph entity indexer"""

    reindex: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.reindex_collection",
        min_priority=50,
    )
    """openaleph collection reindexer"""

    xref: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.xref_collection",
        min_priority=50,
    )
    """openaleph xref collection"""

    load_mapping: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.load_mapping",
        min_priority=90,
    )
    """openaleph load_mapping"""

    flush_mapping: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.flush_mapping",
        min_priority=40,
    )
    """openaleph flush_mapping"""

    export_search: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.export_search",
        max_priority=50,
    )
    """openaleph export_search"""

    export_xref: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.export_xref",
        max_priority=50,
    )
    """openaleph export_xref"""

    update_entity: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.update_entity",
        min_priority=99,
    )
    """openaleph update_entity"""

    prune_entity: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_QUEUE,
        task="aleph.procrastinate.tasks.prune_entity",
        min_priority=99,
    )
    """openaleph prune_entity"""

    cancel_dataset: ServiceSettings = ServiceSettings(
        queue=OPENALEPH_MANAGEMENT_QUEUE,
        task="aleph.procrastinate.tasks.cancel_dataset",
        min_priority=101,
    )
    """openaleph cancel dataset processing"""

class OpenAlephSettings(BaseSettings):
    """
    `openaleph_procrastinate` settings management using
    [pydantic-settings](https://docs.pydantic.dev/latest/concepts/pydantic_settings/)

    Note:
        All settings can be set via environment variables, prepending
        `OPENALEPH_` (except for those with another alias) via runtime or in a
        `.env` file.
    """

    model_config = SettingsConfigDict(
        env_prefix="openaleph_",
        env_nested_delimiter="_",
        env_file=".env",
        nested_model_default_partial_update=True,
        extra="ignore",  # other envs in .env file
    )

    instance: str = Field(default="openaleph", validation_alias="openaleph_instance")
    """Instance identifier"""

    debug: bool = Field(default=False, alias="debug")
    """Debug mode"""

    procrastinate_sync: bool = Field(default=False, alias="procrastinate_sync")
    """Run sync workers (during testing)"""

    db_uri: str = Field(
        default=DEFAULT_DB_URI,
        validation_alias=AliasChoices("openaleph_db_uri", "aleph_database_uri"),
    )
    """OpenAleph database uri"""

    db_pool_size: int = Field(default=5, validation_alias="openaleph_db_pool_size")
    """Max psql pool size per thread"""

    procrastinate_db_uri: str = Field(
        default=DEFAULT_DB_URI,
        validation_alias=AliasChoices(
            "procrastinate_db_uri", "openaleph_db_uri", "aleph_database_uri"
        ),
    )
    """Procrastinate database uri, falls back to OpenAleph database uri"""

    fragments_uri: str = Field(
        default=DEFAULT_DB_URI,
        validation_alias=AliasChoices("ftm_fragments_uri", "ftm_store_uri"),
    )
    """FollowTheMoney Fragments store uri"""

    redis_url: str | None = Field(default=None, validation_alias="redis_url")
    """Redis instance uri"""

    procrastinate_dehydrate_entities: bool = Field(
        default=True, validation_alias="openaleph_procrastinate_dehydrate_entities"
    )
    """Dehydrate entity in job payload, jobs need to re-fetch entity from store"""

    lakehouse: bool = Field(default=False, validation_alias="openaleph_lakehouse")
    """Activate lakehouse storage backend (experimental)"""

    lakehouse_flush_threshold: int = Field(
        default=100_000, validation_alias="openaleph_lakehouse_flush_threshold"
    )
    """Auto-flush the lakehouse journal to parquet after this number of
    processed entities (needs a traced task, see `tracer` module)"""

    trace_entities: bool = Field(
        default=True, validation_alias="openaleph_trace_entities"
    )
    """Track the per-entity status in traced tasks (see `tracer` module). The
    task counters (e.g. for the lakehouse auto-flush) are tracked anyway."""

    @model_validator(mode="after")
    def enforce_lakehouse_payloads(self) -> Self:
        """The lakehouse backend takes entities fully from the job payload
        instead of re-fetching them from the store, so dehydrating is off."""
        if self.lakehouse and self.procrastinate_dehydrate_entities:
            log.warning(
                "Disabling `procrastinate_dehydrate_entities`: the lakehouse "
                "backend needs the full entity in the job payload."
            )
            self.procrastinate_dehydrate_entities = False
        return self

    @property
    def in_memory_db(self) -> bool:
        return (
            self.procrastinate_db_uri.startswith("memory")
            or urlparse(self.procrastinate_db_uri).path == "/:memory:"
        )

class ServiceSettings(BaseSettings):
    """
    Settings for a specific service, like `ingest-file` or `ftm-analyze`
    """

    queue: str
    """queue name"""
    task: str
    """task module path"""
    defer: bool = True
    """enable deferring"""
    max_retries: int = 5
    """Max retries, set to "-1" to enable infinity"""
    dedupe: bool = False
    """Collapse a deferred job with an identical one still waiting in the queue"""
    min_priority: int = MIN_PRIORITY
    """Minimum priority"""
    max_priority: int = MAX_PRIORITY
    """Maximum priority"""

    @property
    def retries(self) -> int | bool:
        if self.max_retries == -1:
            return True
        return max(0, self.max_retries)

    def get_priority(self, priority: int | None = None) -> int:
        """Calculate a random priority between `min_priority` and
        `max_priority`"""
        min_priority = max(priority or MIN_PRIORITY, self.min_priority)
        max_priority = max(min_priority, self.max_priority)
        return random.randint(min_priority, max_priority)

    def get_tracer(self, dataset: str, uri: Uri | None = None) -> "Tracer":
        """Get task status tracer."""
        from openaleph_procrastinate.tracer import get_tracer

        return get_tracer(dataset, self.queue, self.task, uri)