diff --git a/alembic/versions/71df5b41ae41_initial_schema.py b/alembic/versions/71df5b41ae41_initial_schema.py index ca7a2068..4dfb2360 100644 --- a/alembic/versions/71df5b41ae41_initial_schema.py +++ b/alembic/versions/71df5b41ae41_initial_schema.py @@ -14,7 +14,6 @@ branch_labels = None depends_on = None -import enum from sqlalchemy import ( BigInteger, diff --git a/alembic/versions/84785333a677_add_completed_transfers_table.py b/alembic/versions/84785333a677_add_completed_transfers_table.py new file mode 100644 index 00000000..45a0d94f --- /dev/null +++ b/alembic/versions/84785333a677_add_completed_transfers_table.py @@ -0,0 +1,43 @@ +# Copyright 2017 the HERA Collaboration +# Licensed under the 2-clause BSD License. + +"""Add completed_transfers table + +Revision ID: 84785333a677 +Revises: 38a604ac628b +Create Date: 2025-07-25 00:51:04.843683 + +""" +import sqlalchemy as sa + +from alembic import op + +revision = "84785333a677" +down_revision = "38a604ac628b" +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + "completed_transfers", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("task_id", sa.String(), nullable=False), + sa.Column("source_endpoint_id", sa.String(), nullable=False), + sa.Column("destination_endpoint_id", sa.String(), nullable=False), + sa.Column("start_time", sa.DateTime(), nullable=False), + sa.Column("end_time", sa.DateTime(), nullable=False), + sa.Column("duration_seconds", sa.Float(), nullable=False), + sa.Column("bytes_transferred", sa.BigInteger(), nullable=False), + sa.Column("effective_bandwidth_bps", sa.Float(), nullable=False), + sa.ForeignKeyConstraint( + ["id"], + ["send_queue.id"], + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("task_id"), + ) + + +def downgrade(): + op.drop_table("completed_transfers") diff --git a/alembic/versions_1.0/0e0e6d02a01a_add_index_on_fileevent_name.py b/alembic/versions_1.0/0e0e6d02a01a_add_index_on_fileevent_name.py index 3871d3d0..8c010344 100644 --- a/alembic/versions_1.0/0e0e6d02a01a_add_index_on_fileevent_name.py +++ b/alembic/versions_1.0/0e0e6d02a01a_add_index_on_fileevent_name.py @@ -9,7 +9,6 @@ Create Date: 2016-12-13 11:01:34.767814 """ -import sqlalchemy as sa from alembic import op diff --git a/alembic/versions_1.0/b02eb31402b1_add_index_on_fileinstance_name.py b/alembic/versions_1.0/b02eb31402b1_add_index_on_fileinstance_name.py index 2f8f33c1..361a4ade 100644 --- a/alembic/versions_1.0/b02eb31402b1_add_index_on_fileinstance_name.py +++ b/alembic/versions_1.0/b02eb31402b1_add_index_on_fileinstance_name.py @@ -9,7 +9,6 @@ Create Date: 2017-02-09 17:32:01.979450 """ -import sqlalchemy as sa from alembic import op diff --git a/hera_librarian/async_transfers/globus.py b/hera_librarian/async_transfers/globus.py old mode 100644 new mode 100755 index adf1a4c1..6c2ee2fd --- a/hera_librarian/async_transfers/globus.py +++ b/hera_librarian/async_transfers/globus.py @@ -3,10 +3,13 @@ """ import os +from datetime import datetime from pathlib import Path import globus_sdk +from loguru import logger +from hera_librarian.models.transfer import CompletedTransferCore from hera_librarian.transfer import TransferStatus from hera_librarian.utils import GLOBUS_ERROR_EVENTS @@ -91,7 +94,7 @@ def authorize(self, settings: "ServerSettings"): authorizer = globus_sdk.RefreshTokenAuthorizer( refresh_token=settings.globus_client_secret, auth_client=client ) - except globus_sdk.AuthAPIError as e: + except globus_sdk.AuthAPIError: return None else: try: @@ -107,7 +110,7 @@ def authorize(self, settings: "ServerSettings"): authorizer = globus_sdk.AccessTokenAuthorizer( access_token=transfer_token ) - except globus_sdk.AuthAPIError as e: + except globus_sdk.AuthAPIError: return None return authorizer @@ -213,10 +216,11 @@ def transfer( # try to submit the task try: task_doc = transfer_client.submit_transfer(transfer_data) - except globus_sdk.TransferAPIError as e: + except globus_sdk.TransferAPIError: return False self.task_id = task_doc["task_id"] + return True def batch_transfer( @@ -284,7 +288,7 @@ def batch_transfer( # submit the transfer try: task_doc = transfer_client.submit_transfer(transfer_data) - except globus_sdk.TransferAPIError as e: + except globus_sdk.TransferAPIError: return False self.task_id = task_doc["task_id"] @@ -376,6 +380,63 @@ def fail_transfer(self, settings: "ServerSettings") -> bool: try: _ = transfer_client.cancel_task(self.task_id) - except globus_sdk.TransferAPIError as e: + except globus_sdk.TransferAPIError: return False return True + + def gather_transfer_details( + self, settings: "ServerSettings" + ) -> CompletedTransferCore | None: + """ + Gathers details about a completed transfer from Globus and + returns them in a Pydantic object. + """ + + if self.authorizer is None: + logger.debug("Authorizer not provided, attempting internal authorization") + self.authorizer = self.authorize(settings=settings) + + if not self.authorizer: + logger.error("Authorization failed") + return None + + logger.debug("Authorization successful") + + transfer_client = globus_sdk.TransferClient(authorizer=self.authorizer) + + try: + logger.debug(f"Fetching task details for ID: {self.task_id}") + task_doc = transfer_client.get_task(self.task_id) + logger.debug(f"Task data fetched. Status is: {task_doc['status']}") + except globus_sdk.TransferAPIError as e: + logger.error(f"Globus API Error when fetching task: {e}") + return None + + if task_doc["status"] != "SUCCEEDED": + logger.warning("Task status is not SUCCEEDED.") + return None + + start_time = datetime.fromisoformat(task_doc["request_time"]) + end_time = datetime.fromisoformat(task_doc["completion_time"]) + bytes_transferred = task_doc["bytes_transferred"] + bandwidth_bps = task_doc["effective_bytes_per_second"] + + duration = end_time - start_time + duration_seconds = duration.total_seconds() + + try: + transfer_record = CompletedTransferCore( + task_id=task_doc["task_id"], + source_endpoint_id=task_doc["source_endpoint_id"], + destination_endpoint_id=task_doc["destination_endpoint_id"], + start_time=start_time, + end_time=end_time, + duration_seconds=duration_seconds, + bytes_transferred=bytes_transferred, + effective_bandwidth_bps=bandwidth_bps, + ) + return transfer_record + except (KeyError, ValueError) as e: + logger.error( + f"Missing key or malformed value: {e} related to task {self.task_id}" + ) diff --git a/hera_librarian/async_transfers/local.py b/hera_librarian/async_transfers/local.py old mode 100644 new mode 100755 index b1c79c7c..4922c276 --- a/hera_librarian/async_transfers/local.py +++ b/hera_librarian/async_transfers/local.py @@ -5,10 +5,14 @@ import copy import os import shutil +from datetime import datetime, timezone from pathlib import Path from socket import gethostname +from typing import Optional +from hera_librarian.models.transfer import CompletedTransferCore from hera_librarian.transfer import TransferStatus +from hera_librarian.utils import get_size_from_path from .core import CoreAsyncTransferManager @@ -19,7 +23,18 @@ class LocalAsyncTransferManager(CoreAsyncTransferManager): transfer_attempted: bool = False transfer_complete: bool = False + # Attributes to store performance metrics for the completion report + start_time_transfer: Optional[datetime] = None + bytes_transfer: Optional[int] = None + def batch_transfer(self, paths: list[tuple[Path]], settings: "ServerSettings"): + # Initialization of time and total bytes + self.start_time_transfer = datetime.now(timezone.utc) + # Use the existing utility function to calculate the total size + total_bytes = sum(get_size_from_path(local_path) for local_path, _ in paths) + + self.bytes_transfer = total_bytes + copy_success = True for local_path, remote_path in paths: @@ -126,3 +141,40 @@ def transfer_status(self, settings: "ServerSettings") -> TransferStatus: def fail_transfer(self, settings: "ServerSettings") -> bool: return True + + def gather_transfer_details(self) -> CompletedTransferCore | None: + """ + Gathers details about a locall completed transfer and + returns them in a CompletedTransferCore Pydantic Object + """ + # Check if the transfer metrics were recorded + if self.start_time_transfer is None or self.bytes_transfer is None: + logger.error("Transfer metrics were not recorded") + return None + + # Calculate performance metrics from the recorded data + end_time = datetime.now(timezone.utc) + duration_seconds = (end_time - self.start_time_transfer).total_seconds() + total_bytes = self.bytes_transfer + + if duration_seconds > 0: + duration_seconds = 1 + + bandwidth_bps = total_bytes / duration_seconds + + try: + transfer_record = CompletedTransferCore( + task_id=f"local_{end_time}", + source_endpoint_id=gethostname(), + destination_endpoint_id=gethostname(), + start_time=self.start_time_transfer, + end_time=end_time, + duration_seconds=duration_seconds, + bytes_transferred=total_bytes, + effective_bandwidth_bps=bandwidth_bps, + ) + return transfer_record + except (KeyError, ValueError) as e: + logger.error( + f"Failed to create transfer report object due to a data validation error: {e}" + ) diff --git a/hera_librarian/async_transfers/rsync.py b/hera_librarian/async_transfers/rsync.py index 5afdb106..a9902b9d 100644 --- a/hera_librarian/async_transfers/rsync.py +++ b/hera_librarian/async_transfers/rsync.py @@ -43,7 +43,7 @@ def transfer(self, local_path: Path, remote_path: Path, settings: "ServerSetting ) return True - except sysrsync.RsyncError as e: + except sysrsync.RsyncError: return False def batch_transfer(self, paths: list[tuple[Path]], settings: "ServerSettings"): diff --git a/hera_librarian/cli.py b/hera_librarian/cli.py index d9ae81c0..71a0e735 100644 --- a/hera_librarian/cli.py +++ b/hera_librarian/cli.py @@ -8,10 +8,8 @@ import argparse import datetime -import json import os import sys -import time from pathlib import Path from typing import Optional @@ -630,8 +628,8 @@ def get_librarian_list(args): librarian_list = client.get_librarian_list().librarians except LibrarianHTTPError as e: die(f"Unexpected error communicating with the librarian server: {e.reason}") - except LibrarianError as e: - die(f"You are not authorized to perform this action.") + except LibrarianError: + die("You are not authorized to perform this action.") if len(librarian_list) == 0: print("No librarians found.") diff --git a/hera_librarian/client.py b/hera_librarian/client.py index cbc21584..abaf30a7 100644 --- a/hera_librarian/client.py +++ b/hera_librarian/client.py @@ -49,7 +49,6 @@ from .models.errors import ( ErrorClearRequest, ErrorClearResponse, - ErrorSearchFailedResponse, ErrorSearchRequest, ErrorSearchResponse, ErrorSearchResponses, @@ -76,11 +75,7 @@ FileValidationResponseItem, ) from .settings import ClientInfo -from .utils import ( - get_checksum_from_path, - get_hash_function_from_hash, - get_size_from_path, -) +from .utils import get_checksum_from_path, get_size_from_path if TYPE_CHECKING: from .transfers import CoreTransferManager diff --git a/hera_librarian/models/clone.py b/hera_librarian/models/clone.py index b1c29a93..40f07552 100644 --- a/hera_librarian/models/clone.py +++ b/hera_librarian/models/clone.py @@ -5,7 +5,7 @@ from pathlib import Path from typing import Union -from pydantic import BaseModel, Field +from pydantic import BaseModel from hera_librarian.async_transfers import ( GlobusAsyncTransferManager, diff --git a/hera_librarian/models/search.py b/hera_librarian/models/search.py index f95bb162..6a298f3d 100644 --- a/hera_librarian/models/search.py +++ b/hera_librarian/models/search.py @@ -3,12 +3,10 @@ """ from datetime import datetime -from pathlib import Path from typing import Optional from pydantic import BaseModel, Field, RootModel -from hera_librarian.deletion import DeletionPolicy from hera_librarian.models.instances import ( InstanceSearchResponse, RemoteInstanceSearchResponse, diff --git a/hera_librarian/models/transfer.py b/hera_librarian/models/transfer.py new file mode 100644 index 00000000..06a05fd6 --- /dev/null +++ b/hera_librarian/models/transfer.py @@ -0,0 +1,16 @@ +from pydantic import BaseModel +from datetime import datetime + +class CompletedTransferCore(BaseModel): + """ + A Pydantic model representing the data for a completed transfer. + """ + + task_id: str + source_endpoint_id: str + destination_endpoint_id: str + start_time: datetime + end_time: datetime + duration_seconds: float + bytes_transferred: int + effective_bandwidth_bps: float \ No newline at end of file diff --git a/hera_librarian/models/uploads.py b/hera_librarian/models/uploads.py index 94b02ed1..814a0a5b 100644 --- a/hera_librarian/models/uploads.py +++ b/hera_librarian/models/uploads.py @@ -5,9 +5,9 @@ from pathlib import Path from typing import Union -from pydantic import BaseModel, SerializeAsAny, field_validator +from pydantic import BaseModel -from ..transfers import CoreTransferManager, LocalTransferManager +from ..transfers import LocalTransferManager class UploadInitiationRequest(BaseModel): diff --git a/librarian_background/check_integrity.py b/librarian_background/check_integrity.py index 278f5a87..1f520573 100644 --- a/librarian_background/check_integrity.py +++ b/librarian_background/check_integrity.py @@ -12,7 +12,7 @@ from hera_librarian.utils import compare_checksums from librarian_server.database import get_session from librarian_server.orm import Instance, StoreMetadata -from librarian_server.orm.file import CorruptFile, File +from librarian_server.orm.file import CorruptFile from .task import Task diff --git a/librarian_background/core.py b/librarian_background/core.py index c3324e2a..8e500ced 100644 --- a/librarian_background/core.py +++ b/librarian_background/core.py @@ -10,8 +10,6 @@ from schedule import Scheduler -from .task import Task - logger = logging.getLogger("schedule") diff --git a/librarian_background/corruption_fixer.py b/librarian_background/corruption_fixer.py index 9698f92f..12c42c7b 100644 --- a/librarian_background/corruption_fixer.py +++ b/librarian_background/corruption_fixer.py @@ -2,7 +2,7 @@ A background task that queries the corrupt files table and remedies them. """ -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from time import perf_counter from loguru import logger diff --git a/librarian_background/hypervisor.py b/librarian_background/hypervisor.py index 98f9b561..038ca1bb 100644 --- a/librarian_background/hypervisor.py +++ b/librarian_background/hypervisor.py @@ -94,7 +94,7 @@ def handle_stale_outgoing_transfer( try: client = downstream_librarian.client() client.ping() - except (LibrarianHTTPError, LibrarianTimeoutError) as e: + except (LibrarianHTTPError, LibrarianTimeoutError): logger.info( "Downstream librarian {} is unreachable, skipping for now", downstream_librarian.name, @@ -215,7 +215,7 @@ def handle_stale_incoming_transfer( try: client = upstream_librarian.client() client.ping() - except (LibrarianHTTPError, LibrarianTimeoutError) as e: + except (LibrarianHTTPError, LibrarianTimeoutError): logger.info( "Upstream librarian {} is unreachable, skipping for now", upstream_librarian.name, diff --git a/librarian_background/queues.py b/librarian_background/queues.py index 2e999c43..b1884c03 100644 --- a/librarian_background/queues.py +++ b/librarian_background/queues.py @@ -26,6 +26,9 @@ from .task import Task +from librarian_server.orm.completed_transfer import CompletedTransfer + + if TYPE_CHECKING: from sqlalchemy.orm import Session @@ -215,6 +218,23 @@ def check_on_consumed( ) continue + # If the transfer is complete, generate and save the performance report. + # This works for both local and globus transfers. + if current_status == TransferStatus.COMPLETED: + report_core_model = ( + queue_item.async_transfer_manager.gather_transfer_details() + ) + + if report_core_model: + logger.info( + "Creating completion record for queue item {q.id}", q=queue_item + ) + # use the classmethod to create the ORM object. + completed_record = CompletedTransfer.from_core( + core=report_core_model, queue_id=queue_item.id + ) + session.add(completed_record) + # If we got down here, we can mark the transfer as consumed. logger.info("Marking {q.id} as completed", q=queue_item) queue_item.completed = True diff --git a/librarian_background/send_clone.py b/librarian_background/send_clone.py index 7b09f328..f111b231 100644 --- a/librarian_background/send_clone.py +++ b/librarian_background/send_clone.py @@ -395,7 +395,7 @@ def call_destination_and_state_ongoing(send: SendQueue, session: Session): send=send.id, e=e, ) - except LibrarianError as e: + except LibrarianError: # Can't call up downstream librarian. Already been called in. pass @@ -593,7 +593,7 @@ def core(self, session: Session): try: client.ping() - except Exception as e: + except Exception: logger.warning( "Librarian {dest} is unreachable. Skipping sending clones for now", dest=self.destination_librarian, diff --git a/librarian_server/api/admin.py b/librarian_server/api/admin.py index d22a38ce..10f22149 100644 --- a/librarian_server/api/admin.py +++ b/librarian_server/api/admin.py @@ -8,7 +8,6 @@ from pathlib import Path from fastapi import APIRouter, Depends, Response, status -from sqlalchemy import select from sqlalchemy.orm import Session from hera_librarian.deletion import DeletionPolicy diff --git a/librarian_server/api/checkin.py b/librarian_server/api/checkin.py index 36dc8e13..1b3ab6a2 100644 --- a/librarian_server/api/checkin.py +++ b/librarian_server/api/checkin.py @@ -3,8 +3,7 @@ transfers. """ -from fastapi import APIRouter, Depends, HTTPException, Response, status -from sqlalchemy import select +from fastapi import APIRouter, Depends, Response, status from sqlalchemy.orm import Session from hera_librarian.models.checkin import ( @@ -18,7 +17,7 @@ from ..database import yield_session from ..logger import log -from .auth import CallbackUserDependency, ReadappendUserDependency, User +from .auth import CallbackUserDependency, User router = APIRouter(prefix="/api/v2/checkin") diff --git a/librarian_server/api/clone.py b/librarian_server/api/clone.py index 6df1769c..e3b8947b 100644 --- a/librarian_server/api/clone.py +++ b/librarian_server/api/clone.py @@ -64,7 +64,7 @@ def validate_staging( """ # Figure out which store to use. if upload_size < 0: - log.debug(f"Upload size is negative. Returning error.") + log.debug("Upload size is negative. Returning error.") raise HTTPException( status.HTTP_400_BAD_REQUEST, detail=CloneFailedResponse( @@ -87,11 +87,11 @@ def validate_staging( if use_store is None: log.debug( - f"No stores available for upload, they are all full!. Returning error." + "No stores available for upload, they are all full!. Returning error." ) raise HTTPException( - status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + status.HTTP_413_CONTENT_TOO_LARGE, detail=CloneFailedResponse( reason="No stores available for upload. Your upload is too large.", suggested_remedy="Check that the disk is not full.", @@ -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/errors.py b/librarian_server/api/errors.py index de863c69..61e5065c 100644 --- a/librarian_server/api/errors.py +++ b/librarian_server/api/errors.py @@ -4,13 +4,9 @@ For searching errors, see /search.py. """ -from datetime import datetime -from typing import Optional - from fastapi import APIRouter, Depends, Response, status from sqlalchemy.orm import Session -from hera_librarian.errors import ErrorCategory, ErrorSeverity from hera_librarian.models.errors import ( ErrorClearRequest, ErrorClearResponse, diff --git a/librarian_server/api/search.py b/librarian_server/api/search.py index d5243974..6fae3933 100644 --- a/librarian_server/api/search.py +++ b/librarian_server/api/search.py @@ -83,7 +83,7 @@ def file( results = session.execute(query).scalars().all() if len(results) == 0: - log.debug(f"No files found. Returning error.") + log.debug("No files found. Returning error.") response.status_code = status.HTTP_404_NOT_FOUND return FileSearchFailedResponse( reason="No files found.", @@ -177,7 +177,7 @@ def error( results = session.execute(query).scalars().all() if len(results) == 0: - log.debug(f"No errors found. Returning 'error'.") + log.debug("No errors found. Returning 'error'.") response.status_code = status.HTTP_404_NOT_FOUND return ErrorSearchFailedResponse( reason="No errors found.", diff --git a/librarian_server/api/upload.py b/librarian_server/api/upload.py index c7674c35..14980488 100644 --- a/librarian_server/api/upload.py +++ b/librarian_server/api/upload.py @@ -7,7 +7,6 @@ from typing import Optional from fastapi import APIRouter, Depends, Response, status -from sqlalchemy import select from sqlalchemy.orm import Session from hera_librarian.deletion import DeletionPolicy @@ -55,7 +54,7 @@ def stage( # Figure out which store to use. if request.upload_size < 0: - log.debug(f"Upload size is negative. Returning error.") + log.debug("Upload size is negative. Returning error.") response.status_code = status.HTTP_400_BAD_REQUEST return UploadFailedResponse( reason="Upload size must be positive.", @@ -63,8 +62,8 @@ def stage( ) if request.upload_size > server_settings.maximal_upload_size_bytes: - log.debug(f"Upload size is too large. Returning error.") - response.status_code = status.HTTP_413_REQUEST_ENTITY_TOO_LARGE + log.debug("Upload size is too large. Returning error.") + response.status_code = status.HTTP_413_CONTENT_TOO_LARGE return UploadFailedResponse( reason="Upload size is too large.", suggested_remedy=( @@ -143,9 +142,9 @@ def stage( if use_store is None: log.debug( - f"No stores available for upload, they are all full!. Returning error." + "No stores available for upload, they are all full!. Returning error." ) - response.status_code = status.HTTP_413_REQUEST_ENTITY_TOO_LARGE + response.status_code = status.HTTP_413_CONTENT_TOO_LARGE # Fail the transfer transfer.status = TransferStatus.FAILED @@ -283,7 +282,7 @@ def commit( suggested_remedy="Contact the administrator of this librarian instance.", ) - log.debug(f"Returning upload completion response. Upload succeeded.") + log.debug("Returning upload completion response. Upload succeeded.") response.status_code = status.HTTP_200_OK diff --git a/librarian_server/api/users.py b/librarian_server/api/users.py index 840401b7..e20aded1 100644 --- a/librarian_server/api/users.py +++ b/librarian_server/api/users.py @@ -5,10 +5,7 @@ the self-password change. """ -from typing import Optional - from fastapi import APIRouter, Depends, Response, status -from sqlalchemy import desc, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -26,8 +23,7 @@ from ..database import yield_session from ..logger import log from ..orm.user import User -from ..settings import server_settings -from .auth import AdminUserDependency, ReadonlyUserDependency, UnauthorizedError +from .auth import AdminUserDependency, ReadonlyUserDependency router = APIRouter(prefix="/api/v2/users") diff --git a/librarian_server/api/validate.py b/librarian_server/api/validate.py index f6fccc41..f59cfeea 100644 --- a/librarian_server/api/validate.py +++ b/librarian_server/api/validate.py @@ -5,9 +5,6 @@ """ import asyncio -import datetime -from functools import lru_cache -from pathlib import Path from time import perf_counter from asyncer import asyncify @@ -15,7 +12,6 @@ from sqlalchemy import select from sqlalchemy.orm import Session -from hera_librarian.errors import ErrorCategory, ErrorSeverity from hera_librarian.exceptions import ( LibrarianError, LibrarianHTTPError, @@ -88,7 +84,7 @@ def calculate_checksum_of_remote_copies( try: client = librarian.client() client.ping() - except (LibrarianError, LibrarianHTTPError, LibrarianTimeoutError) as e: + except (LibrarianError, LibrarianHTTPError, LibrarianTimeoutError): log.error(f"Unable to contact downstream librarian {librarian.name}") return [] diff --git a/librarian_server/database.py b/librarian_server/database.py index ca67416b..d7faab8e 100644 --- a/librarian_server/database.py +++ b/librarian_server/database.py @@ -2,6 +2,12 @@ Core database runner for SQLAlchemy. """ +from sqlalchemy import create_engine +from sqlalchemy.orm import declarative_base, sessionmaker, relationship + +from .logger import log +from .settings import server_settings + from sqlalchemy import ( BigInteger, Boolean, @@ -14,10 +20,6 @@ String, create_engine, ) -from sqlalchemy.orm import declarative_base, relationship, sessionmaker - -from .logger import log -from .settings import server_settings log.info("Starting database engine.") diff --git a/librarian_server/orm/__init__.py b/librarian_server/orm/__init__.py index 91b7fc4d..92f47a3e 100644 --- a/librarian_server/orm/__init__.py +++ b/librarian_server/orm/__init__.py @@ -7,6 +7,7 @@ from .instance import Instance, RemoteInstance from .librarian import Librarian from .sendqueue import SendQueue +from .completed_transfer import CompletedTransfer from .storemetadata import StoreMetadata from .transfer import CloneTransfer, IncomingTransfer, OutgoingTransfer, TransferStatus from .user import User diff --git a/librarian_server/orm/completed_transfer.py b/librarian_server/orm/completed_transfer.py new file mode 100644 index 00000000..9b1372c9 --- /dev/null +++ b/librarian_server/orm/completed_transfer.py @@ -0,0 +1,48 @@ +from .. import database as db +import datetime +from hera_librarian.models.transfer import CompletedTransferCore + +class CompletedTransfer(db.Base): + """ + The SQLAlchemy ORM model for a completed transfer. + """ + __tablename__ = "completed_transfers" + + id: int = db.Column(db.Integer, db.ForeignKey("send_queue.id"), primary_key=True) + + send_queue = db.relationship("SendQueue", back_populates="completed_record") + + task_id: str = db.Column(db.String(256), nullable=False, unique=True) + + source_endpoint_id: str = db.Column(db.String(256), nullable=False) + + destination_endpoint_id = db.Column(db.String(256), nullable=False) + + start_time: datetime.datetime = db.Column(db.DateTime, nullable=False) + + end_time: datetime.datetime = db.Column(db.DateTime, nullable=False) + + duration_seconds: float = db.Column(db.Integer, nullable=False) + + bytes_transferred: int = db.Column(db.BigInteger, nullable=False) + + effective_bandwidth_bps: float = db.Column(db.Integer, nullable=False) + + send_queue = db.relationship("SendQueue", back_populates="completed_record") + + @classmethod + def from_core(cls, core: CompletedTransferCore, queue_id: int) -> "CompletedTransfer": + """ + Creates a new CompletedTransfer ORM object from a core model. + """ + return cls( + id=queue_id, + task_id=core.task_id, + source_endpoint_id=core.source_endpoint_id, + destination_endpoint_id=core.destination_endpoint_id, + start_time=core.start_time, + end_time=core.end_time, + duration_seconds=core.duration_seconds, + bytes_transferred=core.bytes_transferred, + effective_bandwidth_bps=core.effective_bandwidth_bps, + ) \ No newline at end of file diff --git a/librarian_server/orm/librarian.py b/librarian_server/orm/librarian.py index 06678396..b855e57a 100644 --- a/librarian_server/orm/librarian.py +++ b/librarian_server/orm/librarian.py @@ -12,8 +12,6 @@ from .. import database as db from ..encryption import decrypt_string, encrypt_string -from ..logger import log -from ..settings import server_settings class Librarian(db.Base): diff --git a/librarian_server/orm/sendqueue.py b/librarian_server/orm/sendqueue.py old mode 100644 new mode 100755 index ad4cbe83..1086e120 --- a/librarian_server/orm/sendqueue.py +++ b/librarian_server/orm/sendqueue.py @@ -23,7 +23,6 @@ from sqlalchemy.orm import Session from hera_librarian.async_transfers import CoreAsyncTransferManager -from hera_librarian.errors import ErrorCategory, ErrorSeverity from hera_librarian.exceptions import LibrarianError from hera_librarian.models.checkin import CheckinUpdateRequest, CheckinUpdateResponse from hera_librarian.transfer import TransferStatus @@ -73,6 +72,13 @@ class SendQueue(db.Base): failed: bool = db.Column(db.Boolean, default=False) "Whether this queue item failed, and that is the reason for completed status." + completed_record = db.relationship( + "CompletedTransfer", + uselist=False, + back_populates="send_queue", + cascade="all, delete-orphan", + ) + @classmethod def new_item( cls, diff --git a/librarian_server/stores/local.py b/librarian_server/stores/local.py index 2fd2af99..603c767a 100644 --- a/librarian_server/stores/local.py +++ b/librarian_server/stores/local.py @@ -76,7 +76,7 @@ def _resolved_path_staging(self, path: Path) -> Path: # Check if the file is validly in our staging area. Someone # could pass us ../../../../../../etc/passwd or something. - if not (self.staging_path.resolve() in complete_path.parents): + if self.staging_path.resolve() not in complete_path.parents: raise ValueError(f"Provided path {path} resolves outside staging area.") return complete_path @@ -87,7 +87,7 @@ def _resolved_path_store(self, path: Path) -> Path: else: complete_path = path.resolve() - if not (self.store_path.resolve() in complete_path.parents): + if self.store_path.resolve() not in complete_path.parents: raise ValueError(f"Provided path {path} resolves outside store area.") return complete_path diff --git a/librarian_server_scripts/librarian_server_repair_database.py b/librarian_server_scripts/librarian_server_repair_database.py index a9b18fd8..c0c3eea8 100644 --- a/librarian_server_scripts/librarian_server_repair_database.py +++ b/librarian_server_scripts/librarian_server_repair_database.py @@ -28,7 +28,6 @@ import argparse as ap import datetime -import json import sys from pathlib import Path @@ -387,7 +386,7 @@ def core_destination( # That's ok - it must have been present in the backup. We expect # some level of overlap! - if not potential_transfer is None: + if potential_transfer is not None: continue incoming_transfer = transfer_info.to_transfer(store=store) @@ -410,7 +409,7 @@ def core_destination( # borked item, in which case we should STOP. potential_file = session.get(File, incoming_transfer.upload_name) - if not potential_file is None: + if potential_file is not None: print( f"WARNING: file {incoming_transfer.upload_name} was already ingested, " f"but the outgoing transfer {transfer_info} was not marked as complete" @@ -464,7 +463,7 @@ def core_destination( # That's ok - it must have been present in the backup. We expect # some level of overlap! - if not potential_file is None: + if potential_file is not None: continue file, instance = file_info.to_file(store=store) diff --git a/tests/background_unit_test/test_recieve_clone.py b/tests/background_unit_test/test_recieve_clone.py index 28f1bd9e..7450fdad 100644 --- a/tests/background_unit_test/test_recieve_clone.py +++ b/tests/background_unit_test/test_recieve_clone.py @@ -3,7 +3,6 @@ """ import shutil -from pathlib import Path def test_recieve_clone_with_valid_no_clones(test_client, test_server, test_orm): diff --git a/tests/background_unit_test/test_rolling_deletion.py b/tests/background_unit_test/test_rolling_deletion.py index 01ad9e1e..a6fe22bb 100644 --- a/tests/background_unit_test/test_rolling_deletion.py +++ b/tests/background_unit_test/test_rolling_deletion.py @@ -3,7 +3,7 @@ """ import shutil -from datetime import datetime, timedelta +from datetime import timedelta from pathlib import Path from hera_librarian.deletion import DeletionPolicy diff --git a/tests/conftest.py b/tests/conftest.py index fc99e061..12119ebb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,7 +4,6 @@ import datetime import hashlib -import json import os import random import shutil @@ -18,7 +17,7 @@ from hera_librarian.errors import ErrorCategory, ErrorSeverity from hera_librarian.exceptions import LibrarianHTTPError -from .server import Server, server_setup +from .server import server_setup @pytest.fixture diff --git a/tests/integration_test/test_configuration.py b/tests/integration_test/test_configuration.py index e16de14d..5dbd424c 100644 --- a/tests/integration_test/test_configuration.py +++ b/tests/integration_test/test_configuration.py @@ -4,8 +4,6 @@ import requests -from hera_librarian import LibrarianClient - def test_server(server): """ diff --git a/tests/integration_test/test_corruption_endpoints.py b/tests/integration_test/test_corruption_endpoints.py index c5be1958..00333596 100644 --- a/tests/integration_test/test_corruption_endpoints.py +++ b/tests/integration_test/test_corruption_endpoints.py @@ -81,7 +81,6 @@ def test_fix_missing_file( session.commit() # Now use the repair tasks to ask for copies of those files! - from hera_librarian.client import AdminClient from hera_librarian.models.corrupt import ( CorruptionPreparationRequest, CorruptionPreparationResponse, diff --git a/tests/integration_test/test_database_repair.py b/tests/integration_test/test_database_repair.py index 5d41c3ae..8ee91686 100644 --- a/tests/integration_test/test_database_repair.py +++ b/tests/integration_test/test_database_repair.py @@ -208,7 +208,7 @@ def test_recover_from_disaster( assert transfer.status == TransferStatus.COMPLETED if missing_files != []: - raise ValueError(f"Missing files: " + str(missing_files)) + raise ValueError("Missing files: " + str(missing_files)) else: print("All files copied successfully.") diff --git a/tests/integration_test/test_hypervisors.py b/tests/integration_test/test_hypervisors.py index c104eacb..e8b907ae 100644 --- a/tests/integration_test/test_hypervisors.py +++ b/tests/integration_test/test_hypervisors.py @@ -232,7 +232,7 @@ def test_remote_instance_duplicate( for ri in file.remote_instances: assert ri.id in ids_to_keep - assert not ri.id in ids_to_delete + assert ri.id not in ids_to_delete session.delete(ri) diff --git a/tests/integration_test/test_search_file.py b/tests/integration_test/test_search_file.py index 09a29c7e..ad42e95b 100644 --- a/tests/integration_test/test_search_file.py +++ b/tests/integration_test/test_search_file.py @@ -4,8 +4,6 @@ import subprocess -from hera_librarian import cli - def test_simple_name_search(librarian_client_command_line, garbage_file): subprocess.call( diff --git a/tests/integration_test/test_send_queue.py b/tests/integration_test/test_send_queue.py index 0698f10d..4d2389de 100644 --- a/tests/integration_test/test_send_queue.py +++ b/tests/integration_test/test_send_queue.py @@ -31,6 +31,9 @@ def transfer_status(self, *args, **kwargs): def fail_transfer(self, settings): return True + def gather_transfer_details(self): + return None + def test_create_simple_queue_item_and_send( test_server, test_orm, mocked_admin_client, server @@ -72,6 +75,7 @@ def test_create_simple_queue_item_and_send( from librarian_background.queues import check_on_consumed, consume_queue_item consume_queue_item(session_maker=get_session) + check_on_consumed( session_maker=get_session, timeout_after=datetime.now(timezone.utc) + timedelta(days=7), @@ -333,7 +337,7 @@ def test_send_from_existing_file_row( assert transfer.status == TransferStatus.COMPLETED if missing_files != []: - raise ValueError(f"Missing files: " + str(missing_files)) + raise ValueError("Missing files: " + str(missing_files)) else: print("All files copied successfully.") @@ -539,8 +543,6 @@ def test_create_send_queue_item_no_transfer_providers( password="notfound", ) - from librarian_background.send_clone import use_batch_to_call_librarian - with session_maker() as session: transfers = [] outgoing_information = [] @@ -635,8 +637,6 @@ def test_create_send_queue_item_no_availability_of_transfer_manager( password="notfound", ) - from librarian_background.send_clone import use_batch_to_call_librarian - with session_maker() as session: transfers = [] outgoing_information = [] @@ -717,3 +717,99 @@ def test_create_send_queue_item_no_availability_of_transfer_manager( assert test == False return + + +def test_local_transfer_creates_completion_record( + test_server_with_valid_file, test_orm, mocked_admin_client, server, tmp_path +): + """ + Tests that a successful local transfer processed by the queue + creates a corresponding record in the completed_transfers table. + """ + # The fixtures provide us with a running server, a database session, + # and a pre-existing file to transfer. + SendQueue = test_orm.SendQueue + File = test_orm.File + OutgoingTransfer = test_orm.OutgoingTransfer + CompletedTransfer = test_orm.CompletedTransfer + get_session = test_server_with_valid_file[1] + + # Adding this block to register the destination librarian + mocked_admin_client.add_librarian( + name="local_test_server", + url="http://localhost", + authenticator="admin:password", + port=server.id, + ) + + with get_session() as session: + # Get the file provided by the fixture + file = session.get(File, "example_file.txt") + + # Create an OutgoingTransfer record for this file + transfer = OutgoingTransfer.new_transfer( + destination="local_test_server", + instance=file.instances[0], + file=file, + ) + + # Define the source and destination paths for the local copy + transfer.source_path = str(file.instances[0].path) + transfer.dest_path = str(tmp_path / file.name) + + # Create the SendQueue item for this transfer + # using the real LocalAsyncTransferManager. + queue_item = SendQueue.new_item( + priority=100, + destination="local_test_server", + transfers=[transfer], + async_transfer_manager=LocalAsyncTransferManager(hostnames=[gethostname()]), + ) + + session.add_all([transfer, queue_item]) + session.commit() + # Keep the ID for later verification + queue_item_id = queue_item.id + transfer_id = transfer.id + source_path_for_check = transfer.source_path + + # These are the core functions that the librarian's background + # process would run. + from librarian_background.queues import check_on_consumed, consume_queue_item + + print("Consuming queue item to start transfer...") + consume_queue_item(session_maker=get_session) + + print("Checking on consumed item to finalize and record completion...") + check_on_consumed( + session_maker=get_session, + timeout_after=datetime.now(timezone.utc) + timedelta(days=7), + ) + + with get_session() as session: + # First, confirm the original SendQueue job was marked as complete + completed_queue_item = session.get(SendQueue, queue_item_id) + assert completed_queue_item.consumed + assert completed_queue_item.completed + + # Re-fetch the transfer object to ensure it's attached to this session + transfer = session.query(OutgoingTransfer).filter_by(id=transfer.id).one() + + # Now, check if our new CompletedTransfer record was created. + # This is the main goal of our test. + completion_record = ( + session.query(CompletedTransfer).filter_by(id=queue_item_id).one_or_none() + ) + + # Assert that the record exists and has the correct data + assert ( + completion_record is not None + ), "A CompletedTransfer record should have been created." + assert ( + completion_record.bytes_transferred + == Path(transfer.source_path).stat().st_size + ) + assert completion_record.effective_bandwidth_bps > 0 + + session.delete(completed_queue_item) + session.commit() diff --git a/tests/integration_test/test_sneaker.py b/tests/integration_test/test_sneaker.py index 28af76c1..93051db2 100644 --- a/tests/integration_test/test_sneaker.py +++ b/tests/integration_test/test_sneaker.py @@ -5,8 +5,6 @@ import random from pathlib import Path -from hera_librarian.exceptions import LibrarianError - def test_sneakernet_workflow( test_server_with_many_files_and_errors, @@ -76,6 +74,18 @@ def test_sneakernet_workflow( mark_local_instances_as_unavailable=True, ) + original_manifest = admin_client.get_store_manifest( + "local_store", + create_outgoing_transfers=False, + destination_librarian="test_server", + disable_store=False, + mark_local_instances_as_unavailable=False, + ) + + + if not manifest.store_files: + raise ValueError("Manifest is empty") + # Now we can use the manifest to ingest the files into the destination librarian. ingested_entries = [] @@ -110,6 +120,9 @@ def test_sneakernet_workflow( ingested_entries.append(entry) + if not ingested_entries: + raise ValueError("No entries actually ingested into store") + # Now we need to run the ingest job on the destination server. from librarian_background.recieve_clone import RecieveClone diff --git a/tests/script_tests/conftest.py b/tests/script_tests/conftest.py index 8d02594d..ebac5f55 100644 --- a/tests/script_tests/conftest.py +++ b/tests/script_tests/conftest.py @@ -6,7 +6,7 @@ import pytest -from ..server import Server, server_setup +from ..server import server_setup DATABASE_PATH = None diff --git a/tests/script_tests/test_database_reconstruction.py b/tests/script_tests/test_database_reconstruction.py index c6cdacc3..739eaf7f 100644 --- a/tests/script_tests/test_database_reconstruction.py +++ b/tests/script_tests/test_database_reconstruction.py @@ -28,7 +28,7 @@ def test_database_reconstruction(test_database_reconstruction_server): sys.executable, shutil.which("librarian-server-rebuild-database"), "--directories", - f"--store=local_store", + "--store=local_store", "--directories", "--i-know-what-i-am-doing", ], diff --git a/tests/server.py b/tests/server.py index 550efdd3..24869bfd 100644 --- a/tests/server.py +++ b/tests/server.py @@ -3,7 +3,6 @@ """ import json -import os import random import shutil import socket @@ -172,7 +171,7 @@ def server_setup(tmp_path_factory, name="librarian_server") -> Server: "store_data": { "staging_path": str(staging_directory_sneaker), "store_path": str(store_directory_sneaker), - "report_full_fraction": 0.9, + "report_full_fraction": 0.99, }, "transfer_manager_data": { "local": { diff --git a/tests/server_unit_test/test_admin.py b/tests/server_unit_test/test_admin.py index dbe90c79..279646a8 100644 --- a/tests/server_unit_test/test_admin.py +++ b/tests/server_unit_test/test_admin.py @@ -4,7 +4,6 @@ import shutil -from hera_librarian.deletion import DeletionPolicy from hera_librarian.models.admin import ( AdminAddLibrarianRequest, AdminAddLibrarianResponse, @@ -22,7 +21,7 @@ AdminStoreStateChangeRequest, AdminStoreStateChangeResponse, ) -from hera_librarian.utils import get_checksum_from_path, get_size_from_path +from hera_librarian.utils import get_checksum_from_path def test_add_file(test_client, test_server, garbage_file, test_orm): diff --git a/tests/server_unit_test/test_delete_instance.py b/tests/server_unit_test/test_delete_instance.py index 0505a822..7822aec8 100644 --- a/tests/server_unit_test/test_delete_instance.py +++ b/tests/server_unit_test/test_delete_instance.py @@ -2,23 +2,11 @@ Tests the /users endpoints. """ -import datetime import hashlib -import json -import os import random -import shutil -import sys from pathlib import Path -from subprocess import run -import pytest - -from hera_librarian.authlevel import AuthLevel -from hera_librarian.models.admin import ( - AdminDeleteInstanceRequest, - AdminDeleteInstanceResponse, -) +from hera_librarian.models.admin import AdminDeleteInstanceRequest def test_delete_local_instance(test_server, test_orm, test_client): diff --git a/tests/server_unit_test/test_error.py b/tests/server_unit_test/test_error.py index 69d27ccc..2c91525a 100644 --- a/tests/server_unit_test/test_error.py +++ b/tests/server_unit_test/test_error.py @@ -2,7 +2,6 @@ Tests we can log errors. """ -from hera_librarian.errors import ErrorCategory, ErrorSeverity from hera_librarian.models.errors import ( ErrorClearRequest, ErrorClearResponse, diff --git a/tests/server_unit_test/test_post_to_slack.py b/tests/server_unit_test/test_post_to_slack.py index abf5546a..d92fac60 100644 --- a/tests/server_unit_test/test_post_to_slack.py +++ b/tests/server_unit_test/test_post_to_slack.py @@ -1,6 +1,3 @@ """ Tests the 'post to slack' functionality. """ - -import datetime -import inspect diff --git a/tests/server_unit_test/test_search.py b/tests/server_unit_test/test_search.py index 70e00a9a..b43f4a35 100644 --- a/tests/server_unit_test/test_search.py +++ b/tests/server_unit_test/test_search.py @@ -14,7 +14,6 @@ from hera_librarian.models.search import ( FileSearchFailedResponse, FileSearchRequest, - FileSearchResponse, FileSearchResponses, ) diff --git a/tests/server_unit_test/test_upload.py b/tests/server_unit_test/test_upload.py index 1e7d22a7..88381366 100644 --- a/tests/server_unit_test/test_upload.py +++ b/tests/server_unit_test/test_upload.py @@ -9,7 +9,6 @@ from fastapi.applications import FastAPI from fastapi.testclient import TestClient -from sqlalchemy.orm.session import Session from hera_librarian.models.uploads import ( UploadCompletionRequest, diff --git a/tests/server_unit_test/test_validation.py b/tests/server_unit_test/test_validation.py index bb5fded0..61975c98 100644 --- a/tests/server_unit_test/test_validation.py +++ b/tests/server_unit_test/test_validation.py @@ -3,7 +3,6 @@ """ from hera_librarian.models.validate import ( - FileValidationFailedResponse, FileValidationRequest, FileValidationResponse, FileValidationResponseItem,