openaleph_procrastinate.model

class DatasetJob(Job):
    """
    A job with arbitrary payload bound to a `dataset`.
    The payload always contains an iterable of serialized `EntityProxy` objects
    in the `entities` key. It may contain other payload context data in the
    `context` key.

    There are helpers for accessing archive files or loading entities.
    """

    dataset: str
    batch: str | None = None

    @property
    def log(self) -> BoundLogger:
        return get_logger(
            name=f"{self.dataset}.job",
            dataset=self.dataset,
            queue=self.queue,
            task=self.task,
            batch=self.batch,
        )

    def get_writer(
        self: Self, origin: str = helpers.OPAL_ORIGIN
    ) -> ContextManager[EntityStore]:
        """Get the writer for the dataset of the current job"""
        return helpers.entity_writer(self.dataset, origin)

    def get_queueing_lock(self: Self) -> str | None:
        """A lock that identifies this job by its task, dataset and entities, so
        procrastinate can collapse it with an identical job that is still
        waiting. `None` when the payload carries no entities to key on.

        Returns:
            The queueing lock
        """
        ids = sorted({e["id"] for e in self.payload.get("entities", [])})
        if not ids:
            return None
        key = ",".join(ids)
        if len(key) > MAX_LOCK_KEY:
            # the unique index on `queueing_lock` is a btree, which rejects an
            # entry beyond roughly 2700 bytes
            key = make_data_checksum(ids)
        return f"{self.task}:{self.dataset}:{key}"

    def get_entities(self) -> Generator[EntityProxy, None, None]:
        """
        Get the entities from the payload
        """
        assert "entities" in self.payload, "No entities in payload"
        for data in self.payload["entities"]:
            yield model.get_proxy(data)

    def load_entities(self: Self) -> Generator[EntityProxy, None, None]:
        """Load the entities from the store to refresh it to the latest data"""
        assert "entities" in self.payload, "No entities in payload"

        # if not dehydrating, just get the payload data:
        if not settings.procrastinate_dehydrate_entities:
            yield from self.get_entities()
        else:
            entity_ids = [e["id"] for e in self.payload["entities"]]
            yield from helpers.load_entities(self.dataset, entity_ids)

    # Helpers for file jobs that access the servicelayer archive

    def get_file_references(self) -> Generator[EntityFileReference, None, None]:
        """
        Get file references per entity from this job

        Example:
            ```python
            # process temporary file paths
            for reference in job.get_file_references():
                with reference.get_local_path() as path:
                    subprocess.run(["command", "-i", str(path)])
                # temporary path will be cleaned up when leaving context

            # process temporary file handlers
            for reference in job.get_file_references():
                with reference.open() as handler:
                    do_something(handler.read())
                # temporary path will be cleaned up when leaving context
            ```

        Yields:
            The file references
        """
        for entity in self.get_entities():
            for content_hash in entity.get("contentHash", quiet=True):
                yield EntityFileReference(
                    dataset=self.dataset, entity=entity, content_hash=content_hash
                )

    @classmethod
    def from_entities(
        cls,
        dataset: str,
        queue: str,
        task: str,
        entities: Iterable[EntityProxy],
        dehydrate: bool = settings.procrastinate_dehydrate_entities,
        **context: Any,
    ) -> Self:
        """
        Make a job to process entities for a dataset

        Args:
            dataset: Name of the dataset
            queue: Name of the queue
            task: Python module path of the task
            entities: Entities
            dehydrate: Reduce entity payload to only a reference (tasks should
                re-fetch these entities from the store if they need more data).
                Ignored when `procrastinate_dehydrate_entities` is disabled
            context: Job context
        """
        dehydrate = dehydrate and settings.procrastinate_dehydrate_entities
        if dehydrate:
            entities_ = (make_file_entity(e, quiet=True) for e in entities)
            entities = (e for e in entities_ if e is not None)
        return cls(
            dataset=dataset,
            queue=queue,
            task=task,
            batch=context.pop("batch", None),
            payload={
                "entities": [e.to_dict() for e in entities],
                "context": ensure_dict(context),
            },
        )

class EntityFileReference(BaseModel):
    """
    A file reference (via `content_hash`) from an entity to a file in the
    [archive][openaleph_procrastinate.repository.Archive]
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    dataset: str
    content_hash: str
    entity: EntityProxy

    def open(self: Self) -> ContextManager[VirtualIO]:
        """
        Open the file attached to this job
        """
        return helpers.open_file(self.dataset, self.content_hash)

    def get_localpath(self: Self) -> ContextManager[Path]:
        """
        Get a temporary path for the file attached to this job
        """
        return helpers.get_localpath(self.dataset, self.content_hash)

class Job(JobModel):
    """
    A job with arbitrary payload
    """

    @property
    def context(self) -> dict[str, Any]:
        """Get the context from the payload if any"""
        return ensure_dict(self.payload.get("context")) or {}

    @property
    def log(self) -> BoundLogger:
        return get_logger(name="openaleph.job", queue=self.queue, task=self.task)

    def defer(
        self: Self,
        app: App,
        priority: int | None = None,
        queueing_lock: str | None = None,
    ) -> bool:
        """Defer this job.

        Args:
            app: The procrastinate app instance
            priority: Job priority
            queueing_lock: Refuse the job while one with the same lock is still
                waiting in the queue

        Returns:
            Whether the job was deferred
        """
        self.log.debug("Deferring ...", payload=self.payload)
        data = clean_dict(self.model_dump(mode="json"))
        payload = data.get("payload", {})
        if isinstance(payload.get("entities"), list):
            payload["entities"] = pack_entities(payload["entities"])
        try:
            app.configure_task(
                name=self.task,
                queue=self.queue,
                priority=priority or get_priority(),
                queueing_lock=queueing_lock,
            ).defer(**data)
        except AlreadyEnqueued:
            self.log.info("Already queued, not deferring again", lock=queueing_lock)
            return False
        if settings.debug:
            # option to change synchronousness during test runtime
            _settings = OpenAlephSettings()
            if _settings.procrastinate_sync:
                # run worker synchronously (for testing)
                run_sync_worker(app)
        return True

class JobModel(BaseModel):
    """
    A job with arbitrary payload
    """

    queue: str
    task: str
    payload: dict[str, Any] = {}

    # from procrastinate table:
    id: int | None = None
    status: Status | None = None
    scheduled_at: datetime | None = None