# StructuredStore Integration Guide

This guide shows how to integrate the `StructuredStore` class into the DSOL application for PostgreSQL storage of extracted records.

## Overview

The `StructuredStore` class provides optimized batch storage for extracted records with:
- ✅ Batch optimization (500 records per transaction)
- ✅ JSONB storage for fast JSON queries
- ✅ ON CONFLICT handling for idempotency
- ✅ Retry logic with exponential backoff
- ✅ Performance logging

## Architecture Note

**Current State:** The extraction pipeline (`src/extraction/pipeline.py`) uses **synchronous** SQLAlchemy sessions.

**Future State:** The StructuredStore uses **async** sessions for better performance and scalability.

## Integration Options

### Option 1: Use in Async Contexts (Recommended for New Code)

Use `StructuredStore` in async API endpoints and async Celery tasks.

```python
# src/api/routes/documents.py
from sqlalchemy.ext.asyncio import AsyncSession
from src.knowledge.structured_store import StructuredStore

@router.post("/documents/{document_id}/process-async")
async def process_document_async(
    document_id: UUID,
    session: AsyncSession = Depends(get_async_session)
):
    """Process document using async extraction and storage."""
    # ... async extraction logic ...

    # Store records using StructuredStore
    store = StructuredStore(session)
    count = await store.store_records(records, document_id)

    return {"records_stored": count}
```

### Option 2: Keep Existing Sync Pipeline (Current Implementation)

The current pipeline already has a `_store_records` method (line 709 in `pipeline.py`) that uses synchronous bulk insert.

**Current implementation:**
```python
# src/extraction/pipeline.py (line 709)
def _store_records(self, records: list[ExtractionRecord], document_id: UUID):
    """Store extracted records using sync bulk insert."""
    batch_size = 500
    for i in range(0, len(records), batch_size):
        batch = records[i : i + batch_size]
        batch_data = [...]  # Convert to dicts
        self.db_session.bulk_insert_mappings(db_models.ExtractedRecord, batch_data)
    self.db_session.commit()
```

**Status:** ✅ This works and doesn't need to change yet.

### Option 3: Hybrid Approach (Migration Path)

Create an async version of the pipeline for gradual migration:

```python
# src/extraction/async_pipeline.py (new file)
from sqlalchemy.ext.asyncio import AsyncSession
from src.knowledge.structured_store import StructuredStore

class AsyncExtractionPipeline:
    """Async version of extraction pipeline using StructuredStore."""

    def __init__(self, db_session: AsyncSession):
        self.db_session = db_session
        self.store = StructuredStore(db_session)
        # ... other components ...

    async def process_document(self, document_id: UUID) -> ExtractionSummary:
        """Process document using async pipeline."""
        # PASS 1: Symbol Detection
        # PASS 2: Table Extraction
        records = await self._run_extraction(workbook)

        # Store using StructuredStore
        count = await self.store.store_records(records, document_id)

        # PASS 3: Symbol Resolution
        # PASS 4: Cross-Reference Detection

        return summary
```

## Recommended Integration Steps

### Step 1: Use StructuredStore in New Async Endpoints

```python
# src/api/routes/knowledge.py (new file)
from fastapi import APIRouter, Depends
from sqlalchemy.ext.asyncio import AsyncSession
from src.db.session import get_async_session
from src.knowledge.structured_store import StructuredStore
from src.knowledge.query_helper import StructuredQueryHelper

router = APIRouter(prefix="/knowledge", tags=["knowledge"])

@router.get("/documents/{document_id}/records")
async def get_document_records(
    document_id: UUID,
    item_code: str | None = None,
    session: AsyncSession = Depends(get_async_session)
):
    """Query stored records with optional filtering."""
    query_helper = StructuredQueryHelper(session)

    if item_code:
        records = await query_helper.find_by_exact_field(
            document_id, "Item Code", item_code
        )
    else:
        records = await query_helper.find_by_sheet(document_id, "Sheet1")

    return {
        "document_id": str(document_id),
        "records": [
            {
                "content": r.content,
                "headers": r.headers,
                "sheet": r.sheet_name,
                "row": r.row_number
            }
            for r in records
        ]
    }

@router.post("/documents/{document_id}/store")
async def store_records_endpoint(
    document_id: UUID,
    records: list[dict],  # Your Pydantic model
    session: AsyncSession = Depends(get_async_session)
):
    """Store extracted records."""
    # Convert to ExtractedRecord objects
    extracted_records = [...]  # Your conversion logic

    store = StructuredStore(session)
    count = await store.store_records(extracted_records, document_id)

    return {"success": True, "records_stored": count}
```

### Step 2: Add Query Endpoints Using QueryHelper

