""" Populate SwiftOps codes for existing organizations Run this after applying the migration: 20250128000000_add_swiftops_code.sql Usage: python scripts/populate_org_codes.py """ import sys import os from pathlib import Path # Add src to path sys.path.insert(0, str(Path(__file__).parent.parent / "src")) from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from app.models.client import Client from app.models.contractor import Contractor from app.utils.org_code_generator import generate_org_code def get_database_url(): """Get database URL from environment""" database_url = os.getenv("DATABASE_URL") if not database_url: raise ValueError("DATABASE_URL environment variable not set") return database_url def populate_codes(): """Populate swiftops_code for all existing organizations""" # Create database connection database_url = get_database_url() engine = create_engine(database_url) SessionLocal = sessionmaker(bind=engine) db = SessionLocal() try: print("Starting organization code population...") print("-" * 60) # Process clients clients = db.query(Client).filter(Client.swiftops_code.is_(None)).all() print(f"\nFound {len(clients)} clients without codes") for client in clients: try: code = generate_org_code(client.name, db) client.swiftops_code = code db.flush() # Flush to make code available for uniqueness checks print(f"✓ Client: {client.name[:40]:<40} -> {code}") except Exception as e: print(f"✗ Failed for client '{client.name}': {e}") db.rollback() continue # Commit all client codes db.commit() print(f"\n✓ Successfully generated codes for {len(clients)} clients") # Process contractors contractors = db.query(Contractor).filter(Contractor.swiftops_code.is_(None)).all() print(f"\nFound {len(contractors)} contractors without codes") for contractor in contractors: try: code = generate_org_code(contractor.name, db) contractor.swiftops_code = code db.flush() # Flush to make code available for uniqueness checks print(f"✓ Contractor: {contractor.name[:40]:<40} -> {code}") except Exception as e: print(f"✗ Failed for contractor '{contractor.name}': {e}") db.rollback() continue # Commit all contractor codes db.commit() print(f"\n✓ Successfully generated codes for {len(contractors)} contractors") print("-" * 60) print("\n✅ Code population completed successfully!") print("\nNext steps:") print("1. Verify all organizations have codes:") print(" SELECT name, swiftops_code FROM clients WHERE swiftops_code IS NULL;") print(" SELECT name, swiftops_code FROM contractors WHERE swiftops_code IS NULL;") print("\n2. If all look good, make the field NOT NULL:") print(" ALTER TABLE clients ALTER COLUMN swiftops_code SET NOT NULL;") print(" ALTER TABLE contractors ALTER COLUMN swiftops_code SET NOT NULL;") except Exception as e: print(f"\n❌ Error during population: {e}") db.rollback() raise finally: db.close() if __name__ == "__main__": populate_codes()