From 9233f178a9e0a39584c0a33ee80091869c10dd49 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Thu, 2 Jul 2026 17:29:55 -0400 Subject: [PATCH 1/3] feat: starting implementation of the workers --- pyproject.toml | 1 + src/archivist/__main__.py | 4 + src/archivist/api/archive.py | 17 ++-- src/archivist/api/extract.py | 4 +- src/archivist/database.py | 68 ++++++++++++++++ src/archivist/orm/__init__.py | 1 + src/archivist/orm/archivequeue.py | 128 ++++++++++++++++++++++++++++++ src/archivist/queue.py | 32 ++------ src/archivist/server.py | 23 ++++++ src/archivist/tasks/__init__.py | 0 src/archivist/tasks/archive.py | 25 ++++++ src/archivist/utils.py | 13 +++ 12 files changed, 281 insertions(+), 35 deletions(-) create mode 100644 src/archivist/database.py create mode 100644 src/archivist/orm/__init__.py create mode 100644 src/archivist/orm/archivequeue.py create mode 100644 src/archivist/tasks/__init__.py create mode 100644 src/archivist/tasks/archive.py create mode 100644 src/archivist/utils.py diff --git a/pyproject.toml b/pyproject.toml index 2e2357e..b662569 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,5 +110,6 @@ fixture-parentheses = false mark-parentheses = false [tool.isort] +profile = "black" line_length = 120 skip = ["_version.py"] diff --git a/src/archivist/__main__.py b/src/archivist/__main__.py index 318ecec..9cb86e8 100644 --- a/src/archivist/__main__.py +++ b/src/archivist/__main__.py @@ -132,7 +132,11 @@ def start_server(ctx): import uvicorn + from .database import create_all from .settings import server_settings + from .tasks.archive import start_archive + + create_all() uvicorn.run( "archivist.server:main", diff --git a/src/archivist/api/archive.py b/src/archivist/api/archive.py index 9703223..f5636ff 100644 --- a/src/archivist/api/archive.py +++ b/src/archivist/api/archive.py @@ -2,10 +2,12 @@ from fastapi import Depends, Response from loguru import logger +from sqlalchemy.orm import Session from archivist.api import router -from archivist.core.models import Archive, ManifestFailedResponse, ManifestRequest, ManifestResponse -from archivist.queue import ArchiveQueue, get_archive_queue +from archivist.core.models import ManifestFailedResponse, ManifestRequest, ManifestResponse +from archivist.database import yield_session +from archivist.orm.archivequeue import ArchiveQueue from archivist.settings import Settings, get_settings @@ -13,7 +15,7 @@ def archive( manifest_request: ManifestRequest, response: Response, - queue: ArchiveQueue = Depends(get_archive_queue), + session: Session = Depends(yield_session), settings: Settings = Depends(get_settings), ): """ @@ -36,15 +38,14 @@ def archive( f"Archiving manifest {manifest_id} from librarian '{manifest_request.librarian_name}': {len(manifest_request.store_files)} file(s), {total_size} bytes total" ) - archive = Archive( + item = ArchiveQueue.new_item( manifest_id=str(manifest_id), manifest=manifest_request.model_dump_json(), paths=[entry.instance_path for entry in manifest_request.store_files], - root=settings.storage_root, # Example root path + root=settings.storage_root, ) - - # TODO: Add the archive to a queue for processing - # queue.enqueue_archive(archive) + session.add(item) + session.commit() return ManifestResponse( manifest_id=str(manifest_id), diff --git a/src/archivist/api/extract.py b/src/archivist/api/extract.py index a203e89..8dfccc8 100644 --- a/src/archivist/api/extract.py +++ b/src/archivist/api/extract.py @@ -5,7 +5,7 @@ from archivist.api import router from archivist.core.models import Archive, ManifestFailedResponse, ManifestRequest, ManifestResponse -from archivist.queue import ExtractQueue, get_extract_queue +from archivist.queue import Queue, get_extract_queue from archivist.settings import Settings, get_settings @@ -13,7 +13,7 @@ def extract( manifest_request: ManifestRequest, response: Response, - queue: ExtractQueue = Depends(get_extract_queue), + queue: Queue = Depends(get_extract_queue), settings: Settings = Depends(get_settings), ): """ diff --git a/src/archivist/database.py b/src/archivist/database.py new file mode 100644 index 0000000..d56aeb4 --- /dev/null +++ b/src/archivist/database.py @@ -0,0 +1,68 @@ +""" +Core database runner for SQLAlchemy. +""" + +from loguru import logger +from sqlalchemy import ( + BigInteger, + Boolean, + Column, + DateTime, + Enum, + ForeignKey, + Integer, + PickleType, + String, + create_engine, +) +from sqlalchemy.orm import declarative_base, sessionmaker + +from .settings import server_settings + +logger.info("Starting database engine.") + +# Create the engine and sessionmaker. +engine = create_engine( + server_settings.sqlalchemy_database_uri, + # Required for async and SQLite + connect_args=({"check_same_thread": False} if "sqlite" in server_settings.sqlalchemy_database_uri else {}), +) + +logger.info("Creating database session.") + +SessionMaker = sessionmaker(bind=engine, autocommit=False, autoflush=False) + + +def yield_session() -> SessionMaker: + """ + Yields a new databse session. + """ + + session = SessionMaker() + try: + yield session + finally: + session.close() + + +def get_session() -> SessionMaker: + """ + Returns a new database session. Unlike yield_session, it is + your responsibility to close the session. + """ + + return SessionMaker() + + +Base = declarative_base() + + +def create_all() -> None: + """ + Create all tables that don't already exist. Imports the ORM models + first so their table definitions are registered with `Base`. + """ + + from . import orm # noqa: F401 + + Base.metadata.create_all(engine) diff --git a/src/archivist/orm/__init__.py b/src/archivist/orm/__init__.py new file mode 100644 index 0000000..2446994 --- /dev/null +++ b/src/archivist/orm/__init__.py @@ -0,0 +1 @@ +from .archivequeue import ArchiveQueue # noqa: F401 diff --git a/src/archivist/orm/archivequeue.py b/src/archivist/orm/archivequeue.py new file mode 100644 index 0000000..415d5e1 --- /dev/null +++ b/src/archivist/orm/archivequeue.py @@ -0,0 +1,128 @@ +""" +The ORM definition for the queue of archive jobs awaiting processing. + +When the API receives a manifest to archive, it does not perform the +archive operation synchronously. Instead, it places an item in this +queue (backed by the database, so it survives restarts) and a separate +worker thread picks items up and performs the actual archive operation. +""" + +import datetime + +from sqlalchemy.orm import Session + +from .. import database as db + + +class ArchiveQueue(db.Base): + """ + The queue of archive jobs awaiting processing. + """ + + __tablename__ = "archive_queue" + + id = db.Column(db.Integer, primary_key=True) + "The unique identifier for this item in the queue." + created_time = db.Column(db.DateTime, nullable=False) + "The time that this item was added to the queue." + retries = db.Column(db.Integer, nullable=False, default=0) + "The number of times this item has been tried to be archived." + + manifest_id = db.Column(db.String(256), nullable=False, unique=True) + "The ID of the manifest associated with this archive." + manifest = db.Column(db.PickleType, nullable=False) + "The raw manifest (JSON) that was submitted for archiving." + paths = db.Column(db.PickleType, nullable=False) + "The list of file / directory paths contained in the archive." + root = db.Column(db.String(512), nullable=False) + "The common root filesystem location for building relative paths." + + consumed = db.Column(db.Boolean, default=False) + "Whether this queue item has been picked up by a worker." + consumed_time = db.Column(db.DateTime) + "The time at which the item was picked up. Useful for pruning." + + completed = db.Column(db.Boolean, default=False) + "Whether this queue item has been entirely completed." + completed_time = db.Column(db.DateTime) + "The time at which the queue item was marked as completed." + + failed = db.Column(db.Boolean, default=False) + "Whether this queue item failed, and that is the reason for completed status." + + @classmethod + def new_item(cls, manifest_id: str, manifest: str, paths: list[str], root: str) -> "ArchiveQueue": + """ + Create a new item in the queue from an incoming archive request. + + Parameters + ---------- + manifest_id : str + The ID of the manifest associated with this archive. + manifest : str + The raw manifest (JSON) that was submitted for archiving. + paths : list[str] + The list of file / directory paths contained in the archive. + root : str + The common root filesystem location for building relative paths. + """ + + return cls( + manifest_id=manifest_id, + manifest=manifest, + paths=paths, + root=root, + created_time=datetime.datetime.now(datetime.timezone.utc), + retries=0, + ) + + @classmethod + def dequeue(cls, session: Session) -> "ArchiveQueue | None": + """ + Pop the oldest unconsumed item off the queue and mark it as consumed. + + Parameters + ---------- + session : Session + The database session to use. + """ + + item = session.query(cls).filter_by(consumed=False).order_by(cls.created_time.asc()).first() + + if item is not None: + item.consumed = True + item.consumed_time = datetime.datetime.now(datetime.timezone.utc) + session.commit() + + return item + + def complete(self, session: Session): + """ + Mark this queue item as successfully completed. + + Parameters + ---------- + session : Session + The database session to use. + """ + + self.completed = True + self.completed_time = datetime.datetime.now(datetime.timezone.utc) + + session.commit() + + def fail(self, session: Session): + """ + Mark this queue item as failed. + + Parameters + ---------- + session : Session + The database session to use. + """ + + self.failed = True + self.completed = True + self.completed_time = datetime.datetime.now(datetime.timezone.utc) + + session.commit() diff --git a/src/archivist/queue.py b/src/archivist/queue.py index f6ed0e3..d44a722 100644 --- a/src/archivist/queue.py +++ b/src/archivist/queue.py @@ -2,11 +2,11 @@ from archivist.core.models import Archive -_archive_queue: "ArchiveQueue | None" = None -_extract_queue: "ExtractQueue | None" = None +_archive_queue: "Queue | None" = None +_extract_queue: "Queue | None" = None -class ArchiveQueue: +class Queue: """Thread-safe queue for Archive jobs awaiting processing.""" def __init__(self) -> None: @@ -25,36 +25,18 @@ def task_done(self) -> None: def size(self) -> int: return self._q.qsize() - -class ExtractQueue: - """Thread-safe queue for Extract jobs awaiting processing.""" - - def __init__(self) -> None: - self._q: queue.Queue[Archive] = queue.Queue() - - def enqueue_archive(self, archive: Archive) -> None: - self._q.put(archive) - - def dequeue_archive(self, block: bool = True, timeout: float | None = None) -> Archive: - return self._q.get(block=block, timeout=timeout) - - def task_done(self) -> None: - self._q.task_done() - - @property - def size(self) -> int: return self._q.qsize() -def get_archive_queue() -> ArchiveQueue: +def get_archive_queue() -> Queue: global _archive_queue if _archive_queue is None: - _archive_queue = ArchiveQueue() + _archive_queue = Queue() return _archive_queue -def get_extract_queue() -> ExtractQueue: +def get_extract_queue() -> Queue: global _extract_queue if _extract_queue is None: - _extract_queue = ExtractQueue() + _extract_queue = Queue() return _extract_queue diff --git a/src/archivist/server.py b/src/archivist/server.py index 36d0593..75e493b 100644 --- a/src/archivist/server.py +++ b/src/archivist/server.py @@ -2,12 +2,28 @@ The Archivist v2.0 server. """ +import threading +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from fastapi import FastAPI from .settings import server_settings +_archive_stop_event = threading.Event() + + +def _archive_worker_loop(): + from loguru import logger + + from .tasks.archive import start_archive + + while not _archive_stop_event.is_set(): + try: + start_archive(librarian_name=server_settings.name) + except Exception: + logger.exception("Archive worker iteration failed") + @asynccontextmanager async def slack_post_at_startup_shutdown(app: FastAPI): @@ -18,8 +34,15 @@ async def slack_post_at_startup_shutdown(app: FastAPI): from loguru import logger logger.info("Archivist server starting up") + + archive_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="archive-worker") + archive_pool.submit(_archive_worker_loop) + yield + logger.info("Archivist server shutting down") + _archive_stop_event.set() + archive_pool.shutdown(wait=True) def main() -> FastAPI: diff --git a/src/archivist/tasks/__init__.py b/src/archivist/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/archivist/tasks/archive.py b/src/archivist/tasks/archive.py new file mode 100644 index 0000000..c1f3cea --- /dev/null +++ b/src/archivist/tasks/archive.py @@ -0,0 +1,25 @@ +from time import sleep + +import loguru +from sqlalchemy import inspect + +from archivist.database import get_session +from archivist.orm.archivequeue import ArchiveQueue +from archivist.settings import get_settings + + +def start_archive(librarian_name: str, store_files: bool = True): + """ + Start an archive process by sending a request to the archivist server. + """ + settings = get_settings() + + session = get_session() + try: + archive_item = ArchiveQueue.dequeue(session) + if archive_item is not None: + archive_dict = {c.key: getattr(archive_item, c.key) for c in inspect(archive_item).mapper.column_attrs} + loguru.logger.info(f"Archive Item: {archive_dict}") + finally: + session.close() + sleep(1) diff --git a/src/archivist/utils.py b/src/archivist/utils.py new file mode 100644 index 0000000..3442438 --- /dev/null +++ b/src/archivist/utils.py @@ -0,0 +1,13 @@ +import hashlib +from pathlib import Path + + +def _checksum(path: Path) -> str: + """Compute the sha256 checksum of a file.""" + + hasher = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(chunk) + + return hasher.hexdigest() From 286a8db3c22c0095630a3e8f80ddccefa385e8b0 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Tue, 7 Jul 2026 15:54:59 -0400 Subject: [PATCH 2/3] feat: posix archive works --- src/archivist/__init__.py | 2 +- src/archivist/__main__.py | 20 ++------ src/archivist/api/archive.py | 2 +- src/archivist/api/extract.py | 8 ++-- src/archivist/core/__init__.py | 4 +- src/archivist/core/archive.py | 50 +++++++++++++++----- src/archivist/core/models.py | 14 ------ src/archivist/core/storage.py | 16 ++++--- src/archivist/database.py | 62 +++++++++++++++++------- src/archivist/queue.py | 35 +++++--------- src/archivist/server.py | 22 +++++++-- src/archivist/settings.py | 8 +++- src/archivist/storage/__init__.py | 4 ++ src/archivist/storage/base.py | 49 +++++++++++++++++++ src/archivist/storage/storage_disk.py | 63 +++++++++++++++++++++++++ src/archivist/storage/storage_hpss.py | 20 ++++++++ src/archivist/tasks/archive.py | 68 +++++++++++++++++++++++++-- 17 files changed, 344 insertions(+), 103 deletions(-) create mode 100644 src/archivist/storage/__init__.py create mode 100644 src/archivist/storage/base.py create mode 100644 src/archivist/storage/storage_disk.py create mode 100644 src/archivist/storage/storage_hpss.py diff --git a/src/archivist/__init__.py b/src/archivist/__init__.py index 9e2db2a..d86af70 100644 --- a/src/archivist/__init__.py +++ b/src/archivist/__init__.py @@ -6,4 +6,4 @@ from ._version import __version__ from .core.archive import Archive from .core.registry import RegistryLibrarian, RegistrySqlite -from .core.storage import StorageDisk, StorageHPSS +from .storage import StorageDisk, StorageHPSS diff --git a/src/archivist/__main__.py b/src/archivist/__main__.py index 9cb86e8..4c2d264 100644 --- a/src/archivist/__main__.py +++ b/src/archivist/__main__.py @@ -4,6 +4,7 @@ from pathlib import Path import click +import loguru import requests from archivist.core.models import ManifestEntry, ManifestFailedResponse, ManifestRequest, ManifestResponse @@ -34,12 +35,6 @@ def main(ctx, config): @click.option( "--source-path", type=click.Path(exists=True), required=True, help="The source path of the file to archive" ) -@click.option( - "--dest-path", - type=click.Path(exists=False), - required=True, - help="The destination path where the file should be archived", -) @click.option( "--librarian-name", type=str, @@ -48,18 +43,17 @@ def main(ctx, config): default="librarian", ) @click.pass_context -def archive(ctx, source_path, dest_path, librarian_name): +def archive(ctx, source_path, librarian_name): """Command to archive a file. For example: `archivist -c config archive "--source-path /path/to/source --dest-path /path/to/dest"` """ # Here you would implement the logic to handle the archiving process # For now, we will just print the source and destination paths - click.echo(f"Archiving file from {source_path} to {dest_path}") + click.echo(f"Archiving file from {source_path}") settings = get_settings() source_root = Path(source_path) - dest_root = Path(dest_path) if dest_path is not None else "source_root" if source_root.is_file(): files = [source_root] @@ -74,7 +68,7 @@ def archive(ctx, source_path, dest_path, librarian_name): for file_path in files: stat = file_path.stat() relative_path = file_path.relative_to(walk_root) - + loguru.logger.info(f"Archiving file: {file_path}, relative path: {relative_path}") store_files.append( ManifestEntry( name=file_path.name, @@ -83,7 +77,7 @@ def archive(ctx, source_path, dest_path, librarian_name): checksum=_checksum(file_path), uploader=uploader, source=str(source_root), - instance_path=str(dest_root / relative_path), + instance_path=str(file_path.absolute()), instance_create_time=datetime.now(), instance_available=True, outgoing_transfer_id=0, @@ -132,11 +126,7 @@ def start_server(ctx): import uvicorn - from .database import create_all from .settings import server_settings - from .tasks.archive import start_archive - - create_all() uvicorn.run( "archivist.server:main", diff --git a/src/archivist/api/archive.py b/src/archivist/api/archive.py index f5636ff..4302c8c 100644 --- a/src/archivist/api/archive.py +++ b/src/archivist/api/archive.py @@ -42,7 +42,7 @@ def archive( manifest_id=str(manifest_id), manifest=manifest_request.model_dump_json(), paths=[entry.instance_path for entry in manifest_request.store_files], - root=settings.storage_root, + root=settings.archive_root, ) session.add(item) session.commit() diff --git a/src/archivist/api/extract.py b/src/archivist/api/extract.py index 8dfccc8..2ebe8db 100644 --- a/src/archivist/api/extract.py +++ b/src/archivist/api/extract.py @@ -4,8 +4,10 @@ from fastapi import Depends, Response from archivist.api import router -from archivist.core.models import Archive, ManifestFailedResponse, ManifestRequest, ManifestResponse -from archivist.queue import Queue, get_extract_queue +from archivist.core.archive import Archive +from archivist.core.models import ManifestFailedResponse, ManifestRequest, ManifestResponse + +# from archivist.queue import Queue, get_extract_queue from archivist.settings import Settings, get_settings @@ -13,7 +15,7 @@ def extract( manifest_request: ManifestRequest, response: Response, - queue: Queue = Depends(get_extract_queue), + # queue: Queue = Depends(get_extract_queue), settings: Settings = Depends(get_settings), ): """ diff --git a/src/archivist/core/__init__.py b/src/archivist/core/__init__.py index ee43863..e45246d 100644 --- a/src/archivist/core/__init__.py +++ b/src/archivist/core/__init__.py @@ -1,3 +1 @@ -from .storage import StorageDisk, StorageHPSS - -storage_factory = {"hpss:": StorageHPSS, "posix": StorageDisk} +from archivist.storage import StorageDisk, StorageHPSS, storage_factory # noqa: F401 diff --git a/src/archivist/core/archive.py b/src/archivist/core/archive.py index b8c9198..e0a75f6 100644 --- a/src/archivist/core/archive.py +++ b/src/archivist/core/archive.py @@ -2,7 +2,9 @@ # Full license can be found in the top level "LICENSE" file. """Classes for working with archives.""" -import subprocess +from pathlib import Path + +from .models import ManifestRequest class Archive(object): @@ -16,38 +18,60 @@ class Archive(object): Args: manifest (str): Path to a manifest file listing the files / directories contained in the archive. - paths (list): Alternatively list the full set of files / directories contained - in the archive. root (str): The common root filesystem location for building relative paths for all objects. """ - def __init__(self, manifest=None, paths=None, root=None, type=None): + def __init__( + self, + manifest: dict | None = None, + manifest_id: str | None = None, + local_root: str | None = None, + archive_root: str | None = None, + type: str | None = None, + ): self._manifest = manifest - self._paths = paths - self._root = root + self._manifest_id = manifest_id + self._local_root = local_root + self._archive_root = archive_root self._type = type self._checksum = None # Placeholder for checksum value @property - def paths(self): - return self._paths + def local_root(self): + return self._local_root @property - def root(self): - return self._root + def archive_root(self): + return self._archive_root @property def checksum(self): - """Return the checksum of the archive.""" + return self._checksum + + @property + def manifest(self): + return self._manifest + + @property + def manifest_id(self): + return self._manifest_id + + def __repr__(self): + return f", , , " def create_archive(self): """Create the archive based on the manifest or paths.""" # Placeholder for archive creation logic if self._type == "posix": - # Implement logic for creating a POSIX archive - pass + src_dst_paths = [] + for entry in self._manifest["store_files"]: + src_path = Path(entry["instance_path"]) + dst_path = Path(self._archive_root) / src_path.relative_to(self._local_root) + src_dst_paths.append((src_path, dst_path)) + return src_dst_paths + elif self._type == "hpss": # Implement logic for creating an HPSS archive # Call subprocess to create the archive using HPSS commands diff --git a/src/archivist/core/models.py b/src/archivist/core/models.py index c8c3cd9..55d6e4d 100644 --- a/src/archivist/core/models.py +++ b/src/archivist/core/models.py @@ -48,17 +48,3 @@ class ManifestResponse(BaseModel): class ManifestFailedResponse(BaseModel): error: str "The error message indicating why the manifest request failed." - - -class Archive(BaseModel): - manifest_id: str - "The ID of the manifest associated with this archive." - - manifest: str - "Path to a manifest file listing the files / directories contained in the archive." - - paths: list[str] - "Alternatively list the full set of files / directories contained in the archive." - - root: str - "The common root filesystem location for building relative paths for all objects." diff --git a/src/archivist/core/storage.py b/src/archivist/core/storage.py index c4f24e4..8d19011 100644 --- a/src/archivist/core/storage.py +++ b/src/archivist/core/storage.py @@ -2,12 +2,15 @@ # Full license can be found in the top level "LICENSE" file. """Classes for working with storage systems.""" +from archivist.core.archive import Archive +from archivist.settings import Settings + class Storage(object): """Base class representing an archive storage system.""" - def __init__(self): - pass + def __init__(self, settings: Settings | None = None): + self._settings = settings def _store(self, archive): raise NotImplementedError("Fell through to base class") @@ -62,11 +65,12 @@ class StorageDisk(Storage): """ - def __init__(self, directory): - self._directory = directory - super().__init__() + def __init__(self, settings: Settings | None = None): + super().__init__(settings=settings) + self._storage_proc = None + self._archive = None - def _store(self, archive): + def _store(self, archive: Archive): """Tar the archive files into the storage directory.""" pass diff --git a/src/archivist/database.py b/src/archivist/database.py index d56aeb4..1b4b1c6 100644 --- a/src/archivist/database.py +++ b/src/archivist/database.py @@ -2,6 +2,8 @@ Core database runner for SQLAlchemy. """ +from typing import Generator + from loguru import logger from sqlalchemy import ( BigInteger, @@ -15,43 +17,71 @@ String, create_engine, ) -from sqlalchemy.orm import declarative_base, sessionmaker +from sqlalchemy.orm import Session, declarative_base, sessionmaker -from .settings import server_settings +from .settings import get_settings -logger.info("Starting database engine.") +# Engine and sessionmaker are created lazily, on first use, rather than at +# import time. Settings (in particular the database path / ARCHIVIST_CONFIG_PATH) +# may not be finalized yet when this module is first imported -- e.g. the CLI +# sets ARCHIVIST_CONFIG_PATH from `--config` inside a click callback, which runs +# after all top-level imports have already happened. +_engine = None +_SessionMaker = None -# Create the engine and sessionmaker. -engine = create_engine( - server_settings.sqlalchemy_database_uri, - # Required for async and SQLite - connect_args=({"check_same_thread": False} if "sqlite" in server_settings.sqlalchemy_database_uri else {}), -) -logger.info("Creating database session.") +def get_engine(): + """ + Returns the (lazily created) SQLAlchemy engine, built from the current settings. + """ + + global _engine + + if _engine is None: + settings = get_settings() + logger.info("Starting database engine.") + _engine = create_engine( + settings.sqlalchemy_database_uri, + # Required for async and SQLite + connect_args=({"check_same_thread": False} if "sqlite" in settings.sqlalchemy_database_uri else {}), + ) + + return _engine + + +def get_sessionmaker() -> sessionmaker: + """ + Returns the (lazily created) sessionmaker, bound to the lazily created engine. + """ + + global _SessionMaker + + if _SessionMaker is None: + logger.info("Creating database session.") + _SessionMaker = sessionmaker(bind=get_engine(), autocommit=False, autoflush=False) -SessionMaker = sessionmaker(bind=engine, autocommit=False, autoflush=False) + return _SessionMaker -def yield_session() -> SessionMaker: +def yield_session() -> "Generator[Session, None, None]": """ Yields a new databse session. """ - session = SessionMaker() + session = get_sessionmaker()() try: yield session finally: session.close() -def get_session() -> SessionMaker: +def get_session() -> "Session": """ Returns a new database session. Unlike yield_session, it is your responsibility to close the session. """ - return SessionMaker() + return get_sessionmaker()() Base = declarative_base() @@ -65,4 +95,4 @@ def create_all() -> None: from . import orm # noqa: F401 - Base.metadata.create_all(engine) + Base.metadata.create_all(get_engine()) diff --git a/src/archivist/queue.py b/src/archivist/queue.py index d44a722..37c6daa 100644 --- a/src/archivist/queue.py +++ b/src/archivist/queue.py @@ -1,22 +1,22 @@ import queue -from archivist.core.models import Archive +from archivist.storage.base import Storage -_archive_queue: "Queue | None" = None -_extract_queue: "Queue | None" = None +_status_queue: "Queue | None" = None class Queue: """Thread-safe queue for Archive jobs awaiting processing.""" def __init__(self) -> None: - self._q: queue.Queue[Archive] = queue.Queue() + self._q: queue.Queue[Storage] = queue.Queue() - def enqueue_archive(self, archive: Archive) -> None: - self._q.put(archive) + def enqueue(self, item: Storage) -> None: + self._q.put(item) - def dequeue_archive(self, block: bool = True, timeout: float | None = None) -> Archive: - return self._q.get(block=block, timeout=timeout) + def dequeue(self, block: bool = True, timeout: float | None = None) -> Storage: + item = self._q.get(block=block, timeout=timeout) + return item def task_done(self) -> None: self._q.task_done() @@ -25,18 +25,9 @@ def task_done(self) -> None: def size(self) -> int: return self._q.qsize() - return self._q.qsize() - - -def get_archive_queue() -> Queue: - global _archive_queue - if _archive_queue is None: - _archive_queue = Queue() - return _archive_queue - -def get_extract_queue() -> Queue: - global _extract_queue - if _extract_queue is None: - _extract_queue = Queue() - return _extract_queue +def get_status_queue() -> Queue: + global _status_queue + if _status_queue is None: + _status_queue = Queue() + return _status_queue diff --git a/src/archivist/server.py b/src/archivist/server.py index 75e493b..7679cb6 100644 --- a/src/archivist/server.py +++ b/src/archivist/server.py @@ -25,6 +25,18 @@ def _archive_worker_loop(): logger.exception("Archive worker iteration failed") +def _status_worker_loop(): + from loguru import logger + + from .tasks.archive import process_status_queue + + while not _archive_stop_event.is_set(): + try: + process_status_queue() + except Exception: + logger.exception("Archive worker iteration failed") + + @asynccontextmanager async def slack_post_at_startup_shutdown(app: FastAPI): """ @@ -33,16 +45,20 @@ async def slack_post_at_startup_shutdown(app: FastAPI): """ from loguru import logger + from .database import create_all + logger.info("Archivist server starting up") - archive_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="archive-worker") - archive_pool.submit(_archive_worker_loop) + create_all() + thread_pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix="worker") + thread_pool.submit(_archive_worker_loop) + thread_pool.submit(_status_worker_loop) yield logger.info("Archivist server shutting down") _archive_stop_event.set() - archive_pool.shutdown(wait=True) + thread_pool.shutdown(wait=True) def main() -> FastAPI: diff --git a/src/archivist/settings.py b/src/archivist/settings.py index c425e60..f36743b 100644 --- a/src/archivist/settings.py +++ b/src/archivist/settings.py @@ -41,8 +41,9 @@ class Settings(BaseSettings): port: int = 8080 # storage options - storage_type: str = "posix" - storage_root: str = "/tmp/archivist_storage" + archive_type: str = "posix" + archive_root: str = "/tmp/archivist_storage" + local_root: str = "/tmp/archivist_local" # Database migration settings alembic_config_path: str = "." @@ -76,6 +77,9 @@ class Settings(BaseSettings): checksum_threads: int = 4 checksum_timeout: datetime.timedelta = datetime.timedelta(days=1) + # Number of background threads used to perform filesystem storage copies. + storage_threads: int = 4 + def model_post_init(__context, *args, **kwargs): """ Read sensitive data from their appropriate files. diff --git a/src/archivist/storage/__init__.py b/src/archivist/storage/__init__.py new file mode 100644 index 0000000..54c15e6 --- /dev/null +++ b/src/archivist/storage/__init__.py @@ -0,0 +1,4 @@ +from .storage_disk import StorageDisk +from .storage_hpss import StorageHPSS + +storage_factory = {"hpss:": StorageHPSS, "posix": StorageDisk} diff --git a/src/archivist/storage/base.py b/src/archivist/storage/base.py new file mode 100644 index 0000000..310b896 --- /dev/null +++ b/src/archivist/storage/base.py @@ -0,0 +1,49 @@ +from archivist.core.archive import Archive +from archivist.settings import Settings + + +class Storage(object): + """Base class representing an archive storage system.""" + + def __init__(self, settings: Settings | None = None): + self._settings = settings + + def _store(self, archive): + raise NotImplementedError("Fell through to base class") + + def store(self, archive): + """Store an archive. + + Args: + archive (Archive): The archive to store. + + Returns: + (dict): Metadata specific to the storage system. Used when registering + the archive. + + """ + return self._store(archive) + + def _extract(self, archive_name, storage_info, outdir, paths=None): + raise NotImplementedError("Fell through to base class") + + def extract(self, archive_name, storage_info, outdir, paths=None): + """Extract objects from an archive. + + The archive name and additional storage metadata is used to locate and extract + one or more paths from the archive. `storage_info` should be the same + metadata returned by the call to `store()` and kept in the registry. + + Args: + archive_name (str): The name of the archive to remove. + storage_info (dict): Additional storage-system specific metadata. + outdir (str): The top-level output directory where extracted files + should be placed. + paths (list): Optional list of objects to extract. If None, the + entire archive is extracted. + + Returns: + (None) + + """ + self._extract(archive_name, storage_info, outdir, paths=None) diff --git a/src/archivist/storage/storage_disk.py b/src/archivist/storage/storage_disk.py new file mode 100644 index 0000000..eba398c --- /dev/null +++ b/src/archivist/storage/storage_disk.py @@ -0,0 +1,63 @@ +import os +import shutil +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path + +from loguru import logger + +from archivist.core.archive import Archive +from archivist.database import get_session +from archivist.orm.archivequeue import ArchiveQueue +from archivist.settings import Settings +from archivist.storage.base import Storage + + +class StorageDisk(Storage): + """Simple filesystem storage. + + This class just keeps archives in subdirectories within a top-level location + (``settings.archive_root``). Useful for testing and for archiving to "slow" + storage that provides a POSIX API. + + Copies run in a shared background thread pool (sized by + ``settings.storage_threads``) so that ``_store`` returns immediately; callers + check ``future`` to see when the copy has finished. + """ + + _executor: ThreadPoolExecutor | None = None + + def __init__(self, settings: Settings | None = None): + super().__init__(settings=settings) + self._archive = None + self._future: Future | None = None + + if StorageDisk._executor is None: + max_workers = self._settings.storage_threads if self._settings else 4 + StorageDisk._executor = ThreadPoolExecutor(max_workers=max_workers) + + @property + def future(self) -> Future | None: + """The Future tracking the in-progress (or completed) copy, if any.""" + return self._future + + def _copy_files(self, file_list: list[tuple[Path, Path]]) -> None: + """Copy each path in file_list into dest_dir.""" + for src_path, dst_path in file_list: + os.makedirs(dst_path.parent, exist_ok=True) + if os.path.isdir(src_path): + shutil.copytree(src_path, dst_path, dirs_exist_ok=True) + else: + shutil.copy2(src_path, dst_path) + + def _store(self, archive: Archive) -> "StorageDisk": + """Copy the archive's files into a per-manifest subdirectory of archive_root.""" + self._archive = archive + paths = self._archive.create_archive() + + self._future = StorageDisk._executor.submit(self._copy_files, paths) + + return self + + def _extract(self, archive_name, storage_info, outdir, paths=None): + """Extract files from tar archives.""" + pass diff --git a/src/archivist/storage/storage_hpss.py b/src/archivist/storage/storage_hpss.py new file mode 100644 index 0000000..8989fa4 --- /dev/null +++ b/src/archivist/storage/storage_hpss.py @@ -0,0 +1,20 @@ +from archivist.core.archive import Archive +from archivist.settings import Settings +from archivist.storage.base import Storage + + +class StorageHPSS(Storage): + """Store archives to HPSS tapes.""" + + def __init__(self): + super().__init__() + + def _store(self, archive): + """Use HTAR to archive files. + The tape list is returned for including in the registry entry. + """ + pass + + def _extract(self, archive_name, storage_info, outdir, paths=None): + """Extract files from HTAR archives.""" + pass diff --git a/src/archivist/tasks/archive.py b/src/archivist/tasks/archive.py index c1f3cea..d943a39 100644 --- a/src/archivist/tasks/archive.py +++ b/src/archivist/tasks/archive.py @@ -1,25 +1,85 @@ +import json +import queue from time import sleep import loguru from sqlalchemy import inspect +from archivist.core import storage_factory +from archivist.core.archive import Archive from archivist.database import get_session from archivist.orm.archivequeue import ArchiveQueue +from archivist.queue import get_status_queue from archivist.settings import get_settings -def start_archive(librarian_name: str, store_files: bool = True): +def start_archive(librarian_name: str, store_files: bool = True) -> None: """ Start an archive process by sending a request to the archivist server. + + Called directly from the background worker thread loop (not a FastAPI + route), so dependencies are resolved here rather than via `Depends`. """ + status_queue = get_status_queue() settings = get_settings() - session = get_session() try: archive_item = ArchiveQueue.dequeue(session) if archive_item is not None: archive_dict = {c.key: getattr(archive_item, c.key) for c in inspect(archive_item).mapper.column_attrs} - loguru.logger.info(f"Archive Item: {archive_dict}") + loguru.logger.info(f"Archive Item: {archive_dict}, {type(archive_item)}") + archive = Archive( + manifest=json.loads(archive_dict["manifest"]), + manifest_id=archive_dict["manifest_id"], + local_root=settings.local_root, + archive_root=settings.archive_root, + type=settings.archive_type, + ) + storage = storage_factory[settings.archive_type](settings=settings) + storage_task = storage.store(archive) + status_queue.enqueue(storage_task) finally: session.close() - sleep(1) + + +def process_status_queue() -> None: + """ + Process the status queue to check for completed or failed archive tasks. + + Called directly from the background worker thread loop (not a FastAPI + route), so dependencies are resolved here rather than via `Depends`. + """ + status_queue = get_status_queue() + try: + storage_task = status_queue.dequeue(block=False) + loguru.logger.debug(f"Processing status queue for storage task: {storage_task._archive}") + if storage_task.future.done(): + loguru.logger.debug(f"Storage task for archive {storage_task._archive.manifest_id} future completed.") + session = get_session() + try: + item = session.query(ArchiveQueue).filter_by(manifest_id=storage_task._archive.manifest_id).first() + loguru.logger.debug( + f"Retrieved ArchiveQueue item for manifest_id {storage_task._archive.manifest_id}: {item}" + ) + if item is None: + return + try: + item.complete(session) + loguru.logger.info(f"Archive {storage_task._archive.manifest_id} completed successfully.") + except Exception: + loguru.logger.exception(f"Archive {storage_task._archive.manifest_id} failed to store.") + item.fail(session) + finally: + session.close() + else: + loguru.logger.debug( + f"Storage task for archive {storage_task._archive.manifest_id} future not completed yet." + ) + status_queue.enqueue(storage_task) + + status_queue.task_done() + except queue.Empty: + pass + except Exception as e: + loguru.logger.error(f"Error processing status queue: {e}") + sleep(1) # Sleep for a short duration to avoid busy waiting From 2f2a6e1d36e3886742a4bd3057044206827b7ba0 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Thu, 9 Jul 2026 10:54:13 -0400 Subject: [PATCH 3/3] fix: fixed return --- src/archivist/core/archive.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/archivist/core/archive.py b/src/archivist/core/archive.py index e0a75f6..0d5d5b0 100644 --- a/src/archivist/core/archive.py +++ b/src/archivist/core/archive.py @@ -70,7 +70,7 @@ def create_archive(self): src_path = Path(entry["instance_path"]) dst_path = Path(self._archive_root) / src_path.relative_to(self._local_root) src_dst_paths.append((src_path, dst_path)) - return src_dst_paths + return src_dst_paths elif self._type == "hpss": # Implement logic for creating an HPSS archive