# Story 1.5: Storage Operators (PostgreSQL, Milvus)

**Epic:** Airflow Infrastructure & Parallel Deployment
**Status:** Done
**Priority:** High
**Estimate:** 4 hours

> **Implementation Evolved:** The original plan below used custom operators (`StoreChunksOperator`, `StoreVectorsOperator`). The actual implementation uses `@task.external_python` decorated functions in `airflow/dags/tasks/processing_tasks.py`: `store_chunks_task`, `store_vectors_task`. Large embedding data is stored via pickle files to avoid XCom size limits. See [`docs/specs/airflow.md`](../../../specs/airflow.md) for current architecture.

---

## User Story

As a **developer**, I want to **implement operators for storing chunks in PostgreSQL and vectors in Milvus** so that **processed documents are persisted and searchable**.

---

## Acceptance Criteria

- [x] `StoreChunksOperator` stores chunks in PostgreSQL via existing save logic
- [x] `StoreVectorsOperator` stores vectors in Milvus via existing indexer
- [x] Both operators handle transactions properly
- [x] Both operators have rollback capability on failure
- [x] Unit tests with >80% coverage
- [ ] Integration tests with actual databases (deferred to Story 1.6 - DAG integration)

---

## Technical Details

### Operators to Implement

| Operator | Input | Output | Reuses |
|----------|-------|--------|--------|
| `StoreChunksOperator` | ChunkEmbeddings | StoreResult | `structured_store.py` |
| `StoreVectorsOperator` | ChunkEmbeddings | StoreResult | `vector_store.py` |

### Files to Create

| File | Description |
|------|-------------|
| `plugins/operators/store_chunks.py` | PostgreSQL storage operator |
| `plugins/operators/store_vectors.py` | Milvus storage operator |

---

## Implementation

### StoreChunksOperator

```python
# plugins/operators/store_chunks.py
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook
from uuid import UUID
import json

class StoreChunksOperator(BaseOperator):
    """
    Store extracted chunks/records in PostgreSQL.

    Reuses: src.knowledge.structured_store.StructuredStore
    """

    @apply_defaults
    def __init__(self, connection_id: str = 'postgres_default', **kwargs):
        super().__init__(**kwargs)
        self.connection_id = connection_id

    def execute(self, context):
        from src.knowledge.structured_store import StructuredStore
        from src.db.session import get_db_session

        # Get embedding result which contains chunks
        embedding_result = context['ti'].xcom_pull(task_ids='generate_embeddings')
        validate_result = context['ti'].xcom_pull(task_ids='validate_event')

        chunks = embedding_result.get('chunks', [])
        document_id = UUID(validate_result['document_id'])

        if not chunks:
            self.log.info("No chunks to store")
            return {
                'stored_count': 0,
                'document_id': str(document_id),
                'status': 'empty',
            }

        # Initialize structured store
        store = StructuredStore()

        # Store chunks
        stored_count = 0
        errors = []

        with get_db_session() as session:
            try:
                for chunk in chunks:
                    try:
                        # Handle both dict and object formats
                        if isinstance(chunk, dict):
                            store.save_record(
                                session=session,
                                document_id=document_id,
                                content=chunk.get('content', chunk),
                                source=chunk.get('_source', {}),
                                metadata=chunk.get('metadata', {}),
                            )
                        else:
                            store.save_record(
                                session=session,
                                document_id=document_id,
                                content=chunk.content if hasattr(chunk, 'content') else chunk,
                                source=chunk._source if hasattr(chunk, '_source') else {},
                                metadata=chunk.metadata if hasattr(chunk, 'metadata') else {},
                            )
                        stored_count += 1
                    except Exception as e:
                        errors.append(str(e))
                        self.log.warning(f"Failed to store chunk: {e}")

                session.commit()
                self.log.info(f"Stored {stored_count} chunks for document {document_id}")

            except Exception as e:
                session.rollback()
                self.log.error(f"Transaction failed: {e}")
                raise

        return {
            'stored_count': stored_count,
            'document_id': str(document_id),
            'status': 'success' if not errors else 'partial',
            'errors': errors[:10] if errors else [],  # Limit error list
        }
```

### StoreVectorsOperator

