# Story 3.4a: Extraction Pipeline Orchestrator

Status: review

## Story

As a **system**,
I want **a pipeline that orchestrates the extraction workflow from uploaded document to structured records**,
So that **all extraction components work together seamlessly**.

## Acceptance Criteria

1. **AC3.4a.1:** Load and validate Excel workbook
   - Accept document_id as input parameter
   - Load file from storage path using document metadata
   - Validate file format (openpyxl can load)
   - Update document status to "processing"
   - Handle file not found, invalid format errors

2. **AC3.4a.2:** Orchestrate complete extraction workflow
   - For each sheet in workbook:
     - Detect table boundaries using `TableDetector` (Story 3.1)
     - Extract multi-level headers using `HeaderExtractor` (Story 3.2)
     - Resolve merged cells using `MergedCellResolver` (Story 3.3)
     - Build structured records using `RecordBuilder` (Story 3.4)
   - Process all sheets, collecting results
   - Continue processing on non-critical errors (log warnings)

3. **AC3.4a.3:** Update document status throughout pipeline
   - Start: status = "processing"
   - Success: status = "completed", record_count = total records, processed_at = timestamp
   - Failure: status = "failed", error_message = error details
   - Use database transactions to ensure consistency

4. **AC3.4a.4:** Store extraction results in database
   - Save all `ExtractedRecord` objects to `extracted_records` table
   - Include document_id, sheet_name, row_number, col_range, content, headers
   - Use batch insert for performance (chunks of 100-500 records)
   - Rollback on critical errors, keep partial results on warnings

5. **AC3.4a.5:** Track progress in Redis (Story 2.5 integration)
   - Update progress: `processing:documents:{doc_id}`
   - Format: `{"status": "processing", "current_sheet": 2, "total_sheets": 5, "records": 450}`
   - TTL: 1 hour after completion
   - Allow status endpoint to query progress

6. **AC3.4a.6:** Provide comprehensive structured logging
   - Pipeline started: document_id, filename, sheet_count
   - Sheet processed: sheet_name, tables_detected, records_extracted, duration_ms
   - Pipeline completed: document_id, sheets_processed, total_records, total_duration_ms
   - Errors: error type, context, recovery action
   - Use structlog with all relevant context

7. **AC3.4a.7:** Return extraction summary
   - document_id: UUID
   - sheets_processed: int
   - tables_detected: int
   - records_extracted: int
   - duration_ms: int
   - Success indicator and any warnings

8. **AC3.4a.8:** Handle edge cases gracefully
   - Empty workbook → completed with record_count=0
   - No tables detected in sheet → skip sheet, log warning
   - Partial sheet failures → continue with other sheets
   - Invalid Excel format → fail fast with clear error
   - File access errors → retry once, then fail

## Tasks / Subtasks

- [x] **Task 1: Create ExtractionPipeline class** (AC: 3.4a.1, 3.4a.2)
  - [x] Create `src/extraction/pipeline.py`
  - [x] Implement `ExtractionPipeline` class
  - [x] Implement `process_document(document_id: UUID) -> ExtractionSummary` method
  - [x] Load workbook from file storage
  - [x] Validate Excel format with openpyxl
  - [x] Orchestrate extraction workflow for all sheets

- [x] **Task 2: Integrate extraction components** (AC: 3.4a.2)
  - [x] Import and instantiate TableDetector, HeaderExtractor, MergedCellResolver, RecordBuilder
  - [x] For each sheet: detect tables → extract headers → resolve merges → build records
  - [x] Pass outputs between components (TableBoundary, HeaderInfo dict, resolved_grid, MergeMetadata)
  - [x] Collect all ExtractedRecord objects from all sheets
  - [x] Handle empty sheets and sheets with no tables

- [x] **Task 3: Implement document status updates** (AC: 3.4a.3)
  - [x] Create database session for document updates
  - [x] Update status to "processing" at pipeline start
  - [x] Update status to "completed" or "failed" at pipeline end
  - [x] Set processed_at timestamp, sheet_count, record_count
  - [x] Use transaction management (commit on success, rollback on failure)

- [x] **Task 4: Store extraction results in database** (AC: 3.4a.4)
  - [x] Create SQLAlchemy models for extracted_records table (if not exists)
  - [x] Convert ExtractedRecord objects to ORM models
  - [x] Use batch insert with session.bulk_insert_mappings() for performance
  - [x] Batch size: 500 records per insert
  - [x] Handle database constraints and integrity errors

