POTE/scripts/add_custom_trades.py

#!/usr/bin/env python3
"""
Manually add trades for specific representatives.
Useful when you want to track specific officials or add data from other sources.
"""

import logging
from datetime import datetime, timezone
from decimal import Decimal

from pote.db import get_session
from pote.db.models import Official, Security, Trade

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def add_trade(
    session,
    official_name: str,
    party: str,
    chamber: str,
    state: str,
    ticker: str,
    company_name: str,
    side: str,
    value_min: float,
    value_max: float,
    transaction_date: str,  # YYYY-MM-DD
    disclosure_date: str | None = None,
):
    """Add a single trade to the database."""
    
    # Get or create official
    official = session.query(Official).filter_by(name=official_name).first()
    if not official:
        official = Official(
            name=official_name,
            party=party,
            chamber=chamber,
            state=state,
        )
        session.add(official)
        session.flush()
        logger.info(f"Created official: {official_name}")
    
    # Get or create security
    security = session.query(Security).filter_by(ticker=ticker).first()
    if not security:
        security = Security(ticker=ticker, name=company_name)
        session.add(security)
        session.flush()
        logger.info(f"Created security: {ticker}")
    
    # Create trade
    trade = Trade(
        official_id=official.id,
        security_id=security.id,
        source="manual",
        transaction_date=datetime.strptime(transaction_date, "%Y-%m-%d").date(),
        filing_date=datetime.strptime(disclosure_date, "%Y-%m-%d").date() if disclosure_date else None,
        side=side,
        value_min=Decimal(str(value_min)),
        value_max=Decimal(str(value_max)),
    )
    session.add(trade)
    logger.info(f"Added trade: {official_name} {side} {ticker}")
    
    return trade


def main():
    """Example: Add some trades manually."""
    
    with next(get_session()) as session:
        # Example: Add trades for Elizabeth Warren
        logger.info("Adding trades for Elizabeth Warren...")
        
        add_trade(
            session,
            official_name="Elizabeth Warren",
            party="Democrat",
            chamber="Senate",
            state="MA",
            ticker="AMZN",
            company_name="Amazon.com Inc.",
            side="sell",
            value_min=15001,
            value_max=50000,
            transaction_date="2024-11-15",
            disclosure_date="2024-12-01",
        )
        
        add_trade(
            session,
            official_name="Elizabeth Warren",
            party="Democrat",
            chamber="Senate",
            state="MA",
            ticker="META",
            company_name="Meta Platforms Inc.",
            side="sell",
            value_min=50001,
            value_max=100000,
            transaction_date="2024-11-20",
            disclosure_date="2024-12-05",
        )
        
        # Example: Add trades for Mitt Romney
        logger.info("Adding trades for Mitt Romney...")
        
        add_trade(
            session,
            official_name="Mitt Romney",
            party="Republican",
            chamber="Senate",
            state="UT",
            ticker="BRK.B",
            company_name="Berkshire Hathaway Inc.",
            side="buy",
            value_min=100001,
            value_max=250000,
            transaction_date="2024-10-01",
            disclosure_date="2024-10-15",
        )
        
        session.commit()
        logger.info("✅ All trades added successfully!")
        
        # Show summary
        from sqlalchemy import text
        result = session.execute(text("""
            SELECT o.name, COUNT(t.id) as trade_count
            FROM officials o
            LEFT JOIN trades t ON o.id = t.official_id
            GROUP BY o.name
            ORDER BY trade_count DESC
        """))
        
        print("\n=== Officials Summary ===")
        for row in result:
            print(f"  {row[0]:25s} - {row[1]} trades")


if __name__ == "__main__":
    main()