"""File upload logic.""" import logging import uuid from datetime import datetime, timedelta, timezone from abacus_common_logic.connectors.database import db from botocore.exceptions import BotoCoreError, ClientError from marshmallow import ValidationError from owsresponse import response from abacus_file_upload.connectors import get_s3_connector from abacus_file_upload.constants import ( ABACUS_OUTBOX_EVENT_TYPES, ABACUS_OUTBOX_TARGET_TYPES, EVENT_PROCESSING_STATUSES, INVALID_STATUS_CHANGE, PRESIGNED_URL_EXPIRATION_SECS, UPLOAD_EXPIRATION_SECS, UPLOAD_STATUSES, UPLOAD_TYPE_TO_BUCKET, VALID_FILE_UPLOAD_STATUSES_TRANSITIONS, ) from abacus_file_upload.logic.file_upload_helpers import ( cleanup_s3_upload, get_file_extension, handle_multipart_upload, handle_single_part_upload, handle_upload_operation_error, quarantine_s3_upload, validate_with_upload_config, verify_s3_upload, ) from abacus_file_upload.models import AbacusOutbox, FileUpload, FileUploadConfig from abacus_file_upload.schemas import ( FileUploadDetailSchema, FileUploadDetailVerboseSchema, FileUploadStatusPutSchema, InitiateUploadRequestSchema, InitiateUploadResponseSchema, ) from abacus_file_upload.utils import ( generate_key, validation_error, ) from core.config import Config # Initialize schemas initiate_upload_request_schema = InitiateUploadRequestSchema() initiate_upload_response_schema = InitiateUploadResponseSchema() file_upload_detail_schema = FileUploadDetailSchema() file_upload_detail_verbose_schema = FileUploadDetailVerboseSchema() logger = logging.getLogger(Config.LOGGER_NAME) def cancel_upload(file_key: str) -> response.Response: """Cancel an upload. Args: file_key: File key Returns: Response with cancellation information """ try: file_upload = FileUpload.find_by_file_key(file_key) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) if not file_upload: return response.Response(message=f"Upload not found: '{file_key}'", status=404) if file_upload.upload_status == UPLOAD_STATUSES.COMPLETE: return validation_error('Cannot cancel a completed upload') if file_upload.upload_status == UPLOAD_STATUSES.CANCELLED: # Already cancelled. Return current state return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) try: s3_connector = get_s3_connector() # Clean up S3 resources cleanup_s3_upload(s3_connector, file_upload) # Update database file_upload.update_attributes(upload_status=UPLOAD_STATUSES.CANCELLED) FileUpload.commit_changes() # Refresh to get DB-generated values db.session.refresh(file_upload) # Return full upload record return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) except Exception as e: # S3, database, or unexpected errors return handle_upload_operation_error(file_upload, file_key, e, 'cancel', 500) def complete_upload(file_key: str) -> response.Response: """Complete the upload record after client has completed S3 upload. The client is responsible for completing the S3 upload directly using presigned URLs. This endpoint verifies the upload succeeded and updates the database. Args: file_key: File key Returns: Response with completion information """ try: file_upload = FileUpload.find_by_file_key(file_key) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) if not file_upload: return response.Response(message=f"Upload not found: '{file_key}'", status=404) if file_upload.upload_status == UPLOAD_STATUSES.COMPLETE: # Already completed. Return current state return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) if file_upload.upload_status == UPLOAD_STATUSES.CANCELLED: return validation_error('Cannot complete a cancelled upload') try: s3_connector = get_s3_connector() # Verify S3 upload is complete and valid verify_s3_upload(s3_connector, file_upload) # Update database file_upload.update_attributes( upload_status=UPLOAD_STATUSES.COMPLETE, completed_at=datetime.now(timezone.utc), ) FileUpload.commit_changes() AbacusOutbox.create( target_type=ABACUS_OUTBOX_TARGET_TYPES.FILE_UPLOAD, target_id=file_upload.file_upload_id, event_type=ABACUS_OUTBOX_EVENT_TYPES.FILE_UPLOAD_COMPLETED, status=EVENT_PROCESSING_STATUSES.PENDING, details={'upload_type': file_upload.file_upload_config.upload_type}, ) # Refresh to get DB-generated values db.session.refresh(file_upload) # Return full upload record return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) except ValueError as e: return handle_upload_operation_error(file_upload, file_key, e, 'complete', 400) except Exception as e: # S3, database, or unexpected errors return handle_upload_operation_error(file_upload, file_key, e, 'complete', 500) def download_file( file_key: str, expires_in: int = PRESIGNED_URL_EXPIRATION_SECS ) -> response.Response: """Generate a presigned download URL for an uploaded file. Args: file_key: File key expires_in: URL expiration time in seconds (default: 3600 = 1 hour) Returns: Response with presigned download URL """ try: file_upload = FileUpload.find_by_file_key(file_key) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) if not file_upload: return response.Response(message=f"Upload not found: '{file_key}'", status=404) if file_upload.upload_status != UPLOAD_STATUSES.COMPLETE: return validation_error( f"File not available for download. Upload status: '{file_upload.upload_status}'" ) try: s3_connector = get_s3_connector() download_url = s3_connector.generate_get_presigned_url( file_upload.s3_bucket, file_upload.s3_key, expires_in=expires_in ) return response.Response( message={ 'download_url': download_url, 'expires_in': expires_in, 'file_name': file_upload.original_file_name, 'file_size_bytes': file_upload.file_size_bytes, }, status=200, ) except (ClientError, BotoCoreError) as e: logger.error(f'Failed to generate download URL for {file_key}: {e}') return response.Response( message=f'Failed to generate download URL: {str(e)}', status=500, ) def get_file_upload(file_key: str, verbose: bool = False) -> response.Response: """Get file upload record. Args: file_key: File key verbose: If True, include audit fields (created_at, created_by, etc.) Returns: Response with file upload record """ try: file_upload = FileUpload.find_by_file_key(file_key) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) if not file_upload: return response.Response(message=f"Upload not found: '{file_key}'", status=404) schema = file_upload_detail_verbose_schema if verbose else file_upload_detail_schema return response.Response(message=schema.dump(file_upload), status=200) def initiate_upload(**params) -> response.Response: """Initiate a new file upload. Args: **params: Upload parameters including: - upload_type: Type of upload (matches config) - filename: Original filename - file_size_bytes: File size in bytes - mime_type: Optional MIME type - md5sum: MD5 hash for integrity verification - metadata: Optional metadata dict Returns: Response with upload information and pre-signed URLs """ # Validate input schema try: validated_data = initiate_upload_request_schema.load(params) except ValidationError as e: return validation_error(str(e.messages)) # Extract parameters upload_type = validated_data['upload_type'] filename = validated_data['filename'] file_size_bytes = validated_data['file_size_bytes'] mime_type = validated_data.get('mime_type') md5sum = validated_data.get('md5sum') metadata = validated_data.get('metadata') # Get and validate upload configuration try: config = FileUploadConfig.find_by_upload_type(upload_type) except Exception as e: logger.error(f'DB error finding upload config for type {upload_type}: {e}') return response.Response(message='Failed to retrieve upload config', status=500) if not config: return response.Response( message=f"Upload configuration not found for type: '{upload_type}'", status=404, ) # Validate S3 bucket is configured s3_bucket = UPLOAD_TYPE_TO_BUCKET.get(upload_type) if not s3_bucket: return validation_error( f"S3 bucket not configured for upload type: '{upload_type}'" ) # Get upload metadata (keys, buckets, expiration) expires_at = datetime.now(timezone.utc) + timedelta(seconds=UPLOAD_EXPIRATION_SECS) file_key = str(uuid.uuid4()) is_multipart = file_size_bytes >= (config.multipart_threshold_bytes or float('inf')) s3_key = generate_key( template=config.s3_key_template, file_key=file_key, upload_type=upload_type, file_name=filename, metadata=metadata, upload_time=datetime.now(timezone.utc), ) # S3 metadata keys must be lowercase for signature consistency s3_metadata = { 'filekey': file_key, 'filename': filename, 'md5sum': md5sum, 'uploadtype': upload_type, } try: s3_connector = get_s3_connector() # Validate file type and size against config validate_with_upload_config(config, filename, file_size_bytes) # Handle upload based on type (single-part or multipart) total_parts = 1 multipart_upload_id: str | None = None if is_multipart: upload_setup = handle_multipart_upload( s3_connector, s3_bucket, s3_key, file_size_bytes, config.min_multipart_chunk_size_bytes, s3_metadata, mime_type, ) multipart_upload_id = upload_setup['multipart_upload_id'] total_parts = upload_setup['total_parts'] response_data = { 'file_key': file_key, 'is_multipart': is_multipart, 'expires_at': expires_at.isoformat(), 'chunk_size_bytes': upload_setup['chunk_size_bytes'], 'parts': upload_setup['parts'], 'complete_url': upload_setup['complete_url'], } else: upload_setup = handle_single_part_upload( s3_connector, s3_bucket, s3_key, s3_metadata, md5sum, mime_type, ) response_data = { 'file_key': file_key, 'is_multipart': is_multipart, 'expires_at': expires_at.isoformat(), 'upload_url': upload_setup['upload_url'], 'required_headers': upload_setup['required_headers'], } # Create database record file_type = get_file_extension(filename) FileUpload.build( file_upload_config_id=config.file_upload_config_id, file_key=file_key, original_file_name=filename, file_size_bytes=file_size_bytes, file_type=file_type, mime_type=mime_type, s3_bucket=s3_bucket, s3_key=s3_key, md5sum=md5sum, upload_status=UPLOAD_STATUSES.INIT, multipart_upload_id=multipart_upload_id, total_parts=total_parts, upload_metadata=metadata, expires_at=expires_at, ) # Commit and return FileUpload.commit_changes() return response.Response( message=initiate_upload_response_schema.dump(response_data), status=201 ) except ValueError as e: # Validation errors (e.g., file too large for multipart) db.session.rollback() logger.warning(f'Validation error during upload initiation: {e}') return validation_error(str(e)) except (ClientError, BotoCoreError) as e: # S3 errors db.session.rollback() logger.error(f'S3 error during upload initiation: {e}') return response.Response( message=f'Failed to initiate upload with S3: {str(e)}', status=500 ) except Exception as e: # Unexpected errors db.session.rollback() logger.error(f'Unexpected error during upload initiation: {e}') return response.Response(message=str(e), status=500) def quarantine_infected_upload(file_key: str) -> response.Response: """Move infected files to the quarantine bucket to review. Args: file_key: File key Returns: Response with information """ try: file_upload = FileUpload.find_by_file_key(file_key) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) if not file_upload: return response.Response(message=f"Upload not found: '{file_key}'", status=404) if file_upload.upload_status == UPLOAD_STATUSES.COMPLETE: return validation_error('Cannot quarantine a completed upload') if file_upload.upload_status == UPLOAD_STATUSES.CANCELLED: return validation_error('Cannot quarantine a cancelled upload') if file_upload.upload_status == UPLOAD_STATUSES.QUARANTINED: # Already quarantined. Return current state return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) try: s3_connector = get_s3_connector() # Move file from one S3 bucket to another quarantine_s3_upload(s3_connector, file_upload) # Update database file_upload.update_attributes(upload_status=UPLOAD_STATUSES.QUARANTINED) FileUpload.commit_changes() # Refresh to get DB-generated values db.session.refresh(file_upload) # Return full upload record return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) except (ClientError, BotoCoreError, Exception) as e: # S3, database, or unexpected errors db.session.rollback() return handle_upload_operation_error( file_upload, file_key, e, 'quarantine', 500 ) def update_file_upload_status(file_key: str, params: dict) -> response.Response: """Update the file_upload status. Args: file_key: File key params: Update parameters including: - upload_status: Status of upload Returns: Response with updated status """ try: update_params = FileUploadStatusPutSchema().load(params) upload_status = update_params.get('upload_status') file_upload = FileUpload.find_by_file_key(file_key) if not file_upload: return response.Response( message=f"Upload not found: '{file_key}'", status=404 ) _validate_upload_status_transition(file_upload.upload_status, upload_status) except ValidationError as exc: return validation_error(str(exc)) except Exception as e: logger.error(f'DB error finding upload {file_key}: {e}') return response.Response(message='Failed to retrieve upload', status=500) status_func_mapping = { 'cancelled': cancel_upload, 'complete': complete_upload, 'quarantined': quarantine_infected_upload, } if upload_status in status_func_mapping: return status_func_mapping[upload_status](file_key) try: # Update database file_upload.update_attributes(upload_status=upload_status) FileUpload.commit_changes() # Refresh to get DB-generated values db.session.refresh(file_upload) # Return full upload record return response.Response( message=file_upload_detail_schema.dump(file_upload), status=200 ) except ValidationError as exc: return validation_error(str(exc)) except Exception as e: # database, or unexpected errors db.session.rollback() return handle_upload_operation_error( file_upload, file_key, e, upload_status, 500 ) def _validate_upload_status_transition(current_status: str, new_status: str) -> None: """Validate if upload_status can update from current status to new status. Args: current_status: existing status of file_upload new_status: new status for file_upload """ if current_status == new_status: return is_valid = new_status in VALID_FILE_UPLOAD_STATUSES_TRANSITIONS.get( current_status, [] ) if is_valid is False: raise ValidationError( INVALID_STATUS_CHANGE.format( object_type='FileUpload', current_status=current_status, new_status=new_status, ) )