Add ability to teardown a store during manifest creation

hera_librarian/models/admin.py

@@ -114,6 +114,9 @@ class AdminStoreManifestRequest(BaseModel):
     disable_store: bool = False
     "Whether to disable the store after creating the outgoing transfers."
 
+    mark_local_instances_as_unavailable: bool = False
+    "Mark the local instances as unavailable after creating the outgoing transfers."
+
 
 class AdminStoreManifestResponse(BaseModel):
     librarian_name: str

librarian_server/api/admin.py

@@ -28,6 +28,7 @@
 from ..database import yield_session
 from ..logger import log
 from ..orm import File, Instance, Librarian, OutgoingTransfer, StoreMetadata
+from ..settings import server_settings
 from ..stores import InvertedStoreNames, StoreNames
 from .auth import AdminUserDependency
 
@@ -192,6 +193,31 @@ def store_manifest(
     be a very large request and response, so use with caution. You can then
     ingest the manifest items individually into a different instance of
     the librarian, thus completing the 'sneakernet' process.
+
+    This is a very powerful endpoint, and has the following options
+    configured with its request:
+
+    - create_outgoing_transfers: If true, will create outgoing transfers
+        for each file in the manifest. This is useful for sneakernetting
+        files to a different librarian.
+
+    - destination_librarian: The name of the librarian to send the files to,
+        if create_outgoing_transfers is true. This is required if you are
+        creating outgoing transfers.
+
+    - disable_store: If true, will disable the store after creating the
+        outgoing transfers. This is useful for sneakernetting files to a
+        different librarian, as it allows you to (in one transaction)
+        generate all the outgoing transfers, then disable the store.
+
+    - mark_local_instances_as_unavailable: If true, will mark the local
+        instances as unavailable after creating the outgoing transfers.
+
+    An easy sneakernet workflow is to set all of these to true. This will
+    generate a complete manifest, create outgoing transfers for each file,
+    then disable the store and mark the local instances as unavailable
+    (as we are assuming you are going to then remove the files from the
+    disks entirely).
     """
 
     log.debug(f"Recieved manifest request from {user.username}: {request}.")
@@ -211,6 +237,14 @@ def store_manifest(
             suggested_remedy="Create the store first. Maybe you need to run DB migration?",
         )
 
+    # Now, stop anyone from ingesting any new files if we want the store
+    # to be disabled at the end of this process.
+
+    if request.disable_store:
+        store.enabled = False
+        session.commit()
+        log.info(f"Disabled store {store.name}.")
+
     # If we are going to create outgoing transfers, we need to make sure
     # that the destination librarian exists.
 
@@ -256,7 +290,7 @@ def create_manifest_entry(instance: Instance) -> ManifestEntry:
         else:
             transfer_id = -1
 
-        return ManifestEntry(
+        entry = ManifestEntry(
             name=instance.file.name,
             create_time=instance.file.create_time,
             size=instance.file.size,
@@ -270,7 +304,14 @@ def create_manifest_entry(instance: Instance) -> ManifestEntry:
             outgoing_transfer_id=transfer_id,
         )
 
+        if request.mark_local_instances_as_unavailable:
+            instance.available = False
+            session.commit()
+
+        return entry
+
     response = AdminStoreManifestResponse(
+        librarian_name=server_settings.name,
         store_name=store.name,
         store_files=[
             create_manifest_entry(instance)
@@ -279,39 +320,8 @@ def create_manifest_entry(instance: Instance) -> ManifestEntry:
         ],
     )
 
-    if request.disable_store:
-        store.enabled = False
-
-        try:
-            session.commit()
-        except Exception as e:
-            log.error(
-                f"Failed to disable store {store.name}: {e}. Returning 500, but before "
-                "that, we need to kill off all these transfers we just created."
-            )
-
-            # Kill off all the transfers we just created.
-            if request.create_outgoing_transfers:
-                successfully_killed = 0
-                for file in response.store_files:
-                    if file.outgoing_transfer_id != -1:
-                        transfer = session.query(OutgoingTransfer).get(
-                            file.outgoing_transfer_id
-                        )
-                        transfer.fail_transfer(session)
-                        successfully_killed += 1
-
-                log.error(
-                    f"Killed off {successfully_killed}/{len(response.store_files)} transfers "
-                    "we just created."
-                )
-
-            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
-            return AdminRequestFailedResponse(
-                reason="Failed to disable store.",
-                suggested_remedy="Check the logs for more information.",
-            )
-
-        log.info(f"Disabled store {store.name}.")
+    log.info(
+        f"Generated manifest for store {store.name} containing {len(response.store_files)} files."
+    )
 
     return response

librarian_server/orm/librarian.py

@@ -42,7 +42,9 @@ class Librarian(db.Base):
     "The last time we heard from this librarian (the last time it connected to us)."
 
     @classmethod
-    def new_librarian(self, name: str, url: str, port: int) -> "Librarian":
+    def new_librarian(
+        self, name: str, url: str, port: int, check_connection: bool = True
+    ) -> "Librarian":
         """
         Create a new librarian object.
 
@@ -54,6 +56,9 @@ def new_librarian(self, name: str, url: str, port: int) -> "Librarian":
             The URL of this librarian.
         port : int
             The port of this librarian.
+        check_connection : bool
+            Whether to check the connection to this librarian before
+            returning it (default: True, but turn this off for tests.)
 
         Returns
         -------
@@ -69,6 +74,9 @@ def new_librarian(self, name: str, url: str, port: int) -> "Librarian":
             last_heard=datetime.utcnow(),
         )
 
