openaleph_procrastinate.helpers

@contextmanager
def entity_writer(
    dataset: str, origin: str = OPAL_ORIGIN
) -> Generator[EntityStore, None, None]:
    """
    Get the [`EntityStore`][openaleph_procrastinate.repository.EntityStore] for
    the given `dataset` to write to. It is flushed and closed when leaving the
    context – a store that buffers writes holds a database connection until
    then, so callers that build one themselves have to close it.
    """
    store = get_entity_store(dataset, origin)
    try:
        yield store
    finally:
        store.close()

def get_localpath(dataset: str, content_hash: str) -> ContextManager[Path]:
    """
    Load a file from the archive and store it in a local temporary path for
    further processing. The file is cleaned up after leaving the context.
    [Reference][openaleph_procrastinate.model.DatasetJob.get_file_references]
    """
    return get_archive(dataset).local_path(content_hash)

def load_entities(dataset: str, entity_ids: Iterable[str]) -> EntityProxies:
    """
    Batch retrieve entities from the entity store.
    """
    yield from get_entity_store(dataset).iterate(entity_ids)

def open_file(dataset: str, content_hash: str) -> ContextManager[VirtualIO]:
    """
    Load a file from the archive and store it in a local temporary path for
    further processing. Returns an open file handler. The file is closed and
    cleaned up after leaving the context.
    [Reference][openaleph_procrastinate.model.DatasetJob.get_file_references]
    """
    return get_archive(dataset).open(content_hash)