import os
import shlex
import socket
import subprocess
import tarfile
import tempfile
import shutil
from datetime import datetime, timezone
from typing import List, Dict, Tuple, Optional
from collections import defaultdict
from sqlalchemy.orm import Session
from ccat_ops_db import models
from .config.config import ccat_data_transfer_settings
from .database import DatabaseConnection
from .setup_celery_app import app, make_celery_task
from .utils import (
unique_id,
calculate_checksum,
safe_join,
SSH_OPTS,
)
from .exceptions import DatabaseOperationError, ConfigurationError
from .logging_utils import get_structured_logger
from .boundary import ensure_readable, diagnose_source_missing
from .queue_discovery import route_task_by_location
from .operation_types import OperationType
# Set up global logger
logger = get_structured_logger(__name__)
# Create enhanced task base class for raw data package operations
SQLAlchemyTask = make_celery_task()
def _utcnow() -> datetime:
"""Timezone-aware now, used for operation start/end timestamps."""
return datetime.now(timezone.utc)
[docs]
class RawDataPackageOperations(SQLAlchemyTask):
"""Base class for raw data package operations with error handling.
Packaging runs as a ``PackagingOperation`` (ADR-0003): the ``operation_id``
the base task handles is the ``PackagingOperation.id`` and status lives on
that row. The ``RawDataPackage`` artifact keeps only its lifecycle ``state``.
Retry counting falls to the base ``get_retry_count`` (the uniform
``Operation.retry_count``); no per-artifact override remains.
"""
operation_type = "raw_data_package"
max_retries = 3
[docs]
def mark_in_progress(self, session, operation_id):
"""Flip the PackagingOperation to IN_PROGRESS at task start (ADR-0003).
Single seam, so SCHEDULED (queued) and IN_PROGRESS (running) are
DB-distinguishable on the operation row.
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-run assembly on a
redelivered/recovery-redelivered op (#70/#69 class).
"""
operation = session.query(models.PackagingOperation).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 PackagingOperation for retry."""
try:
operation = session.query(models.PackagingOperation).get(operation_id)
if operation:
operation.status = models.Status.PENDING
operation.retry_count = (operation.retry_count or 0) + 1
# Clear both the free-text last-error and its structured sibling
# on reset (#117); the durable trail stays in OperationFailureEvent.
operation.failure_error_message = None
operation.error_context = None
session.commit()
logger.info(
"packaging operation reset for retry",
operation_id=operation_id,
retry_count=operation.retry_count,
)
except Exception as e:
logger.error(
"failed to reset packaging operation state",
operation_id=operation_id,
error=str(e),
)
[docs]
def mark_permanent_failure(self, session, operation_id, exc):
"""Mark the PackagingOperation as permanently failed."""
try:
operation = session.query(models.PackagingOperation).get(operation_id)
if operation:
operation.status = models.Status.FAILED
operation.end_time = _utcnow()
# Record the last error and its structured breadcrumb (#117).
operation.failure_error_message = str(exc)
operation.error_context = self._current_error_context
session.commit()
logger.error(
"packaging operation marked as permanently failed",
operation_id=operation_id,
error=str(exc),
)
except Exception as e:
logger.error(
"failed to mark packaging operation as failed",
operation_id=operation_id,
error=str(e),
)
@app.task(
base=RawDataPackageOperations,
name="ccat:raw_data_package:create_package",
bind=True,
)
def create_raw_data_package_task(
self,
packaging_operation_id: int,
source_location_id: int,
session: Optional[Session] = None,
) -> bool:
"""
Create a raw data package by assembling files from the source location.
Parameters
----------
self : celery.Task
The Celery task instance.
packaging_operation_id : int
The ID of the PackagingOperation that drives this work. Status lives on
this row; the RawDataPackage it points at is the artifact assembled.
source_location_id : int
The ID of the source 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_raw_data_package_internal(
session, packaging_operation_id, source_location_id
)
else:
return _create_raw_data_package_internal(
session, packaging_operation_id, source_location_id
)
def _create_raw_data_package_internal(
session: Session, packaging_operation_id: int, source_location_id: int
) -> bool:
"""
Internal function to create a raw data package.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
packaging_operation_id : int
The ID of the PackagingOperation driving this work (operation status
lives here; it points at the RawDataPackage to assemble).
source_location_id : int
The ID of the source DataLocation.
Returns
-------
bool
True if the operation was successful.
"""
logger.debug(
"creating raw data package",
packaging_operation_id=packaging_operation_id,
location_id=source_location_id,
)
try:
# The operation row drives status; the package it points at is the
# artifact being assembled.
packaging_operation = session.query(models.PackagingOperation).get(
packaging_operation_id
)
if not packaging_operation:
logger.error(
"packaging operation not found",
packaging_operation_id=packaging_operation_id,
)
return False
if packaging_operation.status == models.Status.COMPLETED:
logger.debug(
"packaging operation already completed",
packaging_operation_id=packaging_operation_id,
)
return True
raw_data_package = session.query(models.RawDataPackage).get(
packaging_operation.raw_data_package_id
)
if not raw_data_package:
logger.error(
"raw data package not found",
raw_package_id=packaging_operation.raw_data_package_id,
)
return False
source_location = session.query(models.DataLocation).get(source_location_id)
if not source_location:
logger.error("source location not found", location_id=source_location_id)
return False
# Get primary buffer for this site
primary_buffer = get_primary_buffer_for_site(session, source_location.site)
if not primary_buffer:
logger.error(
"no active buffer found for site",
site=source_location.site.short_name,
)
return False
# Get all raw data files for this package
raw_data_files = raw_data_package.raw_data_files
if not raw_data_files:
logger.warning(
"no raw data files found for package",
raw_package_id=raw_data_package.id,
)
return False
logger.debug(
"creating archive",
raw_package_id=raw_data_package.id,
file_count=len(raw_data_files),
)
# Produced package copy is captured for copy-anchored lineage (ADR-0003)
# and set in whichever assembly branch runs.
produced_package_copy = None
# Check if we should create the package in the buffer directly
if is_same_host(primary_buffer):
logger.debug("creating package in buffer directly")
physical_raw_data_file_temp_copies = []
# iterate over all relevant raw data files for this package
for raw_data_file in raw_data_files:
logger.debug(
"copying raw data file to buffer",
rel_path=raw_data_file.relative_path,
buffer=primary_buffer.name,
buffer_host=primary_buffer.host,
source=source_location.name,
source_host=source_location.host,
)
# copy raw data files to buffer
# if both localhost and primary buffer are on the same host, copy the file to the buffer
source_path = safe_join(
source_location.path, raw_data_file.relative_path
)
buffer_path = safe_join(
primary_buffer.path, raw_data_file.relative_path
)
logger.debug(
"source and buffer paths",
path=source_path,
buffer_path=buffer_path,
)
if is_same_host(source_location) and is_same_host(primary_buffer):
if source_path != buffer_path:
logger.debug(
"copying raw data file to buffer directly",
rel_path=raw_data_file.relative_path,
buffer=primary_buffer.name,
source=source_location.name,
)
# make sure the directory exists
os.makedirs(os.path.dirname(buffer_path), exist_ok=True)
# Guard the source read so a missing raw file surfaces
# as a diagnosed SourceMissingError (#118).
ensure_readable(
source_path,
side="source",
step="raw_file_buffer_copy",
host=getattr(source_location, "host", None),
diagnose=lambda rdf=raw_data_file: diagnose_source_missing(
session,
models.RawDataFilePhysicalCopy,
fk_attr="raw_data_file_id",
fk_value=rdf.id,
location=source_location,
),
)
shutil.copy2(source_path, buffer_path)
else:
logger.debug(
"raw data file already in buffer",
rel_path=raw_data_file.relative_path,
buffer=primary_buffer.name,
source=source_location.name,
)
else:
os.makedirs(os.path.dirname(buffer_path), exist_ok=True)
# copy the file to the buffer using scp
scp_command = [
"scp",
f"{source_location.user}@{source_location.host}:{source_path}",
f"{primary_buffer.user}@{primary_buffer.host}:{buffer_path}",
]
subprocess.run(
scp_command, check=True, capture_output=True, text=True
)
# create physical copy record at buffer location
physical_copy = models.RawDataFilePhysicalCopy(
raw_data_file=raw_data_file,
data_location=primary_buffer,
checksum=raw_data_file.checksum,
)
session.add(physical_copy)
physical_raw_data_file_temp_copies.append(physical_copy)
session.commit()
# Create package directly in buffer
archive_success = _create_package_archive(
session, raw_data_package, primary_buffer, raw_data_files
)
if not archive_success:
logger.error(
"failed to create package in buffer",
raw_package_id=raw_data_package.id,
)
return False
# Create physical copy record at buffer location
physical_copy = models.RawDataPackagePhysicalCopy(
raw_data_package=raw_data_package,
data_location=primary_buffer,
checksum=raw_data_package.checksum,
)
session.add(physical_copy)
produced_package_copy = physical_copy
for physical_raw_data_file_temp_copy in physical_raw_data_file_temp_copies:
# remove the temp copy we are local
os.remove(physical_raw_data_file_temp_copy.full_path)
session.delete(physical_raw_data_file_temp_copy)
elif is_same_host(source_location):
# we are on the source location, we can create the package directly
archive_success = _create_package_archive(
session, raw_data_package, source_location, raw_data_files
)
if not archive_success:
logger.error(
"failed to create package in source location",
raw_package_id=raw_data_package.id,
)
return False
# Create physical copy record at source location
source_physical_copy = models.RawDataPackagePhysicalCopy(
raw_data_package=raw_data_package,
data_location=source_location,
checksum=raw_data_package.checksum,
)
session.add(source_physical_copy)
session.commit()
# Now we have to transfer it to the buffer
# check if buffer is localhost
# Ensure buffer directory exists
source_path = safe_join(
source_location.path, raw_data_package.relative_path
)
buffer_path = safe_join(
primary_buffer.path, raw_data_package.relative_path
)
if is_same_host(primary_buffer):
os.makedirs(os.path.dirname(buffer_path), exist_ok=True)
else:
success, output = execute_remote_command(
primary_buffer.host,
primary_buffer.user,
f"mkdir -p {shlex.quote(os.path.dirname(buffer_path))}",
)
if not success:
logger.error(
"failed to create buffer directory",
path=os.path.dirname(buffer_path),
error=output,
)
return False
if is_same_host(primary_buffer):
# make sure the directory exists
os.makedirs(os.path.dirname(buffer_path), exist_ok=True)
# Guard the source read so a missing package source surfaces as
# a diagnosed SourceMissingError (#118).
ensure_readable(
source_path,
side="source",
step="raw_package_buffer_copy",
host=getattr(source_location, "host", None),
diagnose=lambda: diagnose_source_missing(
session,
models.RawDataPackagePhysicalCopy,
fk_attr="raw_data_package_id",
fk_value=raw_data_package.id,
location=source_location,
),
)
shutil.copy2(source_path, buffer_path)
else:
# copy the package to the buffer using scp
# make sure the remote directory exists
os.makedirs(os.path.dirname(buffer_path), exist_ok=True)
scp_command = [
"scp",
"-r",
f"{source_location.user}@{source_location.host}:{source_path}",
f"{primary_buffer.user}@{primary_buffer.host}:{buffer_path}",
]
subprocess.run(scp_command, check=True, capture_output=True, text=True)
# create physical copy record at buffer location
buffer_physical_copy = models.RawDataPackagePhysicalCopy(
raw_data_package=raw_data_package,
data_location=primary_buffer,
checksum=raw_data_package.checksum,
)
session.add(buffer_physical_copy)
produced_package_copy = buffer_physical_copy
# remove the package from the source location
os.remove(source_path)
session.delete(source_physical_copy)
else:
logger.error(
"cannot create package in buffer, not on same host as source or "
"buffer, remote to remote transfer is not supported, this celery "
"queue is receiving a task it is not intended to handle"
)
return False
# Operation status lives on the PackagingOperation row (ADR-0003); the
# RawDataPackage keeps only its lifecycle state. The legacy status column
# is left untouched (dropped in #95) — the manager no longer drives it.
packaging_operation.status = models.Status.COMPLETED
packaging_operation.end_time = _utcnow()
raw_data_package.state = models.PackageState.TRANSFERRING
# Copy-anchored lineage (ADR-0003): packaging consumes the source raw-file
# copies and produces the package copy in the buffer.
consumed_copies = _source_raw_file_copies(
session, raw_data_files, source_location
)
produced = [produced_package_copy] if produced_package_copy else []
_record_packaging_lineage(
packaging_operation, consumed=consumed_copies, produced=produced
)
session.commit()
logger.info(
"raw data package assembled and transferred",
raw_package_id=raw_data_package.id,
location=source_location.name,
file_count=len(raw_data_files),
size_bytes=raw_data_package.size,
)
return True
except Exception as e:
logger.error(
"error creating raw data package",
packaging_operation_id=packaging_operation_id,
error=str(e),
)
session.rollback()
return False
def _source_raw_file_copies(
session: Session,
raw_data_files: List[models.RawDataFile],
source_location: models.DataLocation,
) -> List[models.RawDataFilePhysicalCopy]:
"""Return the source-location physical copies of the given raw files.
These are the inputs packaging consumes. A file without a recorded source
copy is simply omitted (lineage is best-effort over what is recorded, not a
correctness gate)."""
file_ids = [f.id for f in raw_data_files]
if not file_ids:
return []
return (
session.query(models.RawDataFilePhysicalCopy)
.filter(
models.RawDataFilePhysicalCopy.raw_data_file_id.in_(file_ids),
models.RawDataFilePhysicalCopy.data_location_id == source_location.id,
)
.all()
)
def _record_packaging_lineage(
operation: models.PackagingOperation,
consumed: List[models.PhysicalCopy],
produced: List[models.PhysicalCopy],
) -> None:
"""Attach copy-anchored consume/produce lineage to a PackagingOperation.
Packaging consumes the input raw-file copies and produces the package copy
(ADR-0003). Idempotent: a copy already linked is not duplicated, so a
re-run does not grow the lineage."""
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)
[docs]
def is_same_host(data_location: models.DiskDataLocation) -> bool:
"""
Check if the current worker is on the same host as the DataLocation.
Parameters
----------
data_location : models.DiskDataLocation
The DataLocation to check against.
Returns
-------
bool
True if the worker is on the same host as the DataLocation.
"""
if data_location.host in ["localhost", "127.0.0.1"]:
return True
# Get the hostname of the current container
container_hostname = socket.gethostname()
logger.debug(
"host check",
container_hostname=container_hostname,
location_host=data_location.host,
)
# If we're in development mode, treat all localhost as same host
if ccat_data_transfer_settings.DEVELOPMENT_MODE_LOCALHOST_ONLY:
logger.debug("development mode: treating all hosts as localhost")
return True
return container_hostname == data_location.host
[docs]
def execute_remote_command(host: str, user: str, command: str) -> Tuple[bool, str]:
"""
Execute a command on a remote host via SSH.
Parameters
----------
host : str
The remote host to execute the command on.
user : str
The user to execute the command as.
command : str
The command to execute.
Returns
-------
Tuple[bool, str]
A tuple containing (success, output/error message)
"""
try:
ssh_command = ["ssh", *SSH_OPTS, f"{user}@{host}", command]
result = subprocess.run(ssh_command, capture_output=True, text=True, check=True)
return True, result.stdout
except subprocess.CalledProcessError as e:
return False, e.stderr
def _create_package_archive(
session: Session,
raw_data_package: models.RawDataPackage,
source_location: models.DataLocation,
raw_data_files: List[models.RawDataFile],
) -> bool:
"""
Create a tar archive from raw data files.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
raw_data_package : models.RawDataPackage
The raw data package to create.
source_location : models.DataLocation
The source data location.
raw_data_files : List[models.RawDataFile]
List of raw data files to include.
Returns
-------
bool
True if successful, False otherwise.
"""
try:
# Ensure source_location is a DiskDataLocation
if not isinstance(source_location, models.DiskDataLocation):
raise ValueError("Source location must be a DiskDataLocation")
# Check if we're on the same host as the source location
if is_same_host(source_location):
# We can work directly with the files
return _create_package_archive_local(
session, raw_data_package, source_location, raw_data_files
)
else:
# We need to work remotely
return _create_package_archive_remote(
session, raw_data_package, source_location, raw_data_files
)
except Exception as e:
logger.error(
"error creating package archive",
raw_package_id=raw_data_package.id,
error=str(e),
)
return False
def _create_package_archive_local(
session: Session,
raw_data_package: models.RawDataPackage,
source_location: models.DiskDataLocation,
raw_data_files: List[models.RawDataFile],
) -> bool:
"""Create package archive when worker is on same host as source location."""
try:
# Ensure the package directory exists
package_dir = os.path.dirname(
safe_join(source_location.path, raw_data_package.relative_path)
)
os.makedirs(package_dir, exist_ok=True)
archive_path = safe_join(
source_location.path, raw_data_package.relative_path
)
total_size = 0
with tarfile.open(archive_path, "w:gz") as tar:
for raw_file in raw_data_files:
file_path = safe_join(source_location.path, raw_file.relative_path)
if os.path.exists(file_path):
# Use the relative_path as the arcname to preserve the hierarchical structure
# This ensures the tar contains the same directory structure as the source
tar.add(file_path, arcname=raw_file.relative_path)
total_size += raw_file.size
logger.debug("added file to archive", rel_path=raw_file.relative_path)
else:
# File not present yet — normal during incremental source scans
logger.debug("raw data file not found, skipping", path=file_path)
# Calculate checksum and update package
checksum = calculate_checksum(archive_path)
if checksum:
raw_data_package.checksum = checksum
raw_data_package.size = total_size
return True
else:
logger.error(
"failed to calculate checksum for archive",
raw_package_id=raw_data_package.id,
)
return False
except Exception as e:
logger.error(
"error creating local package archive",
raw_package_id=raw_data_package.id,
error=str(e),
)
return False
def _create_package_archive_remote(
session: Session,
raw_data_package: models.RawDataPackage,
source_location: models.DiskDataLocation,
raw_data_files: List[models.RawDataFile],
) -> bool:
"""Create package archive when worker is on different host than source location."""
try:
# Create a temporary directory for the archive
with tempfile.TemporaryDirectory() as _:
# Create the archive command
archive_path = os.path.join(
source_location.path, raw_data_package.relative_path
)
# Use tar with -C to change to source directory and preserve relative paths
# This ensures the tar contains the same directory structure as the source
tar_command = f"cd {shlex.quote(source_location.path)} && tar czf {shlex.quote(archive_path)} "
tar_command += " ".join(shlex.quote(f.relative_path) for f in raw_data_files)
# Execute the command remotely
success, output = execute_remote_command(
source_location.host, source_location.user, tar_command
)
if not success:
logger.error(
"failed to create remote archive",
raw_package_id=raw_data_package.id,
error=output,
)
return False
# Calculate checksum remotely (must match local xxHash64 algorithm)
checksum_command = f"xxh64sum {shlex.quote(archive_path)} | cut -d' ' -f1"
success, checksum = execute_remote_command(
source_location.host, source_location.user, checksum_command
)
if not success:
logger.error(
"failed to calculate remote checksum",
raw_package_id=raw_data_package.id,
error=checksum,
)
return False
# Calculate total size remotely
size_command = f"du -b {shlex.quote(archive_path)} | cut -f1"
success, size_str = execute_remote_command(
source_location.host, source_location.user, size_command
)
if not success:
logger.error(
"failed to get remote file size",
raw_package_id=raw_data_package.id,
error=size_str,
)
return False
# Update package with remote information
raw_data_package.checksum = checksum.strip()
raw_data_package.size = int(size_str.strip())
return True
except Exception as e:
logger.error(
"error creating remote package archive",
raw_package_id=raw_data_package.id,
error=str(e),
)
return False
[docs]
def get_unpackaged_raw_data_files(
session: Session, source_location: models.DataLocation
) -> List[models.RawDataFile]:
"""
Retrieve all raw data files from a source location that are not yet assigned to a package.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
source_location : models.DataLocation
The source data location to scan.
Returns
-------
List[models.RawDataFile]
List of unpackaged raw data files.
"""
logger.debug(
"scanning source location",
location=source_location.name,
location_id=source_location.id,
)
# Get all raw data files from this source location that are not yet packaged
unpackaged_files = (
session.query(models.RawDataFile)
.filter(
models.RawDataFile.source_location_id == source_location.id,
models.RawDataFile.raw_data_package_id.is_(None),
)
.all()
)
logger.debug(
"unpackaged files found",
location=source_location.name,
file_count=len(unpackaged_files),
)
return unpackaged_files
[docs]
def group_files_by_execution_and_module(
raw_data_files: List[models.RawDataFile],
) -> Dict[Tuple[int, int], List[models.RawDataFile]]:
"""
Group raw data files by ExecutedObsUnit and InstrumentModule.
Parameters
----------
raw_data_files : List[models.RawDataFile]
List of raw data files to group.
Returns
-------
Dict[Tuple[int, int], List[models.RawDataFile]]
Dictionary mapping (executed_obs_unit_id, instrument_module_id) to list of files.
"""
grouped_files = defaultdict(list)
for raw_file in raw_data_files:
imc = raw_file.instrument_module_configuration
instrument_module_id = imc.instrument_module_id if imc else None
key = (raw_file.executed_obs_unit_id, instrument_module_id)
grouped_files[key].append(raw_file)
logger.debug(
"files grouped into packages",
file_count=len(raw_data_files),
package_count=len(grouped_files),
)
return dict(grouped_files)
[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 resolved",
site=site.short_name,
location=buffer_location.name,
)
else:
logger.warning("no active buffer found for site", site=site.short_name)
return buffer_location
[docs]
def create_raw_data_packages_for_location(
session: Session, source_location: models.DataLocation
) -> None:
"""
Create raw data packages for unpackaged files in a source location.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
source_location : models.DataLocation
The source location to process.
"""
logger.debug(
"processing source location",
location=source_location.name,
location_id=source_location.id,
)
# Get unpackaged files
unpackaged_files = get_unpackaged_raw_data_files(session, source_location)
if not unpackaged_files:
logger.debug("no unpackaged files found", location=source_location.name)
return
# Group files by execution and module
grouped_files = group_files_by_execution_and_module(unpackaged_files)
# Get primary buffer for this site
primary_buffer = get_primary_buffer_for_site(session, source_location.site)
if not primary_buffer:
raise ConfigurationError(
f"No active buffer found for site {source_location.site.short_name}"
)
# Create packages for each group (records PENDING packaging operations; no
# inline dispatch — #70).
for (executed_obs_unit_id, instrument_module_id), files in grouped_files.items():
_create_raw_data_package_entry(
session,
files,
executed_obs_unit_id,
instrument_module_id,
source_location,
primary_buffer,
)
session.commit()
# Poll step: dispatch any PENDING packaging operation exactly once.
schedule_raw_data_package_creation(session, source_location)
logger.info(
"raw data packages scheduled for location",
location=source_location.name,
package_count=len(grouped_files),
)
def _create_raw_data_package_entry(
session: Session,
files: List[models.RawDataFile],
executed_obs_unit_id: int,
instrument_module_id: int,
source_location: models.DataLocation,
target_buffer: models.DataLocation,
) -> models.RawDataPackage:
"""
Create a RawDataPackage entry and its PENDING PackagingOperation.
No task is dispatched here: dispatch is a separate poll step
(``schedule_raw_data_package_creation``). Splitting create from dispatch
closes the inline create-and-dispatch double-dispatch window (#70) — a crash
between the two just leaves a PENDING operation for the next poll to pick up.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
files : List[models.RawDataFile]
List of files to include in the package.
executed_obs_unit_id : int
The executed observation unit ID.
instrument_module_id : int
The instrument module ID.
source_location : models.DataLocation
The source data location.
target_buffer : models.DataLocation
The target buffer location.
Returns
-------
models.RawDataPackage
The created raw data package.
"""
# Get executed obs unit and instrument module for naming
executed_obs_unit = session.query(models.ExecutedObsUnit).get(executed_obs_unit_id)
instrument_module = session.query(models.InstrumentModule).get(instrument_module_id)
if not executed_obs_unit or not instrument_module:
raise ValueError("Invalid executed_obs_unit_id or instrument_module_id")
# Generate unique package name and path
package_id = unique_id()
package_name = (
f"{executed_obs_unit.obs_unit.name}_{instrument_module.name}_{package_id}"
)
relative_path = f"raw_data_packages/{package_name}.tar.gz"
# Create RawDataPackage entry. Lifecycle status lives on the PackagingOperation
# row created below (ADR-0003); RawDataPackage.status was dropped in ops-db#95.
raw_data_package = models.RawDataPackage(
name=package_name,
relative_path=relative_path,
executed_obs_unit_id=executed_obs_unit_id,
instrument_module_id=instrument_module_id,
obs_unit_id=executed_obs_unit.obs_unit_id,
state=models.PackageState.WAITING,
size=sum(f.size for f in files),
checksum="", # Will be calculated during assembly
)
# Associate files with the package
for file in files:
file.raw_data_package = raw_data_package
session.add(raw_data_package)
session.flush() # Get the ID
# Record the packaging operation as PENDING work. Status lives on this row
# (ADR-0003); the poll step flips it to SCHEDULED and dispatches the task.
packaging_operation = models.PackagingOperation(
raw_data_package_id=raw_data_package.id,
status=models.Status.PENDING,
)
session.add(packaging_operation)
session.flush()
# One INFO per raw data package created (assembly NOT dispatched here).
logger.info(
"raw data package created",
raw_package_id=raw_data_package.id,
packaging_operation_id=packaging_operation.id,
location=source_location.name,
file_count=len(files),
size_bytes=raw_data_package.size,
)
return raw_data_package
[docs]
def schedule_raw_data_package_creation(
session: Session, source_location: models.DataLocation
) -> None:
"""Dispatch assembly tasks for this location's PENDING packaging ops.
The standard poll→SCHEDULED→dispatch loop (ADR-0003): each ``PENDING``
``PackagingOperation`` 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
— this is what closes the inline create-and-dispatch double-dispatch window
(#70). The status flip is committed per operation so a crash mid-loop cannot
re-dispatch an already-scheduled operation.
The query is SCOPED to ``source_location``: a packaging op belongs to the
location its package's raw files were observed at
(``RawDataFile.source_location_id``). Without that scope the per-location
poll would scoop a foreign op (e.g. one a recovery reset left ``PENDING``)
and dispatch it to the wrong queue with the wrong ``source_location_id``,
failing its ``is_same_host`` checks on the wrong host. Routing and args
therefore match the location whose ops were selected.
"""
pending_operations = (
session.query(models.PackagingOperation)
.join(
models.RawDataPackage,
models.PackagingOperation.raw_data_package_id == models.RawDataPackage.id,
)
.join(
models.RawDataFile,
models.RawDataFile.raw_data_package_id == models.RawDataPackage.id,
)
.filter(
models.PackagingOperation.status == models.Status.PENDING,
models.RawDataFile.source_location_id == source_location.id,
)
.distinct()
.all()
)
if not pending_operations:
logger.debug("no pending packaging operations", location=source_location.name)
return
queue_name = route_task_by_location(
OperationType.RAW_DATA_PACKAGE_CREATION, source_location
)
for operation in pending_operations:
operation.status = models.Status.SCHEDULED
session.commit()
create_raw_data_package_task.apply_async(
args=[operation.id, source_location.id], queue=queue_name
)
logger.info(
"packaging operation scheduled",
packaging_operation_id=operation.id,
raw_package_id=operation.raw_data_package_id,
location=source_location.name,
queue=queue_name,
)
logger.info(
"packaging operations scheduled",
operation_count=len(pending_operations),
location=source_location.name,
)
[docs]
def get_sites_with_source_locations(session: Session) -> List[models.Site]:
"""
Get all sites that have active SOURCE data locations.
Parameters
----------
session : sqlalchemy.orm.Session
The database session.
Returns
-------
List[models.Site]
List of sites with source locations.
"""
# PRint debug details for all sites before filtering
sites = (
session.query(models.Site)
.join(models.DataLocation)
.filter(
models.DataLocation.location_type == models.LocationType.SOURCE,
models.DataLocation.active == True, # noqa: E712
)
.distinct()
.all()
)
logger.debug("sites with active source locations", site_count=len(sites))
return sites
[docs]
def create_raw_data_packages(verbose: bool = False, session: Session = None) -> None:
"""
Scan all source locations and create raw data packages for unpackaged files.
This function manages the process of creating raw data packages by:
1. Finding all sites with active SOURCE data locations.
2. For each source location, finding unpackaged raw data files.
3. Grouping files by ExecutedObsUnit and InstrumentModule.
4. Creating RawDataPackage entries in the database.
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.
"""
# Own a fresh session per call when none is supplied, so a poisoned/dead
# connection from a prior cycle never carries over to the next.
should_close_session = False
if session is None:
db = DatabaseConnection()
session, _ = db.get_connection()
should_close_session = True
try:
# Get all sites with source locations
source_sites = get_sites_with_source_locations(session)
if not source_sites:
logger.info("no sites with source locations found")
return
logger.debug("processing sites with source locations", site_count=len(source_sites))
for site in source_sites:
logger.debug("processing site", site=site.name, site_id=site.id)
# Get all active source locations for this site
source_locations = (
session.query(models.DataLocation)
.filter(
models.DataLocation.site_id == site.id,
models.DataLocation.location_type == models.LocationType.SOURCE,
models.DataLocation.active == True, # noqa: E712
)
.all()
)
for source_location in source_locations:
try:
create_raw_data_packages_for_location(session, source_location)
except Exception as e:
logger.error(
"error processing source location",
location=source_location.name,
location_id=source_location.id,
error=str(e),
)
session.rollback()
raise DatabaseOperationError(
f"Failed to create raw data packages for {source_location.name}: {str(e)}"
) from e
logger.info("raw data package scan complete")
except Exception as e:
logger.exception("an error occurred while creating raw data packages")
raise RuntimeError("Failed to create raw data packages") from e
finally:
if should_close_session:
session.close()