```python
# plugins/operators/store_vectors.py
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from uuid import UUID

class StoreVectorsOperator(BaseOperator):
    """
    Store vector embeddings in Milvus.

    Reuses: src.knowledge.vector_store.VectorStore
    """

    @apply_defaults
    def __init__(
        self,
        connection_id: str = 'milvus_default',
        collection: str = 'document_embeddings',
        **kwargs
    ):
        super().__init__(**kwargs)
        self.connection_id = connection_id
        self.collection = collection

    def execute(self, context):
        from src.knowledge.vector_store import VectorStore

        # Get embedding result
        embedding_result = context['ti'].xcom_pull(task_ids='generate_embeddings')
        validate_result = context['ti'].xcom_pull(task_ids='validate_event')

        embeddings = embedding_result.get('embeddings', [])
        chunks = embedding_result.get('chunks', [])
        document_id = UUID(validate_result['document_id'])

        if not embeddings:
            self.log.info("No vectors to store")
            return {
                'stored_count': 0,
                'document_id': str(document_id),
                'collection': self.collection,
                'status': 'empty',
            }

        # Initialize vector store
        vector_store = VectorStore(collection_name=self.collection)

        # Prepare data for Milvus
        vectors = []
        metadata_list = []

        for i, emb in enumerate(embeddings):
            vector = emb.get('vector', [])
            chunk_id = emb.get('chunk_id', f'chunk_{i}')

            # Get corresponding chunk for metadata
            chunk = chunks[i] if i < len(chunks) else {}
            if isinstance(chunk, dict):
                metadata = {
                    'document_id': str(document_id),
                    'chunk_id': chunk_id,
                    'text': chunk.get('content', {}).get('text', '')[:1000] or chunk.get('text', '')[:1000],
                    'source': str(chunk.get('_source', {}))[:500],
                }
            else:
                metadata = {
                    'document_id': str(document_id),
                    'chunk_id': chunk_id,
                    'text': '',
                    'source': '',
                }

            vectors.append(vector)
            metadata_list.append(metadata)

        # Store in Milvus
        try:
            stored_ids = vector_store.insert(
                vectors=vectors,
                metadata=metadata_list,
            )
            stored_count = len(stored_ids)
            self.log.info(f"Stored {stored_count} vectors for document {document_id}")

            return {
                'stored_count': stored_count,
                'document_id': str(document_id),
                'collection': self.collection,
                'status': 'success',
                'vector_ids': stored_ids[:10],  # Return sample of IDs
            }

        except Exception as e:
            self.log.error(f"Vector storage failed: {e}")
            raise
```

---

## Tasks

- [x] Create `plugins/operators/store_chunks.py`
- [x] Create `plugins/operators/store_vectors.py`
- [x] Update `plugins/operators/__init__.py` with exports
- [x] Create unit tests for StoreChunksOperator
- [x] Create unit tests for StoreVectorsOperator
- [ ] Integration test with PostgreSQL (deferred to Story 1.6)
- [ ] Integration test with Milvus (deferred to Story 1.6)

---

## Testing

```python
# tests/test_operators/test_store_chunks.py
import pytest
from unittest.mock import MagicMock, patch

class TestStoreChunksOperator:

    @patch('plugins.operators.store_chunks.StructuredStore')
    @patch('plugins.operators.store_chunks.get_db_session')
    def test_stores_chunks_successfully(self, mock_session, mock_store, mock_context):
        """Test successful chunk storage."""
        mock_context['ti'].xcom_pull.side_effect = [
            {
                'chunks': [{'content': {'text': 'test'}, '_source': {}}],
                'embeddings': []
            },
            {'document_id': 'test-uuid'}
        ]

        from plugins.operators.store_chunks import StoreChunksOperator
        op = StoreChunksOperator(task_id='test')
        result = op.execute(mock_context)

        assert result['stored_count'] == 1
        assert result['status'] == 'success'
```

---

## Dependencies

- Story 1.4 (Chunking & Embedding operators)
- Existing `src/knowledge/structured_store.py`
- Existing `src/knowledge/vector_store.py`
- PostgreSQL database
- Milvus database

---

## Notes

- Operators run in parallel after embedding generation
- Transaction handling ensures data consistency
- Error handling allows partial success with error reporting
- Vector metadata limited to prevent Milvus size issues