+        if not check_connection:
+            return librarian
+
         # Before returning it, we should ping it to confirm it exists.
 
         client = librarian.client()

librarian_server/settings.py

@@ -53,6 +53,10 @@ class ServerSettings(BaseSettings):
     variables.
     """
 
+    # Top level name of the server. Should be unique.
+    name: str = "librarian_server"
+
+    # Database settings.
     database_driver: str = "sqlite"
     database_user: Optional[str] = None
     database_password: Optional[str] = None
@@ -61,6 +65,8 @@ class ServerSettings(BaseSettings):
     database: Optional[str] = None
 
     log_level: str = "DEBUG"
+
+    # Display name and description of the site, used in UI only.
     displayed_site_name: str = "Untitled Librarian"
     displayed_site_description: str = "No description set."
 

tests/server_unit_test/test_admin.py

@@ -11,6 +11,7 @@
     AdminRequestFailedResponse,
     AdminStoreListResponse,
     AdminStoreManifestRequest,
+    AdminStoreManifestResponse,
     AdminStoreStateChangeRequest,
     AdminStoreStateChangeResponse,
 )
@@ -223,3 +224,82 @@ def test_store_state_change_no_store(test_client):
 
     assert response.status_code == 400
     response = AdminRequestFailedResponse.model_validate_json(response.content)
+
+
+def test_manifest_generation_and_extra_opts(
+    test_client,
+    test_server_with_many_files_and_errors,
+    test_orm,
+):
+    """
+    Tests that we can generate a manifest and that we can use extra options.
+    """
+
+    get_session = test_server_with_many_files_and_errors[1]
+
+    # First, search the stores.
+    response = test_client.post_with_auth("/api/v2/admin/stores/list", content="")
+    response = AdminStoreListResponse.model_validate_json(response.content).root
+
+    # Add in a librarian
+    with get_session() as session:
+        librarian = test_orm.Librarian.new_librarian(
+            "our_closest_friend",
+            "http://localhost",
+            80,
+            check_connection=False,
+        )
+
+        librarian.authenticator = "password"
+
+        session.add(librarian)
+        session.commit()
+
+    # Now we can try the manifest!
+
+    new_response = test_client.post_with_auth(
+        "/api/v2/admin/stores/manifest",
+        content=AdminStoreManifestRequest(
+            store_name=response[0].name,
+            create_outgoing_transfers=True,
+            destination_librarian="our_closest_friend",
+            disable_store=True,
+            mark_local_instances_as_unavailable=True,
+        ).model_dump_json(),
+    )
+
+    assert new_response.status_code == 200
+
+    new_response = AdminStoreManifestResponse.model_validate_json(new_response.content)
+
+    assert new_response.store_name == response[0].name
+    assert new_response.librarian_name == "librarian_server"
+
+    session = get_session()
+
+    for entry in new_response.store_files:
+        assert entry.outgoing_transfer_id >= 0
+
+        instance = (
+            session.query(test_orm.Instance).filter_by(path=entry.instance_path).first()
+        )
+
+        assert instance.available == False
+
+        transfer = session.get(test_orm.OutgoingTransfer, entry.outgoing_transfer_id)
+
+        assert transfer is not None
+        assert transfer.destination == "our_closest_friend"
+
+        session.delete(transfer)
+
+    store = (
+        session.query(test_orm.StoreMetadata)
+        .filter_by(name=response[0].name)
+        .one_or_none()
+    )
+
+    assert store.enabled == False
+    store.enabled = True
+
+    session.commit()