diff --git a/hera_librarian/cli.py b/hera_librarian/cli.py index d9ae81c..f04d549 100644 --- a/hera_librarian/cli.py +++ b/hera_librarian/cli.py @@ -777,7 +777,7 @@ def set_librarian_transfer(args): Set the transfer status of a librarian. """ - client = get_client(args.conn_name, admin=True) + client = get_client(args.conn_name, admin=args.admin) if args.enabled and args.disabled: die("Cannot set both enabled and disabled.") @@ -790,7 +790,7 @@ def set_librarian_transfer(args): die("You must choose to enable or disable the transfers.") try: - client.set_librarian_transfer( + client.set_librarian_status( librarian_name=args.librarian_name, transfer_status=transfer_status, ) @@ -1760,6 +1760,11 @@ def config_set_librarian_transfer_subparser(sub_parsers): action="store_true", help="Set the librarian to disabled for transfers.", ) + sp.add_argument( + "--admin", + action="store_true", + help="Set the librarian transfer state as an admin.", + ) sp.set_defaults(func=set_librarian_transfer) diff --git a/hera_librarian/client.py b/hera_librarian/client.py index cbc2158..77fe124 100644 --- a/hera_librarian/client.py +++ b/hera_librarian/client.py @@ -56,6 +56,10 @@ ) from .models.ping import PingRequest, PingResponse from .models.search import FileSearchRequest, FileSearchResponse, FileSearchResponses +from .models.transfers import ( + LibrarianTransfersStatusRequest, + LibrarianTransfersStatusResponse, +) from .models.uploads import ( UploadCompletionRequest, UploadInitiationRequest, @@ -548,6 +552,44 @@ def validate_file(self, file_name: str) -> list[FileValidationResponseItem]: return response.root + def set_librarian_status( + self, + librarian_name: str, + transfers_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. + transfers_enabled : bool + Whether transfers to this librarian should be enabled. + + Returns + ------- + bool + The new status. + """ + + try: + response = self.post( + endpoint="/api/v2/transfers/update", + request=LibrarianTransfersStatusRequest( + librarian_name=librarian_name, + transfers_enabled=transfers_enabled, + ), + response=LibrarianTransfersStatusResponse, + ) + except LibrarianHTTPError as e: + if e.status_code == 400 and "Librarian" in e.reason: + raise LibrarianError(e.reason) + else: + raise e + + return response.transfers_enabled + class AdminClient(LibrarianClient): """ @@ -1345,7 +1387,7 @@ def set_librarian_status( try: response = self.post( - endpoint="admin/librarians/transfer_status/set", + endpoint="admin/librarians/transfer_status/change", request=AdminChangeLibrarianTransferStatusRequest( librarian_name=librarian_name, transfers_enabled=transfers_enabled, diff --git a/hera_librarian/models/transfers.py b/hera_librarian/models/transfers.py new file mode 100644 index 0000000..bec127c --- /dev/null +++ b/hera_librarian/models/transfers.py @@ -0,0 +1,67 @@ +""" +Pydantic models for the remote transfer update endpoints +""" + +from datetime import datetime +from typing import Optional + +from pydantic import BaseModel, RootModel + + +class LibrarianTransfersStatusRequest(BaseModel): + """ + A request to report 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 + +class LibrarianTransfersStatusResponse(BaseModel): + """ + A response to a user transfer status request. + """ + "The name of the librarian that was changed." + librarian_name: str + + "Whether the librarian has outbound transfers enabled." + transfers_enabled: bool + +class LocalLibrarianTransfersUpdateRequest(BaseModel): + """ + A request to change the transfer status of a librarian, either + to enable or disable inbound transfers. + """ + "The name of the librarian to change the transfer status from." + librarian_name: str + "Whether to enable or disable outbound transfers." + transfers_enabled: bool + +class LibrarianTransfersUpdateRequest(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 transfers." + transfers_enabled: bool + +class LibrarianTransfersUpdateResponse(BaseModel): + """ + A response to a user change request. + """ + + "The name of the librarian that was changed." + librarian_name: str + + "Whether the librarian has outbound transfers enabled." + transfers_enabled: bool + +class LibrarianTransfersFailedResponse(BaseModel): + reason: str + "The reason why the search failed." + + suggested_remedy: str + "A suggested remedy for the failure." \ No newline at end of file diff --git a/librarian_server/api/admin.py b/librarian_server/api/admin.py index d22a38c..68589d6 100644 --- a/librarian_server/api/admin.py +++ b/librarian_server/api/admin.py @@ -588,7 +588,7 @@ def delete_local_instance( @router.post( path="/librarians/transfer_status/change", - response_model=AdminLibrarianTransferStatusResponse, + response_model=AdminLibrarianTransferStatusResponse | AdminRequestFailedResponse, ) def change_librarian_transfer_status( request: AdminChangeLibrarianTransferStatusRequest, @@ -605,6 +605,12 @@ def change_librarian_transfer_status( session.query(Librarian).filter_by(name=request.librarian_name).one_or_none() ) + log.info( + f"Request from {user.username} to change transfer status of librarian " + f"{request.librarian_name} to {request.transfers_enabled}. " + f"Librarian found: {librarian}" + ) + if librarian is None: response.status_code = status.HTTP_400_BAD_REQUEST return AdminRequestFailedResponse( diff --git a/librarian_server/api/clone.py b/librarian_server/api/clone.py index 6df1769..1ba2a78 100644 --- a/librarian_server/api/clone.py +++ b/librarian_server/api/clone.py @@ -163,7 +163,7 @@ def de_duplicate_file_and_transfer( ) if existing_transfer.status == TransferStatus.ONGOING: - log.info(f"Found existing transfer with status ONGOING. Returning error.") + log.info("Found existing transfer with status ONGOING. Returning error.") raise HTTPException( status.HTTP_425_TOO_EARLY, diff --git a/librarian_server/api/transfers.py b/librarian_server/api/transfers.py new file mode 100644 index 0000000..55b1474 --- /dev/null +++ b/librarian_server/api/transfers.py @@ -0,0 +1,124 @@ +""" +Check in and modify the states of source and destination +transfers. +""" + +from fastapi import APIRouter, Depends, Response, status as status_codes +from sqlalchemy.orm import Session + +from hera_librarian.authlevel import AuthLevel +from hera_librarian.models.transfers import ( + LibrarianTransfersFailedResponse, + LibrarianTransfersStatusRequest, + LibrarianTransfersStatusResponse, + LibrarianTransfersUpdateRequest, + LibrarianTransfersUpdateResponse, + LocalLibrarianTransfersUpdateRequest, +) + +from ..database import yield_session +from ..logger import log +from ..orm.librarian import Librarian +from .auth import CallbackUserDependency, AdminUserDependency + +router = APIRouter(prefix="/api/v2/transfers") + +@router.post("/status", response_model=LibrarianTransfersStatusResponse) +def status( + request: LibrarianTransfersStatusRequest, + response: Response, + user: CallbackUserDependency, + session: Session = Depends(yield_session), +) -> LibrarianTransfersStatusResponse: + """ + Checkin and request the transfer status. + """ + + log.debug(f"Received checkin transfer status request from {user.username}: {request}") + + librarian = ( + session.query(Librarian).filter_by(name=user.username).one_or_none() + ) + + if librarian is None: + response.status_code = status_codes.HTTP_400_BAD_REQUEST + return LibrarianTransfersFailedResponse( + reason=f"Librarian {request.librarian_name} does not exist", + suggested_remedy="Please verify that the requested librarian exists", + ) + + if librarian.name != request.librarian_name: + response.status_code = status_codes.HTTP_403_FORBIDDEN + return LibrarianTransfersFailedResponse( + reason="Cannot check the status of another librarian.", + suggested_remedy="Please verify that you are the librarian making this request", + ) + + response = LibrarianTransfersStatusResponse( + librarian_name=librarian.name, + transfers_enabled=librarian.transfers_enabled, + ) + + log.debug(f"Responding to checkin request with: {response}.") + + return response + + +@router.post("/update", response_model=LibrarianTransfersUpdateResponse | LibrarianTransfersFailedResponse) +def update( + request: LibrarianTransfersUpdateRequest | LocalLibrarianTransfersUpdateRequest, + response: Response, + user: CallbackUserDependency | AdminUserDependency, + session: Session = Depends(yield_session), +) -> LibrarianTransfersUpdateResponse: + """ + Update the transfer status. + """ + + log.debug(f"Received update transfer status request from {user.username}: {request}") + + if user.permission == AuthLevel.ADMIN: + librarian = ( + session.query(Librarian).filter_by(name=request.librarian_name).one_or_none() + ) + update_request = LibrarianTransfersUpdateRequest( + librarian_name=user.username, + transfers_enabled=request.transfers_enabled) + response: LibrarianTransfersUpdateResponse = librarian.client.post( + endpoint="transfers/update", + request=update_request, + response=LibrarianTransfersUpdateResponse, + ) + + else: + + librarian = ( + session.query(Librarian).filter_by(name=user.username).one_or_none() + ) + + if librarian is None: + response.status_code = status_codes.HTTP_400_BAD_REQUEST + return LibrarianTransfersFailedResponse( + reason=f"Librarian {request.librarian_name} does not exist", + suggested_remedy="Please verify that the requested librarian exists", + ) + + if librarian.name != request.librarian_name: + response.status_code = status_codes.HTTP_403_FORBIDDEN + return LibrarianTransfersFailedResponse( + reason="Cannot change the status of another librarian.", + suggested_remedy="Please verify that you are the librarian making this request", + ) + + librarian.transfers_enabled = request.transfers_enabled + + session.commit() + + response = LibrarianTransfersUpdateResponse( + librarian_name=librarian.name, + transfers_enabled=librarian.transfers_enabled, + ) + + log.debug(f"Responding to update request with: {response}.") + + return response diff --git a/tests/conftest.py b/tests/conftest.py index fc99e06..8e28e75 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -91,6 +91,42 @@ def test_server(tmp_path_factory): os.environ[env_var] = env_vars[env_var] +@pytest.fixture(scope="session") +def test_server_disable(tmp_path_factory): + """ + Starts a single 'server' using the test client. + """ + + setup = server_setup(tmp_path_factory, name="test_server_disable") + + env_vars = {x: None for x in setup.env.keys()} + + for env_var in list(env_vars.keys()): + env_vars[env_var] = os.environ.get(env_var, None) + if setup.env[env_var] is not None: + os.environ[env_var] = setup.env[env_var] + + global DATABASE_PATH + DATABASE_PATH = str(setup.database) + + # Before starting, create the DB schema + run([sys.executable, shutil.which("librarian-server-setup")], env=setup.env) + + import librarian_server + import librarian_server.database + + app = librarian_server.main() + get_session = librarian_server.database.get_session + + yield app, get_session, setup + + for env_var in list(env_vars.keys()): + if env_vars[env_var] is None: + del os.environ[env_var] + else: + os.environ[env_var] = env_vars[env_var] + + def create_test_client(test_server, username="admin", password="password"): app, session, setup = test_server diff --git a/tests/integration_test/test_enable_disable_transfers.py b/tests/integration_test/test_enable_disable_transfers.py new file mode 100644 index 0000000..30ac322 --- /dev/null +++ b/tests/integration_test/test_enable_disable_transfers.py @@ -0,0 +1,76 @@ +""" +Integration test for enabling and disabling transfers to a librarian. +""" + +import random +from pathlib import Path + +import pytest + +from hera_librarian.exceptions import LibrarianError + +# What I have to do: +# 1. Create two servers (done via fixtures) +# 2. Register each other as librarians (done) +# 3. Authenticate the clients as admin (done) +# 4. Disable transfers on one of them via the API (to test the API) +# 5. Check the database to ensure the status is updated (to test the database) +# 6. Re-enable transfers on the other one via the API (to test the API) +# 7. Check the database to ensure the status is updated (to test the database) + + +def test_enable_disable_transfers( + test_server_disable, + test_orm, + mocked_admin_client, + server, + admin_client, + librarian_database_session_maker, + tmp_path, +): + # Register downstream and upstream librarians + assert mocked_admin_client.add_librarian( + name="live_server", + url="http://localhost", + authenticator="admin:password", + port=server.id, + ) + assert admin_client.add_librarian( + name="test_server", + url="http://localhost", + authenticator="admin:password", + port=test_server_disable[2].id, + check_connection=False, + ) + + # Disable transfers for test_server using the API + result = admin_client.set_librarian_status( + librarian_name="test_server", + transfers_enabled=False, + ) + assert result is False + + # Verify in the database + with librarian_database_session_maker() as session: + test_librarian = ( + session.query(test_orm.Librarian).filter_by(name="test_server").one() + ) + assert test_librarian.transfers_enabled is False + + # # Re-enable transfers for live_server using the API + result = admin_client.set_librarian_status( + librarian_name="test_server", + transfers_enabled=True, + ) + assert result is True + + # # Verify in the database + with librarian_database_session_maker() as session: + test_librarian = ( + session.query(test_orm.Librarian).filter_by(name="test_server").one() + ) + assert test_librarian.transfers_enabled is True + + # # Clean up + assert mocked_admin_client.remove_librarian(name="live_server") + assert admin_client.remove_librarian(name="test_server") diff --git a/tests/integration_test/test_sneaker.py b/tests/integration_test/test_sneaker.py index 28af76c..a026bfe 100644 --- a/tests/integration_test/test_sneaker.py +++ b/tests/integration_test/test_sneaker.py @@ -22,7 +22,7 @@ def test_sneakernet_workflow( # Representing the source in this test: # server, admin_client, librarian_database_session_maker - # We have to have it this way aruond because we need to be able + # We have to have it this way around because we need to be able # to make a callback to the source server, and we don't have # a mocking capability for that yet.