import os
import json
from datetime import datetime, timezone
from typing import List, Optional, Any, Tuple, Dict
from sqlalchemy.orm import Session, aliased
from ccat_ops_db import models
from .config.config import ccat_data_transfer_settings
from .utils import (
get_redis_connection,
unique_id,
calculate_checksum,
create_archive,
generate_readable_filename,
safe_join,
deduplicate_by_id,
)
from .setup_celery_app import app, make_celery_task
from .exceptions import (
DatabaseOperationError,
)
from .database import DatabaseConnection
from .logging_utils import get_structured_logger
from .queue_discovery import route_task_by_location
from .operation_types import OperationType
from .health_check import HealthCheck
from .operational_config import OperationalConfig
# Set up global logger
logger = get_structured_logger(__name__)
redis_ = get_redis_connection()
def _utcnow() -> datetime:
"""Timezone-aware now, used for operation start/end timestamps."""
return datetime.now(timezone.utc)
[docs]
class DataTransferPackageOperations(make_celery_task()):
"""Base class for data transfer package operations with error handling.
Bundling runs as a ``BundlingOperation`` (ADR-0003): the ``operation_id`` the
base task handles is the ``BundlingOperation.id`` and status lives on that
row. The ``DataTransferPackage`` artifact is no longer driven from its own
``status`` column. Retry counting falls to the base ``get_retry_count``
(the uniform ``Operation.retry_count``); no per-artifact override remains.
"""
operation_type = "data_transfer_package"
max_retries = 3
[docs]
def mark_in_progress(self, session, operation_id):
"""Flip the BundlingOperation to IN_PROGRESS at task start (ADR-0003).
Terminal ops (COMPLETED/FAILED) are left untouched: the base task runs
this seam BEFORE the body, so resurrecting a settled op here would make
the body's COMPLETED early-return miss and re-bundle a redelivered op.
"""
operation = session.query(models.BundlingOperation).get(operation_id)
if operation and operation.status not in (
models.Status.COMPLETED,
models.Status.FAILED,
):
operation.status = models.Status.IN_PROGRESS
operation.start_time = operation.start_time or _utcnow()
[docs]
def reset_state_on_failure(self, session, operation_id, exc):
"""Reset the BundlingOperation for retry."""
try:
operation = session.query(models.BundlingOperation).get(operation_id)
if operation:
operation.status = models.Status.PENDING
operation.retry_count = (operation.retry_count or 0) + 1
operation.failure_error_message = None
# error_context mirrors failure_error_message: NULLed on reset
# (#117); durable trail kept in OperationFailureEvent.
operation.error_context = None
# Return the bundled raw data packages to the transferring state.
package = self._bundle_package(session, operation)
if package:
for raw_package in package.raw_data_packages:
raw_package.state = models.PackageState.TRANSFERRING
session.commit()
# `data` stays the DataTransferPackage id so the (not-yet-migrated)
# transfer:overview consumer correlates created/reset/failed on the
# same key; the operation id rides alongside under a new key.
redis_.publish(
"transfer:overview",
json.dumps(
{
"type": "transfer_package_reset",
"data": package.id if package else None,
"bundling_operation_id": operation_id,
}
),
)
logger.info(
"bundling operation reset for retry",
operation_id=operation_id,
retry_count=operation.retry_count,
)
except Exception as e:
logger.error(
"failed to reset bundling operation state",
operation_id=operation_id,
error=e,
)
[docs]
def mark_permanent_failure(self, session, operation_id, exc):
"""Mark the BundlingOperation as permanently failed."""
try:
operation = session.query(models.BundlingOperation).get(operation_id)
if operation:
operation.status = models.Status.FAILED
operation.end_time = _utcnow()
operation.failure_error_message = str(exc)
# Cache the latest Tier-1 breadcrumb on the row for the UI (#117).
operation.error_context = self._current_error_context
# Mark the bundled raw data packages as failed.
package = self._bundle_package(session, operation)
if package:
for raw_package in package.raw_data_packages:
raw_package.state = models.PackageState.FAILED
session.commit()
# `data` stays the DataTransferPackage id (consumer-compatible);
# the operation id rides alongside under a new key.
redis_.publish(
"transfer:overview",
json.dumps(
{
"type": "transfer_package_failed",
"data": package.id if package else None,
"bundling_operation_id": operation_id,
}
),
)
logger.error(
"bundling operation marked as permanently failed",
operation_id=operation_id,
error=exc,
)
except Exception as e:
logger.error(
"failed to mark bundling operation as failed",
operation_id=operation_id,
error=e,
)
@staticmethod
def _bundle_package(session, operation):
"""Resolve the DataTransferPackage a BundlingOperation produces."""
if operation.data_transfer_package_id is None:
return None
return session.query(models.DataTransferPackage).get(
operation.data_transfer_package_id
)
[docs]
def get_operation_info(self, args, kwargs):
"""Get additional context for bundling tasks."""
if len(args) < 2:
return {}
return {
"bundling_operation_id": str(args[0]),
"buffer_location_id": str(args[1]),
}
@app.task(
base=DataTransferPackageOperations,
name="ccat:data_transfer_package:create_package",
bind=True,
)
def create_data_transfer_package_task(
self,
bundling_operation_id: int,
buffer_location_id: int,
session: Optional[Session] = None,
) -> bool:
"""
Create a data transfer package by assembling raw data packages from the buffer location.
Parameters
----------
self : celery.Task
The Celery task instance.
bundling_operation_id : int
The ID of the BundlingOperation driving this work. Status lives on this
row; the DataTransferPackage it points at is the bundle assembled.
buffer_location_id : int
The ID of the buffer DataLocation.
session : sqlalchemy.orm.Session, optional
A database session for use in testing environments.
Returns
-------
bool
True if the operation was successful.
"""
if session is None:
with self.session_scope() as session:
return _create_data_transfer_package_internal(
session, bundling_operation_id, buffer_location_id
)
else:
return _create_data_transfer_package_internal(
session, bundling_operation_id, buffer_location_id
)
def _create_data_transfer_package_internal(
session: Session, bundling_operation_id: int, buffer_location_id: int
) -> bool:
"""
Internal function to create a data transfer package.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
bundling_operation_id : int
The ID of the BundlingOperation driving this work (operation status
lives here; it points at the DataTransferPackage to assemble).
buffer_location_id : int
The ID of the buffer DataLocation.
Returns
-------
bool
True if the operation was successful.
"""
logger.debug(
"creating transfer package from buffer",
bundling_operation_id=bundling_operation_id,
location_id=buffer_location_id,
)
try:
# The operation row drives status; the bundle it points at is the
# artifact being assembled.
bundling_operation = session.query(models.BundlingOperation).get(
bundling_operation_id
)
if not bundling_operation:
logger.error(
"bundling operation not found",
bundling_operation_id=bundling_operation_id,
)
return False
if bundling_operation.status == models.Status.COMPLETED:
logger.info(
"bundling operation already completed",
bundling_operation_id=bundling_operation_id,
)
return True
data_transfer_package = session.query(models.DataTransferPackage).get(
bundling_operation.data_transfer_package_id
)
if not data_transfer_package:
logger.error(
"data transfer package not found",
package_id=bundling_operation.data_transfer_package_id,
)
return False
bundling_operation.status = models.Status.IN_PROGRESS
bundling_operation.start_time = bundling_operation.start_time or _utcnow()
session.commit()
buffer_location = session.query(models.DataLocation).get(buffer_location_id)
if not buffer_location:
logger.error("buffer location not found", location_id=buffer_location_id)
return False
# Get all raw data packages for this transfer package
raw_data_packages = data_transfer_package.raw_data_packages
if not raw_data_packages:
logger.warning(
"no raw data packages found for transfer package",
package_id=data_transfer_package.id,
)
return False
# Create archive from raw data packages
logger.debug("assembling transfer package archive", raw_package_count=len(raw_data_packages))
archive_success = _create_transfer_package_archive(
session, data_transfer_package, buffer_location, raw_data_packages
)
if archive_success:
# Operation status lives on the BundlingOperation row (ADR-0003); the
# DataTransferPackage's legacy status column is no longer driven here
# (dropped in #95).
bundling_operation.status = models.Status.COMPLETED
bundling_operation.end_time = _utcnow()
# Create physical copy record at buffer location
physical_copy = models.DataTransferPackagePhysicalCopy(
data_transfer_package=data_transfer_package,
data_location=buffer_location,
checksum=data_transfer_package.checksum,
)
session.add(physical_copy)
# Copy-anchored lineage (ADR-0003): bundling consumes the raw-package
# copies in this buffer and produces the bundle copy.
consumed_copies = _buffer_raw_package_copies(
session, raw_data_packages, buffer_location
)
_record_bundling_lineage(
bundling_operation,
consumed=consumed_copies,
produced=[physical_copy],
)
session.commit()
# `data` is the DataTransferPackage id (the key the current consumer
# reads); the operation id rides alongside, consistent with the
# reset/failed events.
redis_.publish(
"transfer:overview",
json.dumps(
{
"type": "transfer_package_created",
"data": data_transfer_package.id,
"bundling_operation_id": bundling_operation.id,
}
),
)
logger.info(
"transfer package created",
transfer_package_id=data_transfer_package.id,
bundling_operation_id=bundling_operation.id,
raw_package_count=len(raw_data_packages),
size_bytes=data_transfer_package.size,
location_id=buffer_location_id,
)
return True
else:
logger.error(
"failed to create package archive",
package_id=data_transfer_package.id,
)
return False
except Exception as e:
logger.error(
"transfer package creation failed",
bundling_operation_id=bundling_operation_id,
error=e,
)
session.rollback()
return False
def _buffer_raw_package_copies(
session: Session,
raw_data_packages: List[models.RawDataPackage],
buffer_location: models.DataLocation,
) -> List[models.RawDataPackagePhysicalCopy]:
"""Return the buffer-location physical copies of the bundled raw packages.
These are the inputs bundling consumes. A package without a recorded buffer
copy is omitted (lineage is best-effort over what is recorded)."""
package_ids = [p.id for p in raw_data_packages]
if not package_ids:
return []
return (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
models.RawDataPackagePhysicalCopy.raw_data_package_id.in_(package_ids),
models.RawDataPackagePhysicalCopy.data_location_id == buffer_location.id,
)
.all()
)
def _record_bundling_lineage(
operation: models.BundlingOperation,
consumed: List[models.PhysicalCopy],
produced: List[models.PhysicalCopy],
) -> None:
"""Attach copy-anchored consume/produce lineage to a BundlingOperation.
Bundling consumes the raw-package copies and produces the bundle copy
(ADR-0003). Idempotent: a copy already linked is not duplicated."""
for copy in consumed:
if copy is not None and copy not in operation.consumed_copies:
operation.consumed_copies.append(copy)
for copy in produced:
if copy is not None and copy not in operation.produced_copies:
operation.produced_copies.append(copy)
def _create_transfer_package_archive(
session: Session,
data_transfer_package: models.DataTransferPackage,
buffer_location: models.DataLocation,
raw_data_packages: List[models.RawDataPackage],
) -> bool:
"""
Create an archive from raw data packages.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
data_transfer_package : models.DataTransferPackage
The data transfer package to create.
buffer_location : models.DataLocation
The buffer data location.
raw_data_packages : List[models.RawDataPackage]
List of raw data packages to include.
Returns
-------
bool
True if successful, False otherwise.
"""
try:
# Ensure the package directory exists
package_dir = os.path.dirname(
safe_join(buffer_location.path, data_transfer_package.relative_path)
)
os.makedirs(package_dir, exist_ok=True)
archive_path = safe_join(
buffer_location.path, data_transfer_package.relative_path
)
total_size = 0
# Use the existing create_archive function which handles different formats
create_archive(raw_data_packages, archive_path, buffer_location.path)
# Calculate total size and checksum
total_size = sum(package.size for package in raw_data_packages)
checksum = calculate_checksum(archive_path)
if checksum:
data_transfer_package.checksum = checksum
data_transfer_package.size = total_size
return True
else:
logger.error("failed to calculate checksum for archive")
return False
except Exception as e:
logger.error("transfer package archive creation failed", error=e)
return False
[docs]
def get_unpackaged_raw_data_packages_in_buffers(
session: Session,
) -> Dict[models.DataLocation, List[models.RawDataPackage]]:
"""
Get all raw data packages in buffer locations that are not yet part of any data transfer package.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
Dict[models.DataLocation, List[models.RawDataPackage]]
Dictionary mapping buffer locations to lists of unpackaged raw data packages.
"""
logger.debug("scanning buffer locations for unpackaged raw data packages")
# Get all active buffer locations
buffer_locations = (
session.query(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.BUFFER,
models.DataLocation.active == True, # noqa: E712
)
.all()
)
logger.debug("buffer locations found", location_count=len(buffer_locations))
unpackaged_by_buffer = {}
for buffer_location in buffer_locations:
# Find raw data packages in this buffer that are not in any data transfer package
# We check physical copies to see what's actually in this buffer
logger.debug(
"checking buffer location",
location=buffer_location.name,
location_id=buffer_location.id,
)
# Now do the full query with all conditions. The join to physical copies
# can return a RawDataPackage more than once, so dedupe by id in Python
# rather than SELECT DISTINCT (which Postgres cannot apply to the json
# error_context column on this model — see deduplicate_by_id).
unpackaged_packages = deduplicate_by_id(
session.query(models.RawDataPackage)
.join(models.RawDataPackagePhysicalCopy)
.filter(
models.RawDataPackagePhysicalCopy.data_location_id
== buffer_location.id,
models.RawDataPackage.data_transfer_package_id == None, # noqa: E711
models.RawDataPackagePhysicalCopy.status
== models.PhysicalCopyStatus.PRESENT,
# models.RawDataPackagePhysicalCopy.deletion_status
# == models.Status.PENDING,
)
.all()
)
logger.debug("unpackaged raw packages found in buffer", package_count=len(unpackaged_packages), location=buffer_location.name)
if unpackaged_packages:
unpackaged_by_buffer[buffer_location] = unpackaged_packages
logger.debug(
"unpackaged raw packages queued for transfer",
package_count=len(unpackaged_packages),
location=buffer_location.name,
location_id=buffer_location.id,
)
total_packages = sum(len(packages) for packages in unpackaged_by_buffer.values())
logger.info("buffer scan complete", total_unpackaged_packages=total_packages)
return unpackaged_by_buffer
[docs]
def group_packages_for_transfer(
raw_data_packages: List[models.RawDataPackage],
max_package_size: Optional[int] = None,
upper_percentage: Optional[float] = 1.1,
lower_percentage: Optional[float] = 0.9,
session: Optional[Session] = None,
op_config: Optional[OperationalConfig] = None,
) -> List[List[models.RawDataPackage]]:
"""
Group raw data packages into data transfer packages of a specified size range.
Parameters
----------
raw_data_packages : List[models.RawDataPackage]
List of raw data packages to be grouped.
max_package_size : int, optional
Maximum size of a data transfer package in bytes. If None, uses the value from settings.
upper_percentage : float, optional
Upper percentage of max_package_size to start a new package. Default is 1.10 (110%).
lower_percentage : float, optional
Lower percentage of max_package_size to consider a package ready. Default is 0.9 (90%).
Returns
-------
List[List[models.RawDataPackage]]
A list of lists, where each inner list represents a data transfer package
containing raw data packages that, when combined, are within the specified size range.
"""
logger.debug("grouping raw data packages into transfer packages")
if max_package_size is None:
# Convert GB to bytes for internal use
if session is not None and op_config is not None:
size_gb = op_config.get(session, "MAXIMUM_DATA_TRANSFER_PACKAGE_SIZE_GB")
else:
size_gb = ccat_data_transfer_settings.MAXIMUM_DATA_TRANSFER_PACKAGE_SIZE_GB
max_package_size = int(size_gb * 1024 * 1024 * 1024)
sorted_packages = sorted(raw_data_packages, key=lambda x: x.size, reverse=True)
data_transfer_package_list = []
current_package = []
current_size = 0
for package in sorted_packages:
if package.size >= max_package_size * lower_percentage:
# Large package goes alone
data_transfer_package_list.append([package])
logger.debug(
"large raw package assigned to its own transfer package",
raw_package_id=package.id,
size_bytes=package.size,
)
elif (current_size + package.size) > max_package_size * upper_percentage:
# Would exceed upper limit, start new package
if current_package:
data_transfer_package_list.append(current_package)
logger.debug("transfer package reached optimal size, starting new group")
current_package, current_size = [], 0
current_package.append(package)
current_size += package.size
else:
# Add to current package
current_package.append(package)
current_size += package.size
# Log progress
if current_package:
package_ready_percentage = (current_size / max_package_size) * 100
logger.debug(
"package fill status",
fill_pct=round(package_ready_percentage),
)
# Add final package if it meets minimum size
if current_package and current_size >= max_package_size * lower_percentage:
data_transfer_package_list.append(current_package)
logger.debug("final group ready", raw_package_count=len(current_package))
logger.info(
"raw packages grouped",
total_raw_packages=len(raw_data_packages),
transfer_package_count=len(data_transfer_package_list),
)
return data_transfer_package_list
[docs]
def discover_automatic_routes(
session: Session,
) -> List[Tuple[models.Site, models.Site]]:
"""
Discover automatic routes from all source sites to all LTA sites.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
List[Tuple[models.Site, models.Site]]
List of (source_site, lta_site) tuples representing automatic routes.
"""
# Get all sites with source locations
source_sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.SOURCE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
# Get all sites with LTA locations
lta_sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.LONG_TERM_ARCHIVE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
automatic_routes = []
for source_site in source_sites:
for lta_site in lta_sites:
if source_site.id != lta_site.id: # Don't route to self
automatic_routes.append((source_site, lta_site))
logger.debug("automatic routes discovered", route_count=len(automatic_routes))
return automatic_routes
[docs]
def discover_secondary_routes(
session: Session,
) -> List[Tuple[models.Site, models.Site]]:
"""
Discover secondary routes between all LTA sites for data replication.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
List[Tuple[models.Site, models.Site]]
List of (lta_site, lta_site) tuples representing secondary routes.
"""
# Get all sites with LTA locations
lta_sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.LONG_TERM_ARCHIVE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
secondary_routes = []
for origin_lta in lta_sites:
for dest_lta in lta_sites:
if origin_lta.id != dest_lta.id: # Don't route to self
secondary_routes.append((origin_lta, dest_lta))
logger.debug("secondary routes discovered", route_count=len(secondary_routes))
return secondary_routes
[docs]
def find_route_overrides(session: Session) -> List[models.DataTransferRoute]:
"""
Find manual route overrides defined in the DataTransferRoute table.
This is a placeholder function for future implementation of route overrides.
Currently just reports what overrides exist without acting on them.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
List[models.DataTransferRoute]
List of manual route overrides (not yet implemented).
"""
# Get all manual route overrides
route_overrides = session.query(models.DataTransferRoute).all()
if route_overrides:
logger.info(
"route overrides found (not yet implemented)",
route_override_count=len(route_overrides),
)
for route in route_overrides:
logger.debug(
"route override",
origin_site=route.origin_site.short_name,
destination_site=route.destination_site.short_name,
route_type=route.route_type,
transfer_method=route.transfer_method,
)
else:
logger.debug("no route overrides found")
return route_overrides
[docs]
def get_primary_buffer_for_site(
session: Session, site: models.Site
) -> Optional[models.DataLocation]:
"""
Get the primary (highest priority) active buffer for a site.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
site : models.Site
The site to get the buffer for.
Returns
-------
Optional[models.DataLocation]
The primary buffer location, or None if no active buffer exists.
"""
buffer_location = (
session.query(models.DataLocation)
.filter(
models.DataLocation.site_id == site.id,
models.DataLocation.location_type == models.LocationType.BUFFER,
models.DataLocation.active == True, # noqa: E712
)
.order_by(models.DataLocation.priority.asc())
.first()
)
if buffer_location:
logger.debug(
"primary buffer found for site",
site=site.short_name,
location=buffer_location.name,
location_id=buffer_location.id,
)
else:
logger.debug("no active buffer found for site", site=site.short_name)
return buffer_location
[docs]
def get_next_lta_site_round_robin(
session: Session,
source_site: models.Site,
automatic_routes: List[Tuple[models.Site, models.Site]],
) -> Optional[models.Site]:
"""
Get the next LTA site for round-robin distribution from a source site.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
source_site : models.Site
The source site.
automatic_routes : List[Tuple[models.Site, models.Site]]
List of automatic routes.
Returns
-------
Optional[models.Site]
The next LTA site in round-robin order.
"""
# Get all LTA sites that this source site can route to
available_lta_sites = [
lta_site
for src_site, lta_site in automatic_routes
if src_site.id == source_site.id
]
if not available_lta_sites:
logger.warning(
"no lta sites available for source site",
site=source_site.short_name,
)
return None
# Use Redis to track round-robin state
redis_key = f"round_robin:source:{source_site.short_name}"
current_index = redis_.get(redis_key)
if current_index is None:
current_index = 0
else:
current_index = int(current_index)
# Get next LTA site
next_lta_site = available_lta_sites[current_index]
# Update round-robin index
next_index = (current_index + 1) % len(available_lta_sites)
redis_.set(redis_key, next_index)
logger.debug(
"round-robin lta site selected",
source_site=source_site.short_name,
destination_site=next_lta_site.short_name,
index=current_index,
)
return next_lta_site
[docs]
def create_primary_data_transfers(
session: Session,
data_transfer_package: models.DataTransferPackage,
source_buffer: models.DataLocation,
) -> None:
"""
Create primary data transfers for a completed DataTransferPackage using automatic route discovery.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
data_transfer_package : models.DataTransferPackage
The completed data transfer package.
source_buffer : models.DataLocation
The source buffer location.
"""
logger.debug(
"creating primary data transfer",
transfer_package_id=data_transfer_package.id,
)
# Discover automatic routes
automatic_routes = discover_automatic_routes(session)
# Find which site this buffer belongs to
source_site = source_buffer.site
# Get the next LTA site using round-robin
destination_lta_site = get_next_lta_site_round_robin(
session, source_site, automatic_routes
)
if not destination_lta_site:
logger.error(
"no destination lta site available",
site=source_site.short_name,
)
return
# Get the primary buffer for the destination LTA site
destination_buffer = get_primary_buffer_for_site(session, destination_lta_site)
if not destination_buffer:
logger.error(
"no active buffer at destination lta site",
site=destination_lta_site.short_name,
)
return
# Create the primary data transfer. The transfer lifecycle status lives on the
# sibling TransferOperation (created later by transfer_manager find-work);
# DataTransfer.status was dropped in ops-db#95.
data_transfer = models.DataTransfer(
data_transfer_package=data_transfer_package,
origin_location=source_buffer,
destination_location=destination_buffer,
data_transfer_method="bbcp", # Default method, could be configurable
)
session.add(data_transfer)
session.commit()
logger.info(
"primary data transfer created",
transfer_package_id=data_transfer_package.id,
origin_site=source_buffer.site.short_name,
origin_location=source_buffer.name,
destination_site=destination_buffer.site.short_name,
destination_location=destination_buffer.name,
)
[docs]
def create_secondary_data_transfers(session: Session) -> None:
"""
Create secondary data transfers between LTA sites for completed DataTransferPackages.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
"""
logger.info("creating secondary data transfers between LTA sites")
# Discover secondary routes
secondary_routes = discover_secondary_routes(session)
# Find completed packages that need secondary transfers. A package is
# "completed" (bundled) when its BundlingOperation is COMPLETED — the
# lifecycle moved off the dropped DataTransferPackage.status (ops-db#95).
completed_packages = (
session.query(models.DataTransferPackage)
.join(
models.BundlingOperation,
models.BundlingOperation.data_transfer_package_id
== models.DataTransferPackage.id,
)
.filter(models.BundlingOperation.status == models.Status.COMPLETED)
.all()
)
new_transfers_count = 0
for package in completed_packages:
# Find which LTA sites already have this package
existing_locations = set()
for physical_copy in package.physical_copies:
if physical_copy.data_location.location_type == models.LocationType.BUFFER:
if physical_copy.data_location.site.id in [
lta[0].id for lta in secondary_routes
]:
existing_locations.add(physical_copy.data_location.site.id)
# Create transfers to LTA sites that don't have the package yet
for origin_lta, dest_lta in secondary_routes:
if (
origin_lta.id in existing_locations
and dest_lta.id not in existing_locations
):
# Check if transfer already exists
origin_location = aliased(models.DataLocation)
destination_location = aliased(models.DataLocation)
existing_transfer = (
session.query(models.DataTransfer)
.join(
origin_location,
models.DataTransfer.origin_location_id == origin_location.id,
)
.join(
destination_location,
models.DataTransfer.destination_location_id
== destination_location.id,
)
.filter(
models.DataTransfer.data_transfer_package_id == package.id,
origin_location.site_id == origin_lta.id,
destination_location.site_id == dest_lta.id,
)
.first()
)
if not existing_transfer:
origin_buffer = get_primary_buffer_for_site(session, origin_lta)
dest_buffer = get_primary_buffer_for_site(session, dest_lta)
if origin_buffer and dest_buffer:
# Status lives on the sibling TransferOperation created
# later by transfer_manager find-work; DataTransfer.status
# was dropped in ops-db#95.
secondary_transfer = models.DataTransfer(
data_transfer_package=package,
origin_location=origin_buffer,
destination_location=dest_buffer,
data_transfer_method="bbcp",
)
session.add(secondary_transfer)
new_transfers_count += 1
logger.debug(
"secondary transfer created",
transfer_package_id=package.id,
origin_site=origin_lta.short_name,
destination_site=dest_lta.short_name,
)
if new_transfers_count > 0:
session.commit()
logger.info("secondary data transfers created", transfer_count=new_transfers_count)
else:
logger.debug("no new secondary transfers needed")
[docs]
def create_data_transfer_packages_for_buffer(
session: Session,
buffer_location: models.DataLocation,
raw_data_packages: List[models.RawDataPackage],
) -> None:
"""
Create data transfer packages for raw data packages in a buffer location.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
buffer_location : models.DataLocation
The buffer location.
raw_data_packages : List[models.RawDataPackage]
List of raw data packages to process.
"""
logger.debug(
"creating transfer packages for buffer",
package_count=len(raw_data_packages),
location=buffer_location.name,
location_id=buffer_location.id,
)
# Group packages into appropriately sized transfer packages
op_config = OperationalConfig()
transfer_package_groups = group_packages_for_transfer(
raw_data_packages, session=session, op_config=op_config
)
for group in transfer_package_groups:
_create_data_transfer_package_entry(session, group, buffer_location)
session.commit()
logger.info(
"transfer packages created for buffer",
transfer_package_count=len(transfer_package_groups),
location=buffer_location.name,
)
def _create_data_transfer_package_entry(
session: Session,
raw_data_packages: List[models.RawDataPackage],
buffer_location: models.DataLocation,
) -> models.DataTransferPackage:
"""
Create a DataTransferPackage database entry and schedule the assembly task.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
raw_data_packages : List[models.RawDataPackage]
List of raw data packages to include.
buffer_location : models.DataLocation
The buffer location.
Returns
-------
models.DataTransferPackage
The created data transfer package.
"""
logger.debug(
"preparing transfer package entry",
raw_package_count=len(raw_data_packages),
)
# Generate unique package name and path
package_id = unique_id()
readable_name = generate_readable_filename(
raw_data_packages[0], # Use the first package for metadata
package_id[:8], # Use only first 8 chars of hash for brevity
file_type="transfer",
extension="tar", # Always use tar.gz extension for consistency
)
# Create DataTransferPackage entry. Lifecycle status / retry_count live on the
# BundlingOperation row created below (ADR-0003); DataTransferPackage.status
# and retry_count were dropped in ops-db#95.
data_transfer_package = models.DataTransferPackage(
hash_id=package_id,
file_name=readable_name,
origin_location=buffer_location,
size=sum(pkg.size for pkg in raw_data_packages),
checksum="", # Will be calculated during assembly
relative_path=f"data_transfer_packages/{readable_name}",
)
# Associate raw data packages with the transfer package
for raw_package in raw_data_packages:
raw_package.data_transfer_package = data_transfer_package
raw_package.state = models.PackageState.TRANSFERRING
session.add(data_transfer_package)
session.flush() # Get the ID
# Record the bundling operation as PENDING work. Status lives on this row
# (ADR-0003); the poll step flips it to SCHEDULED and dispatches the task.
bundling_operation = models.BundlingOperation(
data_transfer_package_id=data_transfer_package.id,
status=models.Status.PENDING,
)
session.add(bundling_operation)
session.commit()
logger.info(
"created data transfer package",
package_id=data_transfer_package.id,
bundling_operation_id=bundling_operation.id,
package_name=readable_name,
raw_package_count=len(raw_data_packages),
)
return data_transfer_package
[docs]
def schedule_data_transfer_package_creation(session: Session) -> None:
"""Dispatch assembly tasks for PENDING bundling operations (find-work loop).
The standard poll→SCHEDULED→dispatch loop (ADR-0003): each ``PENDING``
``BundlingOperation`` is flipped to ``SCHEDULED`` and its assembly task is
queued, keyed by the operation id. Operations already past ``PENDING`` are
skipped, so re-running the poll never queues a second task for the same work.
The status flip is committed per operation before dispatch so a crash
mid-loop cannot re-dispatch an already-scheduled operation.
"""
pending_operations = (
session.query(models.BundlingOperation)
.filter(models.BundlingOperation.status == models.Status.PENDING)
.all()
)
if not pending_operations:
logger.info("no pending bundling operations found")
return
for operation in pending_operations:
package = session.query(models.DataTransferPackage).get(
operation.data_transfer_package_id
)
if not package:
logger.warning(
"bundling operation has no data transfer package",
bundling_operation_id=operation.id,
)
continue
buffer_location = package.origin_location
operation.status = models.Status.SCHEDULED
session.commit()
queue_name = route_task_by_location(
OperationType.DATA_TRANSFER_PACKAGE_CREATION, buffer_location
)
create_data_transfer_package_task.apply_async(
args=[operation.id, buffer_location.id], queue=queue_name
)
logger.info(
"bundling operation scheduled",
bundling_operation_id=operation.id,
transfer_package_id=package.id,
queue=queue_name,
)
logger.info(
"bundling operations scheduled",
operation_count=len(pending_operations),
)
[docs]
def create_data_transfer_packages(
verbose: bool = False, session: Session = None
) -> None:
"""
Scan all buffer locations and create data transfer packages for unpackaged raw data packages.
This function manages the process of creating data transfer packages by:
1. Finding all buffer locations with unpackaged raw data packages.
2. Grouping packages into appropriately sized transfer packages.
3. Creating DataTransferPackage entries in the database.
4. Creating primary and secondary data transfers using automatic route discovery.
5. Scheduling Celery tasks to handle the package assembly.
Args:
verbose (bool, optional): If True, sets logging level to DEBUG. Defaults to False.
session (Session, optional): Database session for testing. If None, creates new session.
Raises:
ConfigurationError: If no active buffer is found for a site.
DatabaseOperationError: If there's an error during database operations.
"""
if session is None:
db = DatabaseConnection()
session, _ = db.get_connection()
try:
# Report any route overrides (placeholder for future implementation)
find_route_overrides(session)
# Get all unpackaged raw data packages by buffer location
unpackaged_by_buffer = get_unpackaged_raw_data_packages_in_buffers(session)
if unpackaged_by_buffer:
logger.info(
"processing buffers with unpackaged data",
buffer_count=len(unpackaged_by_buffer),
)
# Create transfer packages for each buffer location
for buffer_location, raw_data_packages in unpackaged_by_buffer.items():
try:
create_data_transfer_packages_for_buffer(
session, buffer_location, raw_data_packages
)
except Exception as e:
logger.error(
"error processing buffer location",
location=buffer_location.name,
error=e,
)
session.rollback()
raise DatabaseOperationError(
f"Failed to create data transfer packages for {buffer_location.name}: {str(e)}"
) from e
schedule_data_transfer_package_creation(session)
else:
logger.info("no unpackaged raw data packages found in any buffer")
# Always create primary data transfers for completed packages (regardless of new packages)
_create_primary_transfers_for_completed_packages(session)
# Always create secondary data transfers between LTA sites (regardless of new packages)
create_secondary_data_transfers(session)
logger.info("completed data transfer package creation and transfer scheduling")
except Exception as e:
logger.exception("an error occurred while creating data transfer packages")
raise RuntimeError("Failed to create data transfer packages") from e
def _create_primary_transfers_for_completed_packages(session: Session) -> None:
"""
Create primary data transfers for all completed DataTransferPackages that don't have transfers yet.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
"""
logger.debug("checking for completed packages needing primary transfers")
# Find completed packages without any transfers. A package is "completed"
# (bundled) when its BundlingOperation is COMPLETED — the lifecycle moved off
# the dropped DataTransferPackage.status (ops-db#95).
completed_packages = (
session.query(models.DataTransferPackage)
.join(
models.BundlingOperation,
models.BundlingOperation.data_transfer_package_id
== models.DataTransferPackage.id,
)
.filter(models.BundlingOperation.status == models.Status.COMPLETED)
.filter(~models.DataTransferPackage.data_transfers.any())
.all()
)
if not completed_packages:
logger.debug("no completed packages without transfers found")
return
logger.info(
"completed packages need primary transfers",
package_count=len(completed_packages),
)
for package in completed_packages:
# Find the source buffer location for this package
source_buffer = None
for physical_copy in package.physical_copies:
if physical_copy.data_location.location_type == models.LocationType.BUFFER:
# Check if this buffer belongs to a source site
source_locations = (
session.query(models.DataLocation)
.filter(
models.DataLocation.site_id
== physical_copy.data_location.site_id,
models.DataLocation.location_type == models.LocationType.SOURCE,
models.DataLocation.active == True, # noqa: E712
)
.first()
)
if source_locations:
source_buffer = physical_copy.data_location
break
if source_buffer:
try:
create_primary_data_transfers(session, package, source_buffer)
except Exception as e:
logger.error(
"primary transfer creation failed",
transfer_package_id=package.id,
error=e,
)
# Continue with other packages
continue
else:
logger.warning("no source buffer found for completed package", transfer_package_id=package.id)
[docs]
def data_transfer_package_manager_service(verbose: bool = False) -> None:
"""
Main service function for the data transfer package manager.
This service continuously scans buffer locations for unpackaged raw data packages
and creates transfer packages with automatic route discovery.
Args:
verbose (bool): If True, sets logging level to DEBUG. Default is False.
"""
db = DatabaseConnection()
session, _ = db.get_connection()
# Initialize health check
health_check = HealthCheck(
service_type="data_transfer_package",
service_name="data_transfer_package_manager",
)
health_check.start()
try:
while True: # Main service loop
try:
logger.debug("starting transfer package manager cycle")
create_data_transfer_packages(verbose=verbose, session=session)
# Sleep for the configured interval
import time
time.sleep(ccat_data_transfer_settings.PACKAGE_MANAGER_SLEEP_TIME)
except Exception as e:
logger.error("service_loop_error", error=str(e))
import time
time.sleep(10) # Wait before retry
finally:
health_check.stop()
session.close()
# Legacy functions for compatibility (can be removed after full migration)
[docs]
def create_primary_data_transfer_packages(
verbose: bool = False, session: Session = None
) -> None:
"""
Legacy wrapper function for backward compatibility.
This function redirects to the new create_data_transfer_packages function.
Can be removed once all callers are updated.
"""
logger.warning(
"deprecated function called",
deprecated="create_primary_data_transfer_packages",
use_instead="create_data_transfer_packages",
)
create_data_transfer_packages(verbose=verbose, session=session)
def _create_primary_data_transfer_packages_internal(session: Session) -> None:
"""
Legacy internal function for backward compatibility.
This function redirects to the new implementation.
Can be removed once all callers are updated.
"""
logger.warning(
"deprecated function called",
deprecated="_create_primary_data_transfer_packages_internal",
use_instead="create_data_transfer_packages",
)
create_data_transfer_packages(session=session)
# Additional utility functions for the new architecture
[docs]
def get_sites_with_buffer_locations(session: Session) -> List[models.Site]:
"""
Get all sites that have active BUFFER data locations.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
List[models.Site]
List of sites with buffer locations.
"""
sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.BUFFER,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
logger.debug("sites with active buffers found", site_count=len(sites))
return sites
[docs]
def validate_site_configuration(session: Session) -> bool:
"""
Validate that all sites have proper configuration for data transfer.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
bool
True if configuration is valid, False otherwise.
"""
logger.info("validating site configuration for data transfer")
# Check that all source sites have buffers
source_sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.SOURCE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
for site in source_sites:
buffer = get_primary_buffer_for_site(session, site)
if not buffer:
logger.error("source site missing active buffer", site=site.short_name)
return False
# Check that all LTA sites have buffers
lta_sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.LONG_TERM_ARCHIVE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
for site in lta_sites:
buffer = get_primary_buffer_for_site(session, site)
if not buffer:
logger.error("lta site missing active buffer", site=site.short_name)
return False
logger.info("site configuration validation passed")
return True
[docs]
def get_transfer_statistics(session: Session) -> Dict[str, Any]:
"""
Get statistics about the current transfer system state.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
Dict[str, Any]
Dictionary containing transfer statistics.
"""
stats = {}
# Count unpackaged raw data packages
unpackaged_by_buffer = get_unpackaged_raw_data_packages_in_buffers(session)
stats["unpackaged_packages_by_buffer"] = {
loc.name: len(packages) for loc, packages in unpackaged_by_buffer.items()
}
stats["total_unpackaged_packages"] = sum(
len(packages) for packages in unpackaged_by_buffer.values()
)
# Count pending transfer packages. Lifecycle status moved onto the sibling
# operation rows (ADR-0003); DataTransferPackage.status / DataTransfer.status
# were dropped in ops-db#95, so the counts read the BundlingOperation /
# TransferOperation rows instead.
pending_packages = (
session.query(models.BundlingOperation)
.filter_by(status=models.Status.PENDING)
.count()
)
stats["pending_transfer_packages"] = pending_packages
# Count pending transfers
pending_transfers = (
session.query(models.TransferOperation)
.filter_by(status=models.Status.PENDING)
.count()
)
stats["pending_transfers"] = pending_transfers
# Automatic routes
automatic_routes = discover_automatic_routes(session)
stats["automatic_routes_count"] = len(automatic_routes)
secondary_routes = discover_secondary_routes(session)
stats["secondary_routes_count"] = len(secondary_routes)
# Route overrides
route_overrides = find_route_overrides(session)
stats["route_overrides_count"] = len(route_overrides)
logger.debug("transfer statistics computed", **stats)
return stats