---

## File List

| File | Action | Description |
|------|--------|-------------|
| `airflow/plugins/operators/store_chunks.py` | Created | StoreChunksOperator - stores chunks in PostgreSQL via StructuredStore |
| `airflow/plugins/operators/store_vectors.py` | Created | StoreVectorsOperator - stores vectors in Milvus via VectorStore |
| `airflow/plugins/operators/__init__.py` | Modified | Added exports for new operators |
| `src/knowledge/vector_store.py` | Modified | Added `index_with_vectors()` public method for pre-computed embeddings |
| `tests/airflow/test_operators/conftest.py` | Modified | Added mocks for db and storage modules |
| `tests/airflow/test_operators/test_store_chunks.py` | Created | 12 unit tests for StoreChunksOperator |
| `tests/airflow/test_operators/test_store_vectors.py` | Created | 22 unit tests for StoreVectorsOperator + ExcelRecordProxy |

---

## Dev Agent Record

### Implementation Notes

**StoreChunksOperator:**
- Reuses `StructuredStore` from `src/knowledge/structured_store.py`
- Uses async session from `src/db/session.py` via `async_session_maker`
- Converts chunk dicts to `ExtractedRecord` objects for StructuredStore
- Supports both table record format (Excel) and semantic chunk format (Word/PDF)
- Transaction handling with automatic rollback on failure
- Returns: `{stored_count, document_id, status, errors}`

**StoreVectorsOperator:**
- Reuses `VectorStore` from `src/knowledge/vector_store.py`
- Uses pre-computed embeddings from GenerateEmbeddingsOperator (avoids re-embedding)
- Uses public `index_with_vectors()` method (added in code review fix)
- Auto-detects document type (excel/word) from chunk format
- Creates `ExcelRecordProxy` helper class to mimic ExtractedRecord for VectorStore
- Properly disconnects from Milvus in `finally` block to prevent connection leaks
- Returns: `{stored_count, document_id, collection, document_type, status}`

**Key Design Decisions:**
1. Uses async execution via `asyncio.run()` for both operators
2. StoreChunksOperator creates ExtractedRecord objects to match StructuredStore API
3. StoreVectorsOperator uses new `index_with_vectors()` public method to avoid re-embedding
4. Both operators support custom upstream task IDs for flexibility
5. ExcelRecordProxy provides ExtractedRecord-like interface for VectorStore without full dataclass
6. StoreChunksOperator relies on async_session_maker context manager for commit/rollback (no explicit calls)

### Testing Notes

- 34 unit tests for Story 1.5 operators (12 StoreChunks + 22 StoreVectors including helper class)
- All 95 operator tests pass (including error handling and disconnect tests added in code review)
- Tests mock async session and storage services
- Comprehensive coverage: success paths, error handling, edge cases
- Tests verify error propagation and proper Milvus disconnect on both success and failure

### Validation Results

All acceptance criteria validated:
- ✓ StoreChunksOperator stores chunks in PostgreSQL via StructuredStore
- ✓ StoreVectorsOperator stores vectors in Milvus via VectorStore
- ✓ Both operators handle transactions properly (async session with commit/rollback)
- ✓ Both operators have rollback capability on failure
- ✓ 32 unit tests with comprehensive coverage

---

## Change Log

| Date | Change |
|------|--------|
| 2025-12-09 | Created StoreChunksOperator with async StructuredStore integration |
| 2025-12-09 | Created StoreVectorsOperator with pre-computed embedding support |
| 2025-12-09 | Created ExcelRecordProxy helper class for VectorStore compatibility |
| 2025-12-09 | Updated __init__.py exports for new operators |
| 2025-12-09 | Created comprehensive unit tests (32 tests, all passing) |
| 2025-12-09 | Story marked done - all acceptance criteria met |
| 2025-12-09 | Code review: Added public `index_with_vectors()` to VectorStore |
| 2025-12-09 | Code review: Fixed Milvus connection leak (added disconnect in finally) |
| 2025-12-09 | Code review: Removed double-commit in StoreChunksOperator |
| 2025-12-09 | Code review: Removed unused parameters and variables |
| 2025-12-09 | Code review: Added error handling and disconnect tests (34 tests total) |