From 519639714dc54e926598a129de8dc53e1701bb00 Mon Sep 17 00:00:00 2001 From: Josh Borrow Date: Mon, 5 Feb 2024 14:53:35 -0500 Subject: [PATCH] Add background task settings and minor db stuff --- .../versions/71df5b41ae41_initial_schema.py | 20 +- librarian_background/__init__.py | 20 +- librarian_background/bad.py | 2 +- librarian_background/poll.py | 18 -- librarian_background/settings.py | 201 ++++++++++++++++++ pyproject.toml | 4 + .../test_background_settings.py | 33 +++ 7 files changed, 259 insertions(+), 39 deletions(-) delete mode 100644 librarian_background/poll.py create mode 100644 librarian_background/settings.py create mode 100644 tests/background_unit_test/test_background_settings.py diff --git a/alembic/versions/71df5b41ae41_initial_schema.py b/alembic/versions/71df5b41ae41_initial_schema.py index a87cd6d..30087cf 100644 --- a/alembic/versions/71df5b41ae41_initial_schema.py +++ b/alembic/versions/71df5b41ae41_initial_schema.py @@ -151,16 +151,6 @@ def upgrade(): Column("destination_instance_id", Integer, ForeignKey("instances.id")), ) - op.create_table( - "remote_instances", - Column("id", Integer(), primary_key=True, autoincrement=True, unique=True), - Column("file_name", String(256), ForeignKey("files.name"), nullable=False), - Column("store_id", Integer(), nullable=False), - Column("librarian_id", Integer(), ForeignKey("librarians.id"), nullable=False), - Column("copy_time", DateTime(), nullable=False), - Column("sender", String(256), nullable=False), - ) - op.create_table( "librarians", Column("id", Integer(), primary_key=True, autoincrement=True), @@ -173,6 +163,16 @@ def upgrade(): Column("last_heard", DateTime(), nullable=False), ) + op.create_table( + "remote_instances", + Column("id", Integer(), primary_key=True, autoincrement=True, unique=True), + Column("file_name", String(256), ForeignKey("files.name"), nullable=False), + Column("store_id", Integer(), nullable=False), + Column("librarian_id", Integer(), ForeignKey("librarians.id"), nullable=False), + Column("copy_time", DateTime(), nullable=False), + Column("sender", String(256), nullable=False), + ) + op.create_table( "errors", Column("id", Integer(), primary_key=True, autoincrement=True, unique=True), diff --git a/librarian_background/__init__.py b/librarian_background/__init__.py index 2379abb..5e2299c 100644 --- a/librarian_background/__init__.py +++ b/librarian_background/__init__.py @@ -13,23 +13,23 @@ from .check_integrity import CheckIntegrity from .core import SafeScheduler from .create_clone import CreateLocalClone +from .settings import background_settings def background(run_once: bool = False): scheduler = SafeScheduler() # Set scheduling... - scheduler.every(12).hours.do( - CheckIntegrity(name="check_integrity", store_name="local_store", age_in_days=7) - ) - scheduler.every(12).hours.do( - CreateLocalClone( - name="create_clone", - clone_from="local_store", - clone_to="local_clone", - age_in_days=7, - ) + + all_tasks = ( + background_settings.check_integrity + + background_settings.create_local_clone + + background_settings.send_clone + + background_settings.recieve_clone ) + for task in all_tasks: + scheduler.every(task.every.seconds).seconds.do(task.task) + # ...and run it all on startup. scheduler.run_all() diff --git a/librarian_background/bad.py b/librarian_background/bad.py index ca8742b..da62ae0 100644 --- a/librarian_background/bad.py +++ b/librarian_background/bad.py @@ -5,7 +5,7 @@ from .task import Task -class Bad(Task): +class Bad(Task): # pragma: no cover """ A simple background task that polls for new files. """ diff --git a/librarian_background/poll.py b/librarian_background/poll.py deleted file mode 100644 index 6e7e385..0000000 --- a/librarian_background/poll.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -A simple 'example' background task that does nothing but print -to the screen. -""" - -from .task import Task - - -class Poll(Task): - """ - A simple background task that polls for new files. - """ - - name: str = "poll" - - def on_call(self): - print("Polling for new files...") - return diff --git a/librarian_background/settings.py b/librarian_background/settings.py new file mode 100644 index 0000000..2f3e698 --- /dev/null +++ b/librarian_background/settings.py @@ -0,0 +1,201 @@ +""" +Background task settings. +""" + +import abc +import datetime +import os +from pathlib import Path +from typing import TYPE_CHECKING, Any, Optional + +from pydantic import BaseModel +from pydantic_settings import BaseSettings, SettingsConfigDict + +from hera_librarian.deletion import DeletionPolicy + +from .check_integrity import CheckIntegrity +from .create_clone import CreateLocalClone +from .recieve_clone import RecieveClone +from .send_clone import SendClone + +if TYPE_CHECKING: + from .task import Task + + background_settings: "BackgroundSettings" + + +class BackgroundTaskSettings(BaseModel, abc.ABC): + """ + Settings for an individual background task. Generic, should be inherited from for + specific tasks. + """ + + task_name: str + "The name of the task. Used for logging purposes." + + every: datetime.timedelta + "How often to run the task. You can pass in any ``datetime.timedelta`` string, e.g. HH:MM:SS (note leading zeroes are required)." + + @abc.abstractproperty + def task(self) -> "Task": # pragma: no cover + raise NotImplementedError + + +class CheckIntegritySettings(BackgroundTaskSettings): + """ + Settings for the integrity check task. + """ + + age_in_days: int + "The age of the items to check, in days." + + store_name: str + "The name of the store to check." + + @property + def task(self) -> CheckIntegrity: + return CheckIntegrity( + name=self.task_name, + store_name=self.store_name, + age_in_days=self.age_in_days, + ) + + +class CreateLocalCloneSettings(BackgroundTaskSettings): + """ + Settings for the local clone creation task. + """ + + age_in_days: int + "The age of the items to check, in days." + + clone_from: str + "The name of the store to clone from." + + clone_to: str + "The name of the store to clone to." + + @property + def task(self) -> CreateLocalClone: + return CreateLocalClone( + name=self.task_name, + clone_from=self.clone_from, + clone_to=self.clone_to, + age_in_days=self.age_in_days, + ) + + +class SendCloneSettings(BackgroundTaskSettings): + """ + Settings for the clone sending task. + """ + + destination_librarian: str + "The destination librarian for this clone." + + age_in_days: int + "The age of the items to check, in days." + + store_preference: Optional[str] + "The store to send. If None, send all stores." + + @property + def task(self) -> SendClone: + return SendClone( + name=self.task_name, + destination_librarian=self.destination_librarian, + age_in_days=self.age_in_days, + store_preference=self.store_preference, + ) + + +class RecieveCloneSettings(BackgroundTaskSettings): + """ + Settings for the clone receiving task. + """ + + deletion_policy: DeletionPolicy + "The deletion policy for the incoming files." + + @property + def task(self) -> RecieveClone: + return RecieveClone( + name=self.task_name, + deletion_policy=self.deletion_policy, + ) + + +class BackgroundSettings(BaseSettings): + """ + Background task settings, configurable. + """ + + check_integrity: list[CheckIntegritySettings] = [] + "Settings for the integrity check task." + + create_local_clone: list[CreateLocalCloneSettings] = [] + "Settings for the local clone creation task." + + send_clone: list[SendCloneSettings] = [] + "Settings for the clone sending task." + + recieve_clone: list[RecieveCloneSettings] = [] + "Settings for the clone receiving task." + + model_config = SettingsConfigDict(env_prefix="librarian_background_") + + @classmethod + def from_file(cls, config_path: Path | str) -> "BackgroundSettings": + """ + Loads the settings from the given path. + """ + + with open(config_path, "r") as handle: + return cls.model_validate_json(handle.read()) + + +# Automatically create a settings object on use. + +_settings = None + + +def load_settings() -> BackgroundSettings: + """ + Load the settings from the config file. + """ + + global _settings + + try_paths = [ + os.environ.get("LIBRARIAN_BACKGROUND_CONFIG", None), + ] + + for path in try_paths: + if path is not None: + path = Path(path) + else: + continue + + if path.exists(): + _settings = BackgroundSettings.from_file(path) + return _settings + + _settings = BackgroundSettings() + + return _settings + + +def __getattr__(name): + """ + Try to load the settings if they haven't been loaded yet. + """ + + if name == "background_settings": + global _settings + + if _settings is not None: + return _settings + + return load_settings() + + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/pyproject.toml b/pyproject.toml index bb4a4a2..78e81d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,9 @@ dev = [ "pytest-cov", "pytest-xprocess", ] +postgres = [ + "psycopg[binary,pool]", +] [project.scripts] librarian-server-start = "librarian_server_scripts.librarian_server_start:main" @@ -84,6 +87,7 @@ source = [ [tool.coverage.report] exclude_lines = ["pragma: no cover"] +exclude_also = ["if TYPE_CHECKING:"] [tool.isort] profile = "black" diff --git a/tests/background_unit_test/test_background_settings.py b/tests/background_unit_test/test_background_settings.py new file mode 100644 index 0000000..ddcd05a --- /dev/null +++ b/tests/background_unit_test/test_background_settings.py @@ -0,0 +1,33 @@ +""" +Tests our ability to serialize/deserialize the background settings. +""" + +from librarian_background.settings import BackgroundSettings + + +def test_background_settings_full(): + BackgroundSettings.model_validate( + { + "check_integrity": [ + { + "task_name": "check", + "every": "01:00:00", + "age_in_days": 7, + "store_name": "test", + } + ], + "create_local_clone": [ + { + "task_name": "clone", + "every": "22:23:02", + "age_in_days": 7, + "clone_from": "test", + "clone_to": "test", + } + ], + } + ) + + +def test_background_settings_empty(): + BackgroundSettings()