wip:general librarian disable

Author: iparask

hera_librarian/client.py

@@ -29,11 +29,13 @@
 from .models.admin import (
     AdminAddLibrarianRequest,
     AdminAddLibrarianResponse,
+    AdminChangeLibrarianStatusRequest,
     AdminChangeLibrarianTransferStatusRequest,
     AdminCreateFileRequest,
     AdminCreateFileResponse,
     AdminDeleteInstanceRequest,
     AdminDeleteInstanceResponse,
+    AdminLibrarianStatusResponse,
     AdminLibrarianTransferStatusResponse,
     AdminListLibrariansRequest,
     AdminListLibrariansResponse,
@@ -1322,7 +1324,7 @@ def remove_librarian(
 
         return response.success, response.number_of_transfers_removed
 
-    def set_librarian_status(
+    def set_librarian_transfer_status(
         self,
         librarian_name: str,
         transfers_enabled: bool,
@@ -1359,3 +1361,41 @@ def set_librarian_status(
                 raise e
 
         return response.transfers_enabled
+
+    def set_librarian_status(
+        self,
+        librarian_name: str,
+        enabled: bool,
+    ) -> bool:
+        """
+        Set the status of transfers to the librarian.
+
+        Parameters
+        ----------
+        librarian_name : str
+            The name of the librarian to set the status of.
+        enabled : bool
+            Whether the librarian should be enabled.
+
+        Returns
+        -------
+        bool
+            The new status.
+        """
+
+        try:
+            response = self.post(
+                endpoint="admin/librarians/status/set",
+                request=AdminChangeLibrarianStatusRequest(
+                    librarian_name=librarian_name,
+                    enabled=enabled,
+                ),
+                response=AdminLibrarianStatusResponse,
+            )
+        except LibrarianHTTPError as e:
+            if e.status_code == 400 and "Librarian" in e.reason:
+                raise LibrarianError(e.reason)
+            else:
+                raise e
+
+        return response.enabled

hera_librarian/models/admin.py

@@ -167,8 +167,10 @@ class LibrarianListResponseItem(BaseModel):
     available: bool | None
     "Whether the librarian is available or not, only if ping is true."
 
-    enabled: bool
+    transfers_enabled: bool
     "Whether transfers the librarian is enabled or not."
+    enabled: bool
+    "Whether the librarian is enabled or not."
 
 
 class AdminListLibrariansResponse(BaseModel):
@@ -256,6 +258,19 @@ class AdminChangeLibrarianTransferStatusRequest(BaseModel):
     transfers_enabled: bool
 
 
+class AdminChangeLibrarianStatusRequest(BaseModel):
+    """
+    A request to change the transfer status of a librarian, either
+    to enable or disable outbound transfers.
+    """
+
+    "The name of the librarian to change the transfer status of."
+    librarian_name: str
+
+    "Whether to enable or disable outbound communication."
+    enabled: bool
+
+
 class AdminLibrarianTransferStatusResponse(BaseModel):
     """
     A response to a user change request.
@@ -266,3 +281,15 @@ class AdminLibrarianTransferStatusResponse(BaseModel):
 
     "Whether the librarian has outbound transfers enabled."
     transfers_enabled: bool
+
+
+class AdminLibrarianStatusResponse(BaseModel):
+    """
+    A response to a user change request.
+    """
+
+    "The name of the librarian that was changed."
+    librarian_name: str
+
+    "Whether the librarian has communication enabled."
+    enabled: bool

librarian_background/queues.py

@@ -153,7 +153,11 @@ def check_on_consumed(
                 .one_or_none()
             )
 
-            if librarian is None or not librarian.transfers_enabled:
+            if (
+                librarian is None
+                or not librarian.transfers_enabled
+                or not librarian.enabled
+            ):
                 # We can't do anything with this librarian, but there may be other
                 # librarians that are enabled.
                 continue
@@ -258,7 +262,11 @@ def consume_queue_item(session_maker: Callable[[], "Session"]) -> bool:
             .one_or_none()
         )
 
-        if librarian is None or not librarian.transfers_enabled:
+        if (
+            librarian is None
+            or not librarian.transfers_enabled
+            or not librarian.enabled
+        ):
             # We can't do anything with this librarian, but there may be other
             # librarians that are enabled.
             return True

librarian_background/send_clone.py

@@ -580,7 +580,7 @@ def core(self, session: Session):
             return CancelJob
 
         current_time = datetime.datetime.now(datetime.timezone.utc)
-        if not librarian.transfers_enabled and (
+        if (not librarian.transfers_enabled and not librarian.enabled) and (
             current_time - self._skip_timer
         ) > datetime.timedelta(days=self.warn_disabled_timer):
             self._skip_timer = current_time

librarian_server/api/admin.py

@@ -16,11 +16,13 @@
 from hera_librarian.models.admin import (
     AdminAddLibrarianRequest,
     AdminAddLibrarianResponse,
+    AdminChangeLibrarianStatusRequest,
     AdminChangeLibrarianTransferStatusRequest,
     AdminCreateFileRequest,
     AdminCreateFileResponse,
     AdminDeleteInstanceRequest,
     AdminDeleteInstanceResponse,
+    AdminLibrarianStatusResponse,
     AdminLibrarianTransferStatusResponse,
     AdminListLibrariansRequest,
     AdminListLibrariansResponse,
@@ -390,7 +392,8 @@ def list_librarians(
                 url=librarian.url,
                 port=librarian.port,
                 available=ping if request.ping else None,
-                enabled=librarian.transfers_enabled,
+                transfer_enabled=librarian.transfers_enabled,
+                enabled=librarian.enabled,
             )
         )
 
@@ -620,3 +623,39 @@ def change_librarian_transfer_status(
         librarian_name=librarian.name,
         transfers_enabled=librarian.transfers_enabled,
     )
+
+
+@router.post(
+    path="/librarians/status/change",
+    response_model=AdminLibrarianStatusResponse,
+)
+def change_librarian_status(
+    request: AdminChangeLibrarianStatusRequest,
+    user: AdminUserDependency,
+    response: Response,
+    session: Session = Depends(yield_session),
+) -> AdminLibrarianTransferStatusResponse:
+    """
+    Change the transfer status of a librarian. This will enable or disable
+    outbound transfers to the librarian.
+    """
+
+    librarian = (
+        session.query(Librarian).filter_by(name=request.librarian_name).one_or_none()
+    )
+
+    if librarian is None:
+        response.status_code = status.HTTP_400_BAD_REQUEST
+        return AdminRequestFailedResponse(
+            reason=f"Librarian {request.librarian_name} does not exist",
+            suggested_remedy="Please verify that the requested librarian exists",
+        )
+
+    librarian.enabled = request.enabled
+
+    session.commit()
+
+    return AdminLibrarianStatusResponse(
+        librarian_name=librarian.name,
+        enabled=librarian.enabled,
+    )

librarian_server/api/corrupt.py

@@ -135,6 +135,7 @@ def user_and_librarian_validation_flow(
         background_settings.consume_queue
         and background_settings.check_consumed_queue
         and librarian.transfers_enabled
+        and librarian.enabled
         and login_success
     ):
         logger.warning(
@@ -143,7 +144,7 @@ def user_and_librarian_validation_flow(
             librarian.name,
             bool(background_settings.consume_queue),
             bool(background_settings.check_consumed_queue),
-            librarian.transfers_enabled,
+            librarian.transfers_enabled and librarian.enabled,
             login_success,
         )
         raise HTTPException(

librarian_server/api/validate.py

@@ -5,17 +5,13 @@
 """
 
 import asyncio
-import datetime
-from functools import lru_cache
-from pathlib import Path
 from time import perf_counter
 
 from asyncer import asyncify
 from fastapi import APIRouter, Depends, Response, status
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
-from hera_librarian.errors import ErrorCategory, ErrorSeverity
 from hera_librarian.exceptions import (
     LibrarianError,
     LibrarianHTTPError,
@@ -84,6 +80,9 @@ def calculate_checksum_of_remote_copies(
     librarian,
     file_name,
 ):
+    if not librarian.enabled:
+        return []
+
     start = perf_counter()
     try:
         client = librarian.client()

librarian_server/orm/librarian.py

@@ -37,7 +37,9 @@ class Librarian(db.Base):
     authenticator = db.Column(db.String(256), nullable=False)
     "The authenticator so we can connect this librarian. This is encrypted."
     transfers_enabled = db.Column(db.Boolean, nullable=False, default=True)
-
+    "Whether transfers are enabled to this librarian."
+    enabled = db.Column(db.Boolean, nullable=False, default=True)
+    "Whether this librarian is enabled."
     last_seen = db.Column(db.DateTime, nullable=False)
     "The last time we connected to and verified this librarian exists."
     last_heard = db.Column(db.DateTime, nullable=False)

librarian_server_scripts/librarian_change_status.py

@@ -0,0 +1,54 @@
+"""
+A command-line script, for use in containers, to change the behaviour
+of transfers to external librarians.
+"""
+
+import argparse as ap
+
+parser = ap.ArgumentParser(description=("Change the status of an external librarian."))
+
+parser.add_argument(
+    "--librarian",
+    help="Name of the librarian to change the status of.",
+    type=str,
+    required=True,
+)
+
+parser.add_argument(
+    "--enable",
+    help="Enable the librarian.",
+    action="store_true",
+)
+
+parser.add_argument(
+    "--disable",
+    help="Disable the librarian.",
+    action="store_true",
+)
+
+
+def main():
+    args = parser.parse_args()
+
+    if args.enable and args.disable:
+        raise ValueError("Cannot enable and disable at the same time.")
+
+    if not args.enable and not args.disable:
+        raise ValueError("Must enable or disable.")
+
+    from librarian_server.database import get_session
+    from librarian_server.orm import Librarian
+
+    with get_session() as session:
+        librarian = (
+            session.query(Librarian).filter_by(name=args.librarian).one_or_none()
+        )
+        if librarian is None:
+            raise ValueError(f"Librarian {args.librarian} does not exist.")
+
+        if args.enable:
+            librarian.enabled = True
+        elif args.disable:
+            librarian.enabled = False
+
+        session.commit()