- [x] **Task 5: Integrate Redis progress tracking** (AC: 3.4a.5)
  - [x] Create Redis client instance
  - [x] Update progress after each sheet: current_sheet, total_sheets, records_so_far
  - [x] Set progress key with 1-hour TTL
  - [x] Handle Redis connection errors gracefully (don't fail pipeline)

- [x] **Task 6: Add comprehensive structured logging** (AC: 3.4a.6)
  - [x] Log pipeline_started with document_id, filename, sheet_count
  - [x] Log sheet_processed with sheet_name, tables, records, duration_ms
  - [x] Log pipeline_completed with totals and duration
  - [x] Log errors with full context (error type, document, sheet, recovery)
  - [x] Use structlog.get_logger(__name__)

- [x] **Task 7: Create ExtractionSummary model** (AC: 3.4a.7)
  - [x] Add `ExtractionSummary` dataclass to models.py or pipeline.py
  - [x] Fields: document_id, sheets_processed, tables_detected, records_extracted, duration_ms, warnings
  - [x] Return from process_document() method

- [x] **Task 8: Implement error handling** (AC: 3.4a.8)
  - [x] Handle FileNotFoundError → status: failed, error: "FILE_NOT_FOUND"
  - [x] Handle openpyxl.utils.exceptions → status: failed, error: "INVALID_EXCEL_FORMAT"
  - [x] Handle empty workbook → status: completed, record_count: 0
  - [x] Handle partial failures → log warnings, continue processing
  - [x] Use ExtractionError for all extraction-specific errors

- [x] **Task 9: Create Celery task integration** (AC: 3.4a.2)
  - [x] Create `src/workers/extraction_tasks.py` if not exists
  - [x] Create `extract_document` Celery task
  - [x] Task signature: `@celery_app.task(name="extract_document", bind=True)`
  - [x] Call ExtractionPipeline.process_document(document_id)
  - [x] Handle task retries on transient failures
  - [x] Update task state for progress tracking

- [x] **Task 10: Unit tests** (AC: All)
  - [x] Create `tests/extraction/test_pipeline.py`
  - [x] Test ExtractionPipeline.process_document() with mocked components
  - [x] Test document status updates (processing, completed, failed)
  - [x] Test extraction summary generation
  - [x] Test error handling (invalid file, empty workbook, partial failures)
  - [x] Test progress tracking updates

- [x] **Task 11: Integration tests** (AC: All)
  - [x] Add integration tests to `tests/extraction/test_integration.py`
  - [x] Test full pipeline with real Excel files from fixtures
  - [x] Test multiple sheets processing
  - [x] Test database storage of extracted records
  - [x] Test Redis progress updates
  - [x] Verify end-to-end: file → records in database

## Dev Notes

### Architecture Patterns and Constraints

From architecture.md:
- **Module location:** `src/extraction/pipeline.py` (architecture.md:62)
- **Naming conventions:**
  - Module: `pipeline.py` (snake_case)
  - Class: `ExtractionPipeline` (PascalCase)
  - Functions: `process_document()`, `_process_sheet()` (snake_case)
- **Error handling:** Use `ExtractionError(DSOLError)` with message/details
- **Logging:** structlog with context (document_id, filename, sheet_name, duration_ms)
- **Type hints:** Full type annotations required
- **Celery integration:** Worker tasks in `src/workers/extraction_tasks.py`
- **Database:** Use SQLAlchemy session management, transactions, batch inserts
- **Redis:** Progress tracking with TTL, graceful degradation on errors

From epics.md Story 3.4a:
- **Prerequisites:** Stories 3.1, 3.2, 3.3, 3.4 (all in review status)
- **Integration Points:** Document upload (2.2), Status tracking (2.5), All extraction components (3.1-3.4)
- **Transaction Management:** Use database transactions for consistency
- **Progress Tracking:** Redis keys format `processing:documents:{doc_id}`
- **Error Categories:** Invalid format, file not found, empty workbook, partial failures

### Learnings from Previous Story

**From Story 3.4: Record Builder (Status: review)**

- **New Service Created**: `RecordBuilder` class available at `src/extraction/record_builder.py`
  - Use `build_records(resolved_grid, headers, merge_metadata, table_boundary, document_id, filename)` method
  - Returns list[ExtractedRecord] objects ready for database storage
  - Handles data type conversion, row filtering, merge metadata integration automatically

- **New Model Created**: `ExtractedRecord` dataclass in `src/extraction/models.py` (lines 151-215)
  - Fields: content (dict), headers (list[str]), _source (dict), _merge_metadata (optional dict)
  - Has to_dict() method for JSON serialization
  - Comprehensive validation built-in

- **Extraction Pipeline Components Available:**
  - `TableDetector` (Story 3.1): `detect_tables(workbook)` → list[TableBoundary]
  - `HeaderExtractor` (Story 3.2): `extract_headers(sheet, boundary)` → dict[str, HeaderInfo]
  - `MergedCellResolver` (Story 3.3): `resolve_merged_cells(sheet, boundary)` → (resolved_grid, metadata)
  - `RecordBuilder` (Story 3.4): `build_records(...)` → list[ExtractedRecord]

- **Testing Pattern Established:**
  - All extraction components have 100% test pass rate
  - Unit tests + integration tests with real Excel files
  - Sample Excel files in `tests/fixtures/sample_excel/`
  - Follow same testing structure for pipeline tests

- **Data Flow Pattern:**
  ```
  Workbook → TableDetector → TableBoundary[]
  For each TableBoundary:
    Sheet + Boundary → HeaderExtractor → HeaderInfo dict
    Sheet + Boundary → MergedCellResolver → (resolved_grid, MergeMetadata[])
    All inputs → RecordBuilder → ExtractedRecord[]
  ```

- **Error Handling Pattern:**
  - All components use `ExtractionError` for extraction-specific errors
  - All components use structured logging with context
  - All components have comprehensive error handling

- **Known Patterns to Reuse:**
  - Structured logging with context: `logger.info("event_name", key=value, ...)`
  - ExtractionError for failures: `ExtractionError(message="...", details={...})`
  - Full type annotations throughout
  - Dataclass patterns with validation and to_dict()

[Source: docs/sprint-artifacts/3-4-record-builder.md#Dev-Agent-Record]

### Project Structure Notes

Expected module structure after Story 3.4a:
```
src/extraction/
├── __init__.py
├── models.py              # ✓ Stories 3.1-3.4
├── table_detector.py      # ✓ Story 3.1
├── header_extractor.py    # ✓ Story 3.2
├── merged_cell_resolver.py # ✓ Story 3.3
├── record_builder.py      # ✓ Story 3.4
└── pipeline.py            # NEW: Story 3.4a - ExtractionPipeline orchestrator

src/workers/
├── __init__.py
├── celery_app.py          # ✓ Existing Celery config
└── extraction_tasks.py    # NEW: Story 3.4a - Celery task integration

tests/extraction/
├── __init__.py
├── test_table_detector.py         # ✓ Story 3.1
├── test_header_extractor.py       # ✓ Story 3.2
├── test_merged_cell_resolver.py   # ✓ Story 3.3
├── test_record_builder.py         # ✓ Story 3.4
├── test_pipeline.py               # NEW: Story 3.4a unit tests
└── test_integration.py            # MODIFY: Add pipeline integration tests
```

### Prerequisites and Dependencies

**Prerequisites:**
- Story 3.1 (Table Boundary Detection) - ✓ REVIEW status
- Story 3.2 (Multi-Level Header Extraction) - ✓ REVIEW status
- Story 3.3 (Merged Cell Resolution) - ✓ REVIEW status
- Story 3.4 (Record Builder) - ✓ REVIEW status

**Dependencies:**
- openpyxl 3.1.5: Excel file reading
- structlog: Structured logging
- SQLAlchemy: Database ORM and transactions
- Redis: Progress tracking (graceful degradation if unavailable)
- Celery: Async task processing
- All extraction component classes from Stories 3.1-3.4

**Database Tables:**
- `documents`: Update status, record_count, processed_at
- `extracted_records`: Insert all ExtractedRecord objects

**Redis Keys:**
- `processing:documents:{document_id}`: Progress tracking with 1-hour TTL

### Implementation Strategy

**Pipeline Orchestration Flow:**

```python
# 1. Initialize pipeline
pipeline = ExtractionPipeline(
    db_session=session,
    redis_client=redis_client,  # Optional
    storage_path=settings.storage_path
)

# 2. Process document
summary = pipeline.process_document(document_id)

# 3. Internal flow:
def process_document(document_id: UUID) -> ExtractionSummary:
    # Load document metadata
    document = db_session.query(Document).filter_by(id=document_id).first()

    # Update status to processing
    document.status = "processing"
    db_session.commit()

    # Load workbook
    workbook = openpyxl.load_workbook(document.file_path, data_only=True)

    # Initialize extraction components
    table_detector = TableDetector()
    header_extractor = HeaderExtractor()
    merged_cell_resolver = MergedCellResolver()
    record_builder = RecordBuilder()

    # Process all sheets
    all_records = []
    for sheet in workbook.worksheets:
        # Detect tables
        boundaries = table_detector.detect_tables_in_sheet(sheet)

        for boundary in boundaries:
            # Extract headers
            headers = header_extractor.extract_headers(sheet, boundary)

            # Resolve merged cells
            resolved_grid, merge_metadata = merged_cell_resolver.resolve_merged_cells(
                sheet, boundary
            )

            # Build records
            records = record_builder.build_records(
                resolved_grid=resolved_grid,
                headers=headers,
                merge_metadata=merge_metadata,
                table_boundary=boundary,
                document_id=document_id,
                filename=document.filename
            )

            all_records.extend(records)

        # Update progress in Redis
        update_progress(document_id, current_sheet, total_sheets, len(all_records))

    # Store records in database (batch insert)
    store_records_batch(all_records, db_session)

    # Update document status to completed
    document.status = "completed"
    document.record_count = len(all_records)
    document.processed_at = datetime.now(timezone.utc)
    db_session.commit()

    # Return summary
    return ExtractionSummary(
        document_id=document_id,
        sheets_processed=len(workbook.worksheets),
        tables_detected=total_tables,
        records_extracted=len(all_records),
        duration_ms=duration
    )
```

**Celery Task Integration:**

```python
# src/workers/extraction_tasks.py

@celery_app.task(name="extract_document", bind=True)
def extract_document(self, document_id: str):
    """Async task to extract document using ExtractionPipeline."""
    from src.extraction.pipeline import ExtractionPipeline
    from src.db.session import SessionLocal

    session = SessionLocal()
    try:
        pipeline = ExtractionPipeline(db_session=session)
        summary = pipeline.process_document(UUID(document_id))
        return summary.to_dict()
    except Exception as e:
        logger.exception("extraction_task_failed", document_id=document_id, error=str(e))
        raise
    finally:
        session.close()
```

### References

- [Source: docs/epics.md#Story-3.4a] - Original story definition with acceptance criteria
- [Source: docs/architecture.md#Project-Structure] - Module location (line 62)
- [Source: docs/architecture.md#Implementation-Patterns] - Error handling, logging, type hints
- [Source: docs/sprint-artifacts/3-1-table-boundary-detection.md] - TableDetector API
- [Source: docs/sprint-artifacts/3-2-multi-level-header-extraction.md] - HeaderExtractor API
- [Source: docs/sprint-artifacts/3-3-merged-cell-resolution.md] - MergedCellResolver API
- [Source: docs/sprint-artifacts/3-4-record-builder.md] - RecordBuilder API, ExtractedRecord model

## Dev Agent Record

### Context Reference

No context file - proceeded with story file and codebase exploration

### Agent Model Used

Claude Sonnet 4.5 (claude-sonnet-4-5-20250929)

### Debug Log References

All tests passed (130/130 extraction tests):
- 22 pipeline unit tests
- 6 pipeline integration tests
- All existing tests remain passing

### Completion Notes List

**Implementation Summary:**
- Created ExtractionPipeline class (src/extraction/pipeline.py, 618 lines)
- Implemented complete orchestration workflow:
  - Document loading and validation
  - Sheet-by-sheet table processing
  - Component integration: TableDetector → HeaderExtractor → MergedCellResolver → RecordBuilder
  - Database storage with batch inserts (500 records/batch)
  - Redis progress tracking with graceful degradation
  - Comprehensive structured logging
  - Robust error handling for all edge cases
- Created Celery task integration (src/workers/extraction_tasks.py, 81 lines)
- Wrote 22 unit tests covering all pipeline methods
- Wrote 6 integration tests for end-to-end validation
- All 130 extraction tests pass (100% pass rate)

**Key Implementation Details:**
- ExtractionSummary dataclass returns: document_id, sheets_processed, tables_detected, records_extracted, duration_ms, warnings
- Pipeline handles empty workbooks, missing tables, partial failures gracefully
- Document status lifecycle: pending → processing → completed/failed
- Progress tracking format: `processing:documents:{doc_id}` with 1-hour TTL
- Error handling: FileNotFoundError, InvalidFileException, ExtractionError
- Logging: pipeline_started, sheet_processed, pipeline_completed, errors with full context

**Integration Points:**
- Uses existing TableDetector, HeaderExtractor, MergedCellResolver, RecordBuilder
- Integrates with Document ORM model for status updates
- Stores results in ExtractedRecord ORM table
- Celery task queue: "extraction" with retry support (max 3 retries)

### File List

| Status | File Path |
|--------|-----------|
| NEW | src/extraction/pipeline.py (618 lines - ExtractionPipeline class, ExtractionSummary dataclass) |
| NEW | src/workers/extraction_tasks.py (81 lines - extract_document Celery task) |
| MODIFIED | src/workers/celery_app.py (added extraction_tasks to includes, added extraction queue) |
| NEW | tests/extraction/test_pipeline.py (530 lines - 22 unit tests) |
| MODIFIED | tests/extraction/test_integration.py (added 6 pipeline integration tests, 174 lines added) |