from typing import List, Optional
import datetime
import subprocess
from sqlalchemy.orm import Session, selectinload
from sqlalchemy import and_
import os
import shutil
import tempfile
from .database import DatabaseConnection
from .setup_celery_app import app, make_celery_task
from .decorators import track_metrics
from .utils import (
get_s3_client,
get_s3_key_for_package,
resolve_s3_method,
resolve_coscine_config,
)
from ccat_ops_db import models
from .logging_utils import get_structured_logger
from .queue_discovery import route_task_by_location
from .operation_types import OperationType
logger = get_structured_logger(__name__)
[docs]
class StagingTask(make_celery_task()):
"""Base class for staging tasks.
Staging is on the uniform Operation row (ADR-0003, #154): the task operates
on a single ``StagingOperation`` (one per package), so retry / recovery /
circuit-breaker all key on the operation, and a failing package never resets
its siblings. The identity is the frozen breadcrumb string ``staging`` (==
``OperationKind.STAGING``).
No ``get_retry_count`` override: the legacy StagingJob counter is gone, so
the base uniform reader (``Operation.retry_count``) is used (#75/#154).
"""
operation_type = "staging"
[docs]
def __init__(self):
super().__init__()
self.max_retries = 3
[docs]
def mark_in_progress(self, session, staging_operation_id):
"""Mark the StagingOperation IN_PROGRESS at task start (ADR-0003 seam).
Overrides the base no-op: this is the single place SCHEDULED (queued)
becomes IN_PROGRESS (running) for staging, giving the DB an honest
running signal that stuck-detection and the UI read.
"""
staging_operation = session.get(
models.StagingOperation, staging_operation_id
)
if staging_operation:
staging_operation.status = models.Status.IN_PROGRESS
[docs]
def reset_state_on_failure(self, session, staging_operation_id, exc):
"""Reset a StagingOperation to PENDING for retry, bumping its own
uniform ``retry_count`` — only this package retries, not the job.
Re-syncs the owning legacy StagingJob.status afterwards: the internal
marked the job mirror FAILED when this package failed, so on a retryable
failure the mirror must leave FAILED (re-derived from the children) for
the retry window — otherwise the API/explorer would read a stuck FAILED
while the package is actually retrying.
"""
staging_operation = session.get(
models.StagingOperation, staging_operation_id
)
if staging_operation:
staging_operation.status = models.Status.PENDING
staging_operation.failure_error_message = str(exc)
staging_operation.retry_count += 1
_sync_owning_job_status(session, staging_operation)
session.commit()
[docs]
def mark_permanent_failure(self, session, staging_operation_id, exc):
"""Mark a StagingOperation permanently FAILED."""
staging_operation = session.get(
models.StagingOperation, staging_operation_id
)
if staging_operation:
staging_operation.status = models.Status.FAILED
staging_operation.failure_error_message = str(exc)
session.commit()
[docs]
def get_operation_info(self, args, kwargs):
"""Get information about the operation for logging."""
staging_operation_id = args[0] if args else kwargs.get("staging_operation_id")
return {
"staging_operation_id": staging_operation_id,
"operation_type": "staging",
}
@app.task(
base=StagingTask,
name="ccat:data_transfer:staging",
bind=True,
)
def stage_data_task(
self,
staging_operation_id: int,
session: Optional[Session] = None,
) -> None:
"""Celery task to stage a single package via its StagingOperation."""
if session is None:
with self.session_scope() as session:
return _stage_operation_internal(session, staging_operation_id)
return _stage_operation_internal(session, staging_operation_id)
@track_metrics(
operation_type="staging",
additional_tags={
"transfer_method": "s3",
},
)
def _stage_operation_internal(session: Session, staging_operation_id: int) -> None:
"""Stage the single package behind one ``StagingOperation`` (#154).
Each package is an independent operation, so this drives exactly one of them
to its own terminal state and re-raises on failure — the base task hook then
resets/retries (or permanently fails) just this operation, leaving the
siblings untouched. On success the operation is COMPLETED here; the derived
group/job status is recomputed from all children so the legacy
``StagingJob.status`` (read by the staging API/explorer) stays in step.
"""
staging_operation = session.get(
models.StagingOperation, staging_operation_id
)
if not staging_operation:
logger.error(
"staging operation not found",
staging_operation_id=staging_operation_id,
)
return
try:
start_time = datetime.datetime.now()
_stage_single_package(session, staging_operation)
staging_operation.status = models.Status.COMPLETED
staging_operation.start_time = start_time
staging_operation.end_time = datetime.datetime.now()
_sync_owning_job_status(session, staging_operation)
session.commit()
logger.info(
"staging operation completed",
staging_operation_id=staging_operation_id,
raw_package_id=staging_operation.raw_data_package_id,
)
except Exception as e:
logger.error(
"staging operation failed",
staging_operation_id=staging_operation_id,
raw_package_id=staging_operation.raw_data_package_id,
error=e,
)
# Mark this package failed and re-sync the derived job status before the
# base failure hook runs its own reset/retry on this operation alone.
staging_operation.status = models.Status.FAILED
staging_operation.failure_error_message = str(e)
_sync_owning_job_status(session, staging_operation)
session.commit()
raise
def _stage_single_package(
session: Session, staging_operation: models.StagingOperation
) -> bool:
"""Download + unpack the one package behind a StagingOperation.
The per-package work formerly inlined in the all-or-nothing job loop, lifted
into a single-operation unit so one package can succeed or fail on its own.
Returns True on success; raises on any failure (the caller maps the raise to
the operation's FAILED state).
"""
raw_data_package = session.get(
models.RawDataPackage, staging_operation.raw_data_package_id
)
if not raw_data_package:
raise ValueError(
f"Raw data package {staging_operation.raw_data_package_id} not found"
)
origin_location = session.get(
models.DataLocation, staging_operation.origin_location_id
)
destination_location = session.get(
models.DataLocation, staging_operation.destination_location_id
)
# Skip if already staged at destination (idempotent re-run).
if _check_existing_copies(
session, destination_location, raw_data_package.id
):
logger.debug(
"package already staged at destination, skipping",
staging_operation_id=staging_operation.id,
raw_package_id=raw_data_package.id,
)
return True
physical_copy = _get_physical_copy(
session, origin_location.id, raw_data_package.id
)
if not physical_copy:
raise ValueError(
f"No physical copy found in origin location {origin_location.id} "
f"for package {raw_data_package.id}"
)
final_destination_path = _construct_destination_path(
destination_location, raw_data_package
)
# Download to a temporary location first, then unpack to the final dest.
# Use just the filename to avoid double raw_data_packages segments.
package_filename = os.path.basename(raw_data_package.relative_path)
temp_package_path = os.path.join(
destination_location.path,
"raw_data_packages",
package_filename,
)
result = _execute_polymorphic_copy(physical_copy, temp_package_path)
if not result:
raise ValueError(f"Failed to stage package {raw_data_package.id}")
if temp_package_path.endswith(".tar.gz"):
_unpack_file(temp_package_path, final_destination_path)
_check_raw_data_files(raw_data_package, final_destination_path)
_create_raw_data_file_physical_copies(
session, destination_location, raw_data_package, True
)
_mark_package_as_staged_and_cleanup(
session, destination_location, raw_data_package
)
return True
def _sync_owning_job_status(
session: Session, staging_operation: models.StagingOperation
) -> None:
"""Re-derive the legacy StagingJob.status from this operation's group.
The group is the source of truth (status derived from its children); the
legacy job mirror is what the staging API/explorer read, so it must follow.
A NULL group (operation created outside a dispatch) is a no-op.
"""
group = staging_operation.operation_group
if group is None:
return
staging_job = _find_job_for_group(session, group)
if staging_job is not None:
_sync_job_status_from_group(session, staging_job, group)
def _get_physical_copy(
session: Session, data_location_id: int, raw_data_package_id: int
) -> Optional[models.PhysicalCopy]:
"""Get the physical copy for a raw data package in a specific data location."""
logger.debug(
"looking for physical copy",
location_id=data_location_id,
raw_package_id=raw_data_package_id,
)
# First, let's see what physical copies exist for this package
all_copies = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
models.RawDataPackagePhysicalCopy.raw_data_package_id == raw_data_package_id
)
.all()
)
logger.debug(
"physical copies found for package",
raw_package_id=raw_data_package_id,
copy_count=len(all_copies),
)
for copy in all_copies:
logger.debug(
"physical copy detail",
physical_copy_id=copy.id,
location_id=copy.data_location_id,
status=str(copy.status),
)
# Now look for the specific one we need
physical_copy = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
and_(
models.RawDataPackagePhysicalCopy.data_location_id == data_location_id,
models.RawDataPackagePhysicalCopy.raw_data_package_id
== raw_data_package_id,
models.RawDataPackagePhysicalCopy.status
== models.PhysicalCopyStatus.PRESENT,
)
)
.first()
)
if physical_copy:
logger.debug(
"matching physical copy found",
physical_copy_id=physical_copy.id,
status=str(physical_copy.status),
)
else:
logger.debug(
"no matching physical copy found",
location_id=data_location_id,
raw_package_id=raw_data_package_id,
)
return physical_copy
def _construct_destination_path(
destination_location: models.DataLocation, raw_data_package: models.RawDataPackage
) -> str:
"""Construct the destination path for staged data."""
logger.debug(
"constructing destination path",
destination_location_type=type(destination_location).__name__,
destination_location=destination_location.name,
rel_path=raw_data_package.relative_path,
)
if isinstance(destination_location, models.DiskDataLocation):
# For staging, we want to unpack directly into the destination location root
# This will recreate the hierarchical structure from the source location
# The tar file contains the full relative paths (e.g., CHAI/LFA/filename.fits)
# so unpacking into the root will create the correct hierarchy
path = destination_location.path
logger.debug(
"destination path resolved",
destination_location_path=destination_location.path,
path=path,
)
return path
else:
logger.debug(
"unsupported destination location type",
destination_location_type=type(destination_location).__name__,
)
raise ValueError(
f"Unsupported destination location type: {type(destination_location)}"
)
def _execute_polymorphic_copy(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the copy operation based on the storage types of source and destination."""
try:
# Get the source location from the physical copy
source_location = physical_copy.data_location
# For now, we'll assume the destination is a disk location for staging
# In the future, this could be extended to support staging to other storage types
if isinstance(source_location, models.S3DataLocation):
return _execute_s3_download(physical_copy, destination_path)
# elif isinstance(source_location, models.DiskDataLocation):
# # Check if this is a local or remote copy
# if source_location.host and source_location.host != "localhost":
# return _execute_remote_copy(physical_copy, destination_path)
# else:
# return _execute_local_copy(physical_copy, destination_path)
# elif isinstance(source_location, models.TapeDataLocation):
# return _execute_tape_to_disk_copy(physical_copy, destination_path)
else:
raise ValueError(
f"Unsupported source location type: {type(source_location)}"
)
except Exception as e:
logger.error("polymorphic copy failed", error=e)
raise
def _execute_tape_to_disk_copy(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute copy from tape to disk."""
try:
# This is a placeholder for tape operations
# In a real implementation, this would use tape library commands
logger.warning("tape to disk copy not yet implemented")
raise NotImplementedError("Tape to disk copy not yet implemented")
except Exception as e:
logger.error("tape to disk copy failed", error=e)
raise
def _execute_remote_copy(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the remote copy using the physical copy's full path."""
try:
source_location = physical_copy.data_location
if not isinstance(source_location, models.DiskDataLocation):
raise ValueError(f"Expected DiskDataLocation, got {type(source_location)}")
remote_host = source_location.host
remote_user = source_location.user or "ccat"
remote_path = f"{remote_user}@{remote_host}:{physical_copy.full_path}"
# Create destination directory if it doesn't exist
os.makedirs(os.path.dirname(destination_path), exist_ok=True)
subprocess.run(["scp", remote_path, destination_path], check=True)
return True
except Exception as e:
logger.error("remote copy failed", error=e)
raise
def _execute_local_copy(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the local copy using the physical copy's full path."""
try:
# Create destination directory if it doesn't exist
os.makedirs(os.path.dirname(destination_path), exist_ok=True)
shutil.copy(physical_copy.full_path, destination_path)
return True
except Exception as e:
logger.error("local copy failed", error=e)
raise
def _execute_s3_download(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the S3 download using the physical copy's location information.
Routes to either boto3 or Coscine download based on the location's
access_method (falling back to the deployment-wide S3_METHOD).
"""
method = resolve_s3_method(physical_copy.data_location)
if method == "boto3":
return _execute_boto3_s3_download(physical_copy, destination_path)
elif method == "coscine":
return _execute_coscine_s3_download(physical_copy, destination_path)
else:
raise ValueError(f"Unsupported s3 access method: {method}")
def _execute_boto3_s3_download(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the S3 download using boto3."""
try:
source_location = physical_copy.data_location
if not isinstance(source_location, models.S3DataLocation):
raise ValueError(f"Expected S3DataLocation, got {type(source_location)}")
# Use the source location's own endpoint/credentials so a download honours
# the per-location connection config (falls back to defaults internally).
site_name = source_location.site.short_name if source_location.site else None
s3_client = get_s3_client(source_location, site_name)
# Get bucket name and object key from the S3 location
bucket_name = source_location.bucket_name
# Use the shared function to construct the S3 key consistently
object_key = get_s3_key_for_package(
source_location, physical_copy.raw_data_package
)
logger.debug(
"s3 download details",
bucket_name=bucket_name,
object_key=object_key,
destination_path=destination_path,
)
# Create the destination directory if it doesn't exist
os.makedirs(os.path.dirname(destination_path), exist_ok=True)
# Download the file
s3_client.download_file(bucket_name, object_key, destination_path)
return True
except Exception as e:
logger.error("s3 download failed", error=e)
raise
def _execute_coscine_s3_download(
physical_copy: models.RawDataPackagePhysicalCopy,
destination_path: str,
) -> bool:
"""Execute the S3 download using Coscine API.
Parameters
----------
physical_copy : models.RawDataPackagePhysicalCopy
The physical copy record containing source location information
destination_path : str
Local path where the file should be downloaded
Returns
-------
bool
True if download was successful
Raises
------
ValueError
If Coscine configuration is not set or download fails
"""
import coscine
try:
source_location = physical_copy.data_location
if not isinstance(source_location, models.S3DataLocation):
raise ValueError(f"Expected S3DataLocation, got {type(source_location)}")
# Resolve Coscine config per-location (project/resource from the location,
# token from the env), falling back to the deployment-wide settings.
site_name = source_location.site.short_name if source_location.site else None
coscine_api_token, coscine_project, coscine_resource = resolve_coscine_config(
source_location, site_name
)
if (
not coscine_api_token
or coscine_api_token == "none"
or not coscine_project
or coscine_project == "none"
or not coscine_resource
or coscine_resource == "none"
):
raise ValueError(
"COSCINE configuration is not set. Set COSCINE_PROJECT / "
"COSCINE_RESOURCE on the data location (or in settings) and the "
"COSCINE_API_TOKEN in the env."
)
# Construct the object key (same as upload)
object_key = get_s3_key_for_package(
source_location, physical_copy.raw_data_package
)
# Sanitize key to match upload convention (archive_manager.py:1097)
object_key = object_key.replace("@", "At_").replace(":", "-")
logger.debug(
"coscine download details",
project=coscine_project,
resource=coscine_resource,
object_key=object_key,
destination_path=destination_path,
)
# Initialize Coscine client
client = coscine.ApiClient(coscine_api_token)
project = client.project(coscine_project)
resource = project.resource(coscine_resource)
# Create the destination directory if it doesn't exist
os.makedirs(os.path.dirname(destination_path), exist_ok=True)
# Download the specific file from the resource
# Note: Depending on the Coscine API version, this might need adjustment
# If there's a file-specific download method, use that instead
try:
# Try to download a specific file if the API supports it
file_obj = resource.file(object_key)
file_obj.download(destination_path)
except AttributeError:
# Fallback: Download entire resource to temp location and extract the file
with tempfile.TemporaryDirectory() as temp_dir:
resource.download(path=temp_dir)
# Find and copy the specific file to destination
source_file = os.path.join(temp_dir, object_key)
if not os.path.exists(source_file):
raise ValueError(
f"File {object_key} not found in downloaded resource"
)
shutil.copy(source_file, destination_path)
logger.info(
"coscine download succeeded",
object_key=object_key,
destination_path=destination_path,
)
return True
except Exception as e:
logger.error(
"coscine download failed",
destination_path=destination_path,
error=e,
)
raise ValueError(f"COSCINE S3 download failed: {str(e)}")
def _check_raw_data_files(
raw_data_package: models.RawDataPackage, destination_path: str
) -> None:
"""Check the completeness of the raw data files.
This function checks if all files from the raw data package exist in the destination
directory by comparing their full relative paths. The files should be unpacked with
their complete hierarchical structure (e.g., CHAI/LFA/filename.fits).
Parameters
----------
raw_data_package : models.RawDataPackage
The raw data package containing the list of expected files
destination_path : str
The root directory where files should be found
Raises
------
ValueError
If any expected files are missing or if there's an error scanning the directory
"""
try:
# Get the expected file paths from the raw data package
expected_files = {
file.relative_path for file in raw_data_package.raw_data_files
}
if not expected_files:
raise ValueError("No files listed in raw data package")
# For staging, the destination_path is the root directory where files are unpacked
# Files should be found at their full relative paths (e.g., CHAI/LFA/filename.fits)
search_path = destination_path
# Check for missing files by testing each expected path
missing_files = []
for expected_path in expected_files:
full_path = os.path.join(search_path, expected_path)
if not os.path.exists(full_path):
missing_files.append(expected_path)
else:
logger.debug("file verified", rel_path=expected_path)
if missing_files:
raise ValueError(
f"Missing files in destination directory: {sorted(missing_files)}"
)
logger.info(
"raw data files verified",
raw_package_id=raw_data_package.id,
file_count=len(expected_files),
search_path=search_path,
)
except Exception as e:
if isinstance(e, ValueError):
raise
raise ValueError(f"Error checking raw data files: {str(e)}")
def _unpack_file(file_path: str, destination_dir: str) -> None:
"""Unpack the file directly into the destination directory, preserving the hierarchical structure.
Parameters
----------
file_path : str
Path to the tar.gz file to unpack
destination_dir : str
Directory to unpack the files into
Raises
------
ValueError
If unpacking fails or if the file is not a tar.gz
"""
if not file_path.endswith(".tar.gz"):
raise ValueError(f"Expected .tar.gz file, got: {file_path}")
try:
# Create the destination directory if it doesn't exist
os.makedirs(destination_dir, exist_ok=True)
logger.debug("unpacking archive", path=file_path, destination_dir=destination_dir)
# First, let's see what's in the tar.gz file to understand the structure
list_result = subprocess.run(
["tar", "-tzf", file_path],
capture_output=True,
text=True,
check=True,
)
if list_result.returncode != 0:
raise ValueError(
f"Failed to list contents of {file_path}: {list_result.stderr}"
)
tar_contents = list_result.stdout.strip().split("\n")
logger.debug("tar contents listed", path=file_path, entry_count=len(tar_contents))
if not tar_contents or not tar_contents[0]:
raise ValueError(f"Tar file {file_path} appears to be empty")
# Extract the tar file directly into the destination directory
# The tar file contains the full relative paths (e.g., CHAI/LFA/filename.fits)
# so unpacking will recreate the correct hierarchical structure
result = subprocess.run(
[
"tar",
"-xzf",
file_path,
"-C",
destination_dir,
"--overwrite",
],
capture_output=True,
text=True,
check=True,
)
if result.returncode != 0:
raise ValueError(f"Failed to unpack {file_path}: {result.stderr}")
logger.info("archive unpacked", path=file_path, destination_dir=destination_dir)
except subprocess.CalledProcessError as e:
raise ValueError(f"Failed to unpack {file_path}: {e.stderr}")
except Exception as e:
raise ValueError(f"Error unpacking {file_path}: {str(e)}")
def _get_staging_job(
session: Session, staging_job_id: int
) -> Optional[models.StagingJob]:
"""Get a staging job by ID."""
return session.query(models.StagingJob).get(staging_job_id)
def _create_raw_data_file_physical_copies(
session: Session,
destination_location: models.DataLocation,
raw_data_package: models.RawDataPackage,
success: bool,
) -> None:
"""Create physical copy records for all RawDataFiles in the package.
After unpacking a RawDataPackage, we create physical copy records for each
RawDataFile so they can be tracked and deleted when no longer needed. Keyed
on the destination location directly (the operation owns it), not the legacy
StagingJob (#154).
"""
if not success:
return
# Create physical copies for each RawDataFile
for raw_data_file in raw_data_package.raw_data_files:
# Check if physical copy already exists
existing_copy = (
session.query(models.RawDataFilePhysicalCopy)
.filter(
and_(
models.RawDataFilePhysicalCopy.raw_data_file_id == raw_data_file.id,
models.RawDataFilePhysicalCopy.data_location_id
== destination_location.id,
models.RawDataFilePhysicalCopy.status
== models.PhysicalCopyStatus.PRESENT,
)
)
.first()
)
if not existing_copy:
physical_copy = models.RawDataFilePhysicalCopy(
raw_data_file_id=raw_data_file.id,
data_location_id=destination_location.id,
status=models.PhysicalCopyStatus.PRESENT,
created_at=datetime.datetime.now(datetime.timezone.utc),
)
session.add(physical_copy)
session.commit()
logger.info(
"raw data file physical copies created",
raw_package_id=raw_data_package.id,
file_count=len(raw_data_package.raw_data_files),
)
def _mark_package_as_staged_and_cleanup(
session: Session,
destination_location: models.DataLocation,
raw_data_package: models.RawDataPackage,
) -> None:
"""Mark RawDataPackage as STAGED and delete the physical package file.
After unpacking and creating RawDataFile physical copies, we mark the package
as STAGED and remove the physical package file to save space. Keyed on the
destination location directly (#154).
"""
# Find or create the RawDataPackage physical copy record
package_physical_copy = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
and_(
models.RawDataPackagePhysicalCopy.raw_data_package_id
== raw_data_package.id,
models.RawDataPackagePhysicalCopy.data_location_id
== destination_location.id,
)
)
.first()
)
if not package_physical_copy:
# Create new record if it doesn't exist
package_physical_copy = models.RawDataPackagePhysicalCopy(
raw_data_package_id=raw_data_package.id,
data_location_id=destination_location.id,
status=models.PhysicalCopyStatus.STAGED,
created_at=datetime.datetime.now(datetime.timezone.utc),
)
session.add(package_physical_copy)
else:
# Update existing record to STAGED
package_physical_copy.status = models.PhysicalCopyStatus.STAGED
# Delete the physical package file
# For staging, the package file is stored in a temporary location
# We need to find where the original package file was downloaded
package_file_path = None
# Look for the package file in the destination location's raw_data_packages directory
if isinstance(destination_location, models.DiskDataLocation):
# Use just the filename to match the temporary path construction
package_filename = os.path.basename(raw_data_package.relative_path)
package_file_path = os.path.join(
destination_location.path,
"raw_data_packages",
package_filename,
)
if package_file_path and os.path.exists(package_file_path):
try:
os.remove(package_file_path)
logger.info(
"physical package file deleted",
raw_package_id=raw_data_package.id,
path=package_file_path,
)
except OSError as e:
logger.warning(
"failed to delete physical package file",
raw_package_id=raw_data_package.id,
path=package_file_path,
error=e,
)
else:
logger.debug(
"package file not found at expected location",
raw_package_id=raw_data_package.id,
path=package_file_path,
)
session.commit()
logger.info(
"package marked as staged",
raw_package_id=raw_data_package.id,
location=destination_location.name,
)
def _get_pending_staging_jobs(session: Session) -> List[models.StagingJob]:
"""Get all pending staging jobs."""
return (
session.query(models.StagingJob)
.filter(models.StagingJob.status == models.Status.PENDING)
.all()
)
def _check_existing_copies(
session: Session,
destination_location: models.DataLocation,
raw_data_package_id: int,
) -> bool:
"""Check if a specific package already exists at the destination.
For RawDataPackages, we check if they are STAGED (unpacked and physical file removed)
rather than PRESENT, since we unpack packages and remove the physical file after staging.
Keyed on the destination location directly (#154).
"""
destination_location_id = destination_location.id
logger.debug(
"checking existing copies",
raw_package_id=raw_data_package_id,
destination_location_id=destination_location_id,
)
# Check if the package is already staged (STAGED status means unpacked and ready)
package_staged = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
and_(
models.RawDataPackagePhysicalCopy.raw_data_package_id
== raw_data_package_id,
models.RawDataPackagePhysicalCopy.data_location_id
== destination_location_id,
models.RawDataPackagePhysicalCopy.status
== models.PhysicalCopyStatus.STAGED,
)
)
.first()
is not None
)
if package_staged:
logger.debug(
"package already staged",
raw_package_id=raw_data_package_id,
)
return True
# Also check if package is currently present (in case it was staged but not yet marked as STAGED)
package_present = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
and_(
models.RawDataPackagePhysicalCopy.raw_data_package_id
== raw_data_package_id,
models.RawDataPackagePhysicalCopy.data_location_id
== destination_location_id,
models.RawDataPackagePhysicalCopy.status
== models.PhysicalCopyStatus.PRESENT,
)
)
.first()
is not None
)
if package_present:
logger.debug(
"package present at destination, verifying file",
raw_package_id=raw_data_package_id,
)
# Check that the file is really present and not empty
physical_copy = (
session.query(models.RawDataPackagePhysicalCopy)
.filter(
and_(
models.RawDataPackagePhysicalCopy.raw_data_package_id
== raw_data_package_id,
models.RawDataPackagePhysicalCopy.data_location_id
== destination_location_id,
)
)
.first()
)
if (
os.path.exists(physical_copy.full_path)
and os.path.getsize(physical_copy.full_path) > 0
):
logger.debug(
"package file confirmed present",
raw_package_id=raw_data_package_id,
)
return True
else:
logger.debug(
"package file missing or empty, fixing db record",
raw_package_id=raw_data_package_id,
)
# Fix database entry - mark as deleted since file is gone
physical_copy.status = models.PhysicalCopyStatus.DELETED
session.commit()
return False
else:
logger.debug(
"package not found at destination",
raw_package_id=raw_data_package_id,
)
return False
def _dispatch_staging_job(
session: Session, staging_job: models.StagingJob
) -> models.StagingOperationGroup:
"""Fan a StagingJob out into one StagingOperation per package (#154).
Creates a StagingOperationGroup and one child StagingOperation per
RawDataPackage, dispatches one celery task per operation (keyed on the
operation id so each package retries/recovers on its own), and syncs the
legacy StagingJob.status from the derived group status. Returns the group.
"""
queue_name = route_task_by_location(
OperationType.STAGING, staging_job.destination_data_location
)
group = models.StagingOperationGroup(active=True)
session.add(group)
session.flush()
operations = []
for raw_data_package in staging_job.raw_data_packages:
operation = models.StagingOperation(
raw_data_package_id=raw_data_package.id,
origin_location_id=staging_job.origin_data_location_id,
destination_location_id=staging_job.destination_data_location_id,
status=models.Status.SCHEDULED,
operation_group=group,
)
session.add(operation)
operations.append(operation)
session.flush()
for operation in operations:
stage_data_task.apply_async(
args=[operation.id],
queue=queue_name,
)
# Take the job OUT of the pending set at dispatch (same pattern as
# transfer_manager): the find-work loop re-selects PENDING jobs, and an
# all-SCHEDULED group derives to PENDING, so mirroring the derived status
# here would re-dispatch the job every poll cycle (#70 double-dispatch).
# SCHEDULED is set explicitly — the honest pre-start value (it is in the
# explorer's pending bucket) — diverging deliberately from the derived
# mirror only for this queued-but-not-started window. Once any child moves,
# _sync_owning_job_status takes the derived status back over.
staging_job.status = models.Status.SCHEDULED
session.commit()
logger.debug(
"staging job dispatched",
staging_job_id=staging_job.id,
operation_count=len(operations),
destination=staging_job.destination_data_location.name,
queue=queue_name,
)
return group
def _sync_job_status_from_group(
session: Session,
staging_job: models.StagingJob,
group: models.OperationGroup,
) -> None:
"""Mirror the group's DERIVED status onto the legacy StagingJob.status.
The group is the source of truth (status derived from its child operations,
never set independently). The legacy ``StagingJob.status`` column is the
backward-compatible mirror the staging API/explorer query for pending/done
counts. This NEVER touches ``StagingJob.active`` — the deletion retention
claim is honored unchanged by deletion_manager (#154).
"""
staging_job.status = group.status
def _find_job_for_group(
session: Session, group: models.OperationGroup
) -> Optional[models.StagingJob]:
"""Resolve the legacy StagingJob a StagingOperationGroup was dispatched for.
No FK links the two (ops-db is not edited here), so the correlation is the
one a single dispatch guarantees: a dispatch creates exactly one operation
per job package, so the job shares the group's destination AND has the SAME
package set. The match is exact (not subset) to stay unambiguous when one
job's packages are a subset of another's on a resubmit. Returns None when no
such job exists (e.g. an operation created outside the dispatch flow).
"""
operations = group.operations
if not operations:
return None
package_ids = {op.raw_data_package_id for op in operations}
destination_ids = {op.destination_location_id for op in operations}
if len(destination_ids) != 1:
return None
destination_id = next(iter(destination_ids))
candidates = (
session.query(models.StagingJob)
.options(selectinload(models.StagingJob.raw_data_packages))
.filter(models.StagingJob.destination_data_location_id == destination_id)
.all()
)
for job in candidates:
job_package_ids = {pkg.id for pkg in job.raw_data_packages}
if package_ids == job_package_ids:
return job
return None
[docs]
def process_staging_jobs(verbose: bool = False, session: Session = None) -> None:
"""Main function to process all pending staging jobs."""
if session is None:
logger.debug("no session provided, creating new one")
db = DatabaseConnection()
session, _ = db.get_connection()
should_close_session = True
else:
should_close_session = False
try:
pending_jobs = _get_pending_staging_jobs(session)
logger.info("processing pending staging jobs", job_count=len(pending_jobs))
for job in pending_jobs:
logger.debug(
"staging job details",
staging_job_id=job.id,
status=str(job.status),
retry_count=job.retry_count,
failure_error_message=job.failure_error_message,
origin_location=job.origin_data_location.name,
destination_location=job.destination_data_location.name,
package_count=len(job.raw_data_packages),
)
try:
# Fan the job out into per-package StagingOperations.
_dispatch_staging_job(session, job)
logger.info("staging job dispatched", staging_job_id=job.id)
except Exception as e:
logger.error(
"staging job dispatch failed",
staging_job_id=job.id,
error=e,
)
job.status = models.Status.FAILED
job.failure_error_message = str(e)
session.commit()
finally:
if should_close_session:
session.close()
[docs]
def get_processing_locations_for_site(
session: Session, site: models.Site
) -> List[models.DataLocation]:
"""Get all processing locations for a specific site."""
return (
session.query(models.DataLocation)
.filter(
models.DataLocation.site_id == site.id,
models.DataLocation.location_type == models.LocationType.PROCESSING,
models.DataLocation.active == True, # noqa: E712
)
.all()
)
[docs]
def get_sites_with_processing_locations(session: Session) -> List[models.Site]:
"""Get all sites that have processing locations."""
return (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.PROCESSING,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
