stock-oracle/app/api/v1/endpoints/migration.py

"""
Data migration endpoints
"""

from datetime import datetime
import time
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Header
from sqlalchemy.ext.asyncio import AsyncSession
import httpx

from app.core.database import get_db
from app.core.config import settings
from app.schemas.financial import MigrationRequest, MigrationResponse, ErrorType
from app.models.financial import FinancialData, CalculatedMetrics, Company

router = APIRouter()

async def verify_migration_key(x_api_key: Optional[str] = Header(None)):
    """Verify migration API key"""
    if not settings.ALLOW_MIGRATION:
        raise HTTPException(
            status_code=403,
            detail={
                "error_type": ErrorType.AUTHENTICATION_ERROR,
                "message": "Migration endpoint is disabled"
            }
        )
    
    if x_api_key != settings.MIGRATION_API_KEY:
        raise HTTPException(
            status_code=401,
            detail={
                "error_type": ErrorType.AUTHENTICATION_ERROR,
                "message": "Invalid migration API key"
            }
        )

@router.post(
    "/migrate",
    response_model=MigrationResponse,
    summary="Migrate data from another instance",
    description="""
    Migrate financial data from another SEC Investment API instance.
    
    This endpoint allows you to:
    - Transfer all data from one instance to another
    - Migrate specific tickers only
    - Migrate data within specific date ranges
    
    Requires valid migration API key in X-API-Key header.
    """
)
async def migrate_data(
    request: MigrationRequest,
    db: AsyncSession = Depends(get_db),
    _: None = Depends(verify_migration_key)
):
    """Migrate data from another instance"""
    
    start_time = time.time()
    total_records = 0
    migrated_records = 0
    failed_records = 0
    errors = []
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            # Set up headers for source API
            headers = {"X-API-Key": request.api_key}
            
            # Get list of tickers to migrate
            if request.tickers:
                tickers = request.tickers
            else:
                # Get all tickers from source
                response = await client.get(
                    f"{request.source_url}/api/v1/companies",
                    headers=headers
                )
                if response.status_code == 200:
                    companies = response.json()
                    tickers = [c["ticker"] for c in companies]
                else:
                    raise ValueError("Failed to fetch company list from source")
            
            # Migrate each ticker
            for ticker in tickers:
                try:
                    # Build query parameters
                    params = {"ticker": ticker}
                    if request.start_date:
                        params["start_date"] = request.start_date.isoformat()
                    if request.end_date:
                        params["end_date"] = request.end_date.isoformat()
                    
                    # Fetch financial data
                    response = await client.get(
                        f"{request.source_url}/api/v1/financial/data/{ticker}",
                        headers=headers,
                        params=params
                    )
                    
                    if response.status_code == 200:
                        data = response.json()
                        
                        # Store company info
                        company_info = data.get("company", {})
                        if company_info:
                            company = await db.get(Company, ticker)
                            if not company:
                                company = Company(
                                    ticker=ticker,
                                    name=company_info.get("name"),
                                    cik=company_info.get("cik"),
                                    sector=company_info.get("sector"),
                                    industry=company_info.get("industry"),
                                    business_description=company_info.get("business_description")
                                )
                                db.add(company)
                        
                        # Store financial data
                        for fd in data.get("financial_data", []):
                            total_records += 1
                            # Check if exists
                            existing = await db.get(
                                FinancialData,
                                (ticker, fd["period_date"], fd["period_type"])
                            )
                            if not existing:
                                fin_data = FinancialData(**fd, ticker=ticker)
                                db.add(fin_data)
                                migrated_records += 1
                        
                        # Store calculated metrics
                        for cm in data.get("calculated_metrics", []):
                            total_records += 1
                            # Check if exists
                            existing = await db.get(
                                CalculatedMetrics,
                                (ticker, cm["calculation_date"], cm["period_date"])
                            )
                            if not existing:
                                metrics = CalculatedMetrics(**cm, ticker=ticker)
                                db.add(metrics)
                                migrated_records += 1
                        
                        await db.commit()
                        
                    else:
                        failed_records += 1
                        errors.append({
                            "ticker": ticker,
                            "error": f"HTTP {response.status_code}: {response.text}"
                        })
                        
                except Exception as e:
                    failed_records += 1
                    errors.append({
                        "ticker": ticker,
                        "error": str(e)
                    })
                    await db.rollback()
            
            duration = time.time() - start_time
            
            return MigrationResponse(
                status="completed" if failed_records == 0 else "completed_with_errors",
                total_records=total_records,
                migrated_records=migrated_records,
                failed_records=failed_records,
                errors=errors[:10],  # Limit errors to first 10
                duration_seconds=round(duration, 2)
            )
            
        except Exception as e:
            return MigrationResponse(
                status="failed",
                total_records=total_records,
                migrated_records=migrated_records,
                failed_records=failed_records,
                errors=[{"error": str(e)}],
                duration_seconds=round(time.time() - start_time, 2)
            )

@router.get(
    "/migration/export/{ticker}",
    summary="Export data for migration",
    description="Export financial data for a specific ticker (used by migration process)"
)
async def export_data(
    ticker: str,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
    db: AsyncSession = Depends(get_db),
    _: None = Depends(verify_migration_key)
):
    """Export data for migration"""
    
    # This endpoint would be used by the migration process
    # Implementation depends on specific needs
    
    return {
        "ticker": ticker,
        "message": "Export endpoint for migration",
        "note": "This would return raw data for migration purposes"
    }