```python
@router.get("/documents/{document_id}/search")
async def search_records(
    document_id: UUID,
    field: str,
    value: str,
    session: AsyncSession = Depends(get_async_session)
):
    """Search records by field value."""
    query_helper = StructuredQueryHelper(session)
    records = await query_helper.find_by_exact_field(document_id, field, value)

    return {
        "matches": len(records),
        "records": [r.content for r in records]
    }

@router.get("/documents/{document_id}/symbols")
async def get_resolved_records(
    document_id: UUID,
    session: AsyncSession = Depends(get_async_session)
):
    """Get records with resolved symbol content."""
    query_helper = StructuredQueryHelper(session)
    records = await query_helper.find_with_resolved_content(document_id)

    return {
        "count": len(records),
        "records": [
            {
                "content": r.content,
                "resolved": r.resolved_content
            }
            for r in records
        ]
    }
```

### Step 3: Keep Existing Sync Pipeline Unchanged

**No changes needed** to `src/extraction/pipeline.py` - it continues to work with synchronous sessions.

The existing `_store_records` method (line 709) handles storage for Celery tasks and synchronous contexts.

## Testing Integration

### Unit Tests with Mock Session

```python
# tests/api/test_knowledge_routes.py
import pytest
from unittest.mock import AsyncMock

@pytest.mark.asyncio
async def test_store_records_endpoint(client, mock_db_session):
    """Test async storage endpoint."""
    response = await client.post(
        f"/knowledge/documents/{document_id}/store",
        json={"records": [...]},
    )

    assert response.status_code == 200
    assert response.json()["records_stored"] == 10
```

### Integration Tests with Real Database

```python
# tests/integration/test_knowledge_integration.py
import pytest
import os

pytestmark = pytest.mark.integration

@pytest.mark.asyncio
async def test_full_storage_query_cycle(integration_db_session, test_document):
    """Test complete storage and query cycle."""
    # Store records
    store = StructuredStore(integration_db_session)
    records = [...]  # Sample records
    count = await store.store_records(records, test_document.id)

    assert count == len(records)

    # Query records
    query_helper = StructuredQueryHelper(integration_db_session)
    found = await query_helper.find_by_exact_field(
        test_document.id,
        "Item Code",
        "2024-4_0001"
    )

    assert len(found) == 1
    assert found[0].content["Item Code"] == "2024-4_0001"
```

## Environment Setup

### Development Environment

```bash
# .env
DATABASE_URL=postgresql+asyncpg://user:pass@localhost:5432/dsol
REDIS_URL=redis://localhost:6379/0
LOG_LEVEL=DEBUG
```

### Test Environment

```bash
# .env.test
TEST_DATABASE_URL=postgresql+asyncpg://user:pass@localhost:5432/dsol_test

# Run tests
export TEST_DATABASE_URL="postgresql+asyncpg://user:pass@localhost:5432/dsol_test"
python -m pytest tests/knowledge/ -v

# Run integration tests specifically
python -m pytest tests/knowledge/test_structured_store_integration.py -v
```

### Docker Setup

```yaml
# docker-compose.yml (already exists)
services:
  postgres:
    image: postgres:18
    environment:
      POSTGRES_DB: dsol
      POSTGRES_USER: dsol_user
      POSTGRES_PASSWORD: dsol_password
    ports:
      - "5432:5432"

  postgres_test:
    image: postgres:18
    environment:
      POSTGRES_DB: dsol_test
      POSTGRES_USER: dsol_user
      POSTGRES_PASSWORD: dsol_password
    ports:
      - "5433:5432"
```

## Performance Monitoring

### Structured Logging

All operations log performance metrics:

```python
# Example log output
{
    "event": "store_records_completed",
    "component": "structured_store",
    "document_id": "uuid",
    "records_stored": 1000,
    "duration_ms": 2450,
    "records_per_second": 408,
    "timestamp": "2025-11-28T10:30:00Z"
}
```

### Query Performance Targets

- **Exact field lookup:** < 50ms (uses GIN index)
- **JSONB containment:** < 100ms (uses GIN index)
- **Sheet filter:** < 20ms (uses B-tree index)
- **Batch insert (1000 records):** < 3s

## Migration Path (Future)

When ready to migrate the pipeline to async:

1. **Create async_pipeline.py** - Async version of ExtractionPipeline
2. **Update Celery tasks** - Use async tasks (Celery 5.3+ supports async)
3. **Replace _store_records** - Use StructuredStore instead of bulk_insert_mappings
4. **Add async endpoints** - Gradually migrate API endpoints to async
5. **Remove sync pipeline** - After full migration and testing

## Summary

**Current State (Story 4.1):**
- ✅ StructuredStore implemented and tested
- ✅ QueryHelper for JSONB queries
- ✅ Integration tests ready
- ✅ Existing sync pipeline unchanged

**Next Steps:**
- Use StructuredStore in new async endpoints
- Add query endpoints for record retrieval
- Consider async pipeline migration (Story 4.x or later)

**Files Created:**
- `src/knowledge/structured_store.py` - Async storage class
- `src/knowledge/query_helper.py` - Query helper class
- `tests/knowledge/*` - Comprehensive test suite
- `docs/integration-guide-structured-store.md` - This guide