#!/usr/bin/env python3 """ Batch upload documents to SeaweedFS and trigger Airflow DAG with progress monitoring. Uploads all supported files from a folder to SeaweedFS S3-compatible storage, triggers the document_extraction_dag for each file, and monitors DAG run status. Progress is tracked in a JSON file that updates in real-time. Usage: # Upload and trigger DAG for all files in a folder python scripts/batch_upload_and_monitor.py path/to/folder/ # Upload only (no DAG trigger) python scripts/batch_upload_and_monitor.py path/to/folder/ --no-trigger # Resume monitoring from a previous progress file python scripts/batch_upload_and_monitor.py --resume progress_20240101_120000.json # Custom bucket and prefix python scripts/batch_upload_and_monitor.py path/to/folder/ --bucket my-bucket --prefix inbox/ # Dry run python scripts/batch_upload_and_monitor.py path/to/folder/ --dry-run # Set polling interval for DAG monitoring python scripts/batch_upload_and_monitor.py path/to/folder/ --poll-interval 30 """ import argparse import json import os import sys import time import uuid from datetime import datetime, timezone from pathlib import Path import boto3 import requests from botocore.config import Config from dotenv import load_dotenv load_dotenv() SUPPORTED_EXTENSIONS = { '.pdf': 'application/pdf', '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', '.doc': 'application/msword', '.docm': 'application/vnd.ms-word.document.macroEnabled.12', '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation', '.ppt': 'application/vnd.ms-powerpoint', '.xls': 'application/vnd.ms-excel', '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', '.xlsm': 'application/vnd.ms-excel.sheet.macroEnabled.12', } # DAG run terminal states TERMINAL_STATES = {'success', 'failed'} # States that mean the run is still going RUNNING_STATES = {'queued', 'running'} class ProgressTracker: """Track upload and DAG run progress in a JSON file.""" def __init__(self, progress_file: str): self.progress_file = progress_file self.data = { 'created_at': datetime.now(timezone.utc).isoformat(), 'updated_at': None, 'summary': { 'total': 0, 'pending': 0, 'uploading': 0, 'uploaded': 0, 'dag_triggered': 0, 'running': 0, 'success': 0, 'failed': 0, 'skipped': 0, }, 'files': {}, } def load(self) -> bool: """Load progress from file. Returns True if loaded successfully.""" try: with open(self.progress_file) as f: self.data = json.load(f) return True except (FileNotFoundError, json.JSONDecodeError): return False def save(self): """Save progress to file.""" self.data['updated_at'] = datetime.now(timezone.utc).isoformat() self._recalculate_summary() with open(self.progress_file, 'w') as f: json.dump(self.data, f, indent=2, default=str) def _recalculate_summary(self): """Recalculate summary counts from file entries.""" summary = { 'total': 0, 'pending': 0, 'uploading': 0, 'uploaded': 0, 'dag_triggered': 0, 'running': 0, 'success': 0, 'failed': 0, 'skipped': 0, } for entry in self.data['files'].values(): summary['total'] += 1 status = entry.get('status', 'pending') if status in summary: summary[status] += 1 self.data['summary'] = summary def add_file(self, file_path: str, file_name: str, size_bytes: int): """Add a file entry with pending status.""" self.data['files'][file_name] = { 'file_path': file_path, 'file_name': file_name, 'size_bytes': size_bytes, 'status': 'pending', 'document_id': None, 'bucket': None, 'object_key': None, 'dag_run_id': None, 'error': None, 'upload_time': None, 'trigger_time': None, 'completion_time': None, } def update_status(self, file_name: str, status: str, **kwargs): """Update the status and extra fields for a file.""" if file_name in self.data['files']: self.data['files'][file_name]['status'] = status for key, value in kwargs.items(): self.data['files'][file_name][key] = value self.save() def get_files_by_status(self, *statuses: str) -> list[dict]: """Get all file entries matching any of the given statuses.""" return [ entry for entry in self.data['files'].values() if entry.get('status') in statuses ] def print_summary(self): """Print a formatted summary table.""" s = self.data['summary'] print('\n' + '=' * 60) print(' PROGRESS SUMMARY') print('=' * 60) print(f" Total files: {s['total']}") print(f" Pending: {s['pending']}") print(f" Uploading: {s['uploading']}") print(f" Uploaded: {s['uploaded']}") print(f" DAG triggered: {s['dag_triggered']}") print(f" Running: {s['running']}") print(f" Success: {s['success']}") print(f" Failed: {s['failed']}") print(f" Skipped: {s['skipped']}") print('=' * 60) # Show failed files failed = self.get_files_by_status('failed') if failed: print('\n FAILED FILES:') for entry in failed: print(f" - {entry['file_name']}: {entry.get('error', 'unknown error')}") print() def get_s3_client(): """Create S3 client for SeaweedFS.""" endpoint = os.getenv('SEAWEEDFS_ENDPOINT') if not endpoint: host = os.getenv('SEAWEEDFS_HOST', 'localhost') port = os.getenv('SEAWEEDFS_PORT', '8333') use_ssl = os.getenv('SEAWEEDFS_USE_SSL', '').lower() == 'true' or \ os.getenv('S3_USE_HTTPS', '').lower() == 'true' protocol = 'https' if use_ssl else 'http' endpoint = f"{protocol}://{host}:{port}" access_key = os.getenv('SEAWEEDFS_ACCESS_KEY_ID', 'dummy') secret_key = os.getenv('SEAWEEDFS_SECRET_ACCESS_KEY', 'dummy') return boto3.client( 's3', endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, config=Config(signature_version='s3v4'), region_name='us-east-1', ) def ensure_bucket_exists(s3_client, bucket_name: str) -> bool: """Ensure the bucket exists, create if not.""" try: s3_client.head_bucket(Bucket=bucket_name) return True except s3_client.exceptions.ClientError as e: error_code = e.response.get('Error', {}).get('Code', '') if error_code == '404': try: s3_client.create_bucket(Bucket=bucket_name) print(f" Created bucket: {bucket_name}") return True except Exception as create_err: print(f" Failed to create bucket: {create_err}") return False else: print(f" Error checking bucket: {e}") return False def get_content_type(file_path: Path) -> str: """Get content type for a file.""" ext = file_path.suffix.lower() return SUPPORTED_EXTENSIONS.get(ext, 'application/octet-stream') def is_supported_file(file_path: Path) -> bool: """Check if file is supported for extraction.""" return file_path.suffix.lower() in SUPPORTED_EXTENSIONS def collect_files(path: Path) -> list[Path]: """Collect all supported files from a directory recursively.""" if path.is_file(): return [path] if is_supported_file(path) else [] if path.is_dir(): files = [] for ext in SUPPORTED_EXTENSIONS: files.extend(path.rglob(f'*{ext}')) return sorted(files) return [] def upload_file(s3_client, file_path: Path, bucket: str, prefix: str) -> dict: """Upload a single file to SeaweedFS. Returns upload info.""" object_key = f"{prefix}/{file_path.name}" if prefix else file_path.name object_key = object_key.lstrip('/') content_type = get_content_type(file_path) document_id = str(uuid.uuid4()) with open(file_path, 'rb') as f: s3_client.put_object( Bucket=bucket, Key=object_key, Body=f, ContentType=content_type, Metadata={'document_id': document_id}, ) return { 'document_id': document_id, 'bucket': bucket, 'object_key': object_key, 'content_type': content_type, } def get_airflow_token(airflow_url: str, username: str, password: str) -> str | None: """Get JWT token from Airflow 3.x API.""" url = f"{airflow_url}/auth/token" try: response = requests.post( url, json={'username': username, 'password': password}, headers={'Content-Type': 'application/json'}, timeout=30, ) if response.status_code in (200, 201): return response.json().get('access_token') else: print(f" Failed to get Airflow token: {response.status_code} - {response.text}") return None except Exception as e: print(f" Failed to get Airflow token: {e}") return None def trigger_dag(upload_info: dict, airflow_url: str, token: str) -> str | None: """Trigger the document_extraction_dag. Returns dag_run_id or None.""" dag_id = 'document_extraction_dag' url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns" payload = { 'logical_date': None, 'conf': { 'document_id': upload_info['document_id'], 'bucket': upload_info['bucket'], 'object_key': upload_info['object_key'], }, } try: response = requests.post( url, json=payload, headers={ 'Content-Type': 'application/json', 'Authorization': f'Bearer {token}', }, timeout=30, ) if response.status_code in (200, 201): return response.json().get('dag_run_id') else: print(f" Failed to trigger DAG: {response.status_code} - {response.text}") return None except Exception as e: print(f" Failed to trigger DAG: {e}") return None def check_dag_run_status(dag_run_id: str, airflow_url: str, token: str) -> str | None: """Check status of a DAG run. Returns state string or None on error.""" dag_id = 'document_extraction_dag' url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}" try: response = requests.get( url, headers={'Authorization': f'Bearer {token}'}, timeout=30, ) if response.status_code == 200: return response.json().get('state') else: return None except Exception: return None def phase_upload(s3_client, tracker: ProgressTracker, bucket: str, prefix: str, dry_run: bool): """Phase 1: Upload all pending files to SeaweedFS.""" pending = tracker.get_files_by_status('pending') if not pending: print(' No pending files to upload.') return print(f'\n Uploading {len(pending)} file(s)...') for entry in pending: file_path = Path(entry['file_path']) file_name = entry['file_name'] if dry_run: print(f" [DRY RUN] Would upload: {file_name}") tracker.update_status(file_name, 'uploaded', document_id=str(uuid.uuid4()), bucket=bucket, object_key=f"{prefix}/{file_name}".lstrip('/')) continue tracker.update_status(file_name, 'uploading') try: info = upload_file(s3_client, file_path, bucket, prefix) tracker.update_status( file_name, 'uploaded', document_id=info['document_id'], bucket=info['bucket'], object_key=info['object_key'], upload_time=datetime.now(timezone.utc).isoformat(), ) size = entry['size_bytes'] print(f" Uploaded: {file_name} ({size:,} bytes)") except Exception as e: tracker.update_status(file_name, 'failed', error=f"Upload error: {e}") print(f" FAILED: {file_name} - {e}") def phase_trigger(tracker: ProgressTracker, airflow_url: str, token: str, dry_run: bool): """Phase 2: Trigger DAG runs for all uploaded files.""" uploaded = tracker.get_files_by_status('uploaded') if not uploaded: print(' No uploaded files to trigger.') return print(f'\n Triggering DAG for {len(uploaded)} file(s)...') for entry in uploaded: file_name = entry['file_name'] if dry_run: print(f" [DRY RUN] Would trigger DAG for: {file_name}") tracker.update_status(file_name, 'success', dag_run_id='dry-run') continue upload_info = { 'document_id': entry['document_id'], 'bucket': entry['bucket'], 'object_key': entry['object_key'], } dag_run_id = trigger_dag(upload_info, airflow_url, token) if dag_run_id: tracker.update_status( file_name, 'dag_triggered', dag_run_id=dag_run_id, trigger_time=datetime.now(timezone.utc).isoformat(), ) print(f" Triggered: {file_name} -> {dag_run_id}") else: tracker.update_status(file_name, 'failed', error='Failed to trigger DAG') print(f" FAILED: {file_name} - could not trigger DAG") def phase_monitor(tracker: ProgressTracker, airflow_url: str, token: str, poll_interval: int): """Phase 3: Monitor DAG runs until all reach terminal state.""" active = tracker.get_files_by_status('dag_triggered', 'running') if not active: print(' No active DAG runs to monitor.') return print(f'\n Monitoring {len(active)} DAG run(s) (poll every {poll_interval}s)...') while True: active = tracker.get_files_by_status('dag_triggered', 'running') if not active: break for entry in active: file_name = entry['file_name'] dag_run_id = entry['dag_run_id'] state = check_dag_run_status(dag_run_id, airflow_url, token) if state is None: continue # Transient error, retry next cycle if state in TERMINAL_STATES: tracker.update_status( file_name, state, completion_time=datetime.now(timezone.utc).isoformat(), error='DAG run failed' if state == 'failed' else None, ) icon = 'OK' if state == 'success' else 'FAIL' print(f" [{icon}] {file_name} -> {state}") elif state in RUNNING_STATES: if entry['status'] != 'running': tracker.update_status(file_name, 'running') # Recheck still_active = tracker.get_files_by_status('dag_triggered', 'running') if not still_active: break time.sleep(poll_interval) tracker.save() def main(): parser = argparse.ArgumentParser( description='Batch upload documents to SeaweedFS, trigger Airflow DAG, and monitor progress', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) parser.add_argument( 'path', type=Path, nargs='?', help='Folder or file to upload', ) parser.add_argument( '--bucket', '-b', default='documents', help='S3 bucket name (default: documents)', ) parser.add_argument( '--prefix', '-p', default='', help='Key prefix for uploaded files', ) parser.add_argument( '--no-trigger', action='store_true', help='Upload only, do not trigger DAG', ) parser.add_argument( '--dry-run', '-n', action='store_true', help='Show what would be done without actually doing it', ) parser.add_argument( '--resume', '-r', type=str, default=None, help='Resume from a previous progress JSON file', ) parser.add_argument( '--poll-interval', type=int, default=15, help='Seconds between DAG status checks (default: 15)', ) parser.add_argument( '--progress-file', type=str, default=None, help='Custom path for the progress JSON file', ) parser.add_argument( '--airflow-url', default=os.getenv('AIRFLOW_URL', 'http://localhost:8081'), help='Airflow API URL (default: http://localhost:8081)', ) parser.add_argument( '--airflow-user', default=os.getenv('_AIRFLOW_WWW_USER_USERNAME', 'admin'), help='Airflow username', ) parser.add_argument( '--airflow-password', default=os.getenv('_AIRFLOW_WWW_USER_PASSWORD', 'admin'), help='Airflow password', ) args = parser.parse_args() # --- Resume mode --- if args.resume: print(f"Resuming from: {args.resume}") tracker = ProgressTracker(args.resume) if not tracker.load(): print(f"Error: cannot load progress file: {args.resume}") sys.exit(1) tracker.print_summary() else: # --- New upload mode --- if not args.path: parser.error('path is required unless --resume is used') if not args.path.exists(): print(f"Error: path does not exist: {args.path}") sys.exit(1) files = collect_files(args.path) if not files: print(f"No supported files found in: {args.path}") print(f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS.keys()))}") sys.exit(1) # Create progress file timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') progress_file = args.progress_file or f"progress_{timestamp}.json" tracker = ProgressTracker(progress_file) for f in files: tracker.add_file(str(f), f.name, f.stat().st_size) tracker.save() print(f"Found {len(files)} supported file(s)") print(f"Progress file: {progress_file}") print(f"Bucket: {args.bucket}") if args.prefix: print(f"Prefix: {args.prefix}") # --- Phase 1: Upload --- s3_client = None if not args.dry_run: s3_client = get_s3_client() if not ensure_bucket_exists(s3_client, args.bucket): print("Failed to ensure bucket exists") sys.exit(1) phase_upload(s3_client, tracker, args.bucket, args.prefix, args.dry_run) tracker.print_summary() if args.no_trigger: print("\n--no-trigger set. Stopping after upload.") return # --- Phase 2: Trigger DAG --- token = None if not args.dry_run: token = get_airflow_token(args.airflow_url, args.airflow_user, args.airflow_password) if not token: print("Failed to get Airflow token. Files are uploaded but DAG not triggered.") print(f"Re-run with: --resume {tracker.progress_file}") sys.exit(1) phase_trigger(tracker, args.airflow_url, token, args.dry_run) tracker.print_summary() if args.dry_run: print("\n[DRY RUN] Complete.") return # --- Phase 3: Monitor --- print('\n--- Monitoring DAG runs (Ctrl+C to stop, resume later with --resume) ---') try: phase_monitor(tracker, args.airflow_url, token, args.poll_interval) except KeyboardInterrupt: print(f"\n\nInterrupted. Resume with: python {sys.argv[0]} --resume {tracker.progress_file}") tracker.save() sys.exit(0) tracker.print_summary() print(f"\nProgress saved to: {tracker.progress_file}") if __name__ == '__main__': main()