#!/usr/bin/env python3 """ Migration script to assign existing vehicles to their organization's default garage (branch). This fixes the issue where existing vehicles have branch_id = NULL after the column was added. Logic: 1. For each Asset with owner_org_id or operator_org_id 2. Find the default Branch (Garage) for that Organization (is_main = True) 3. Update Asset.branch_id to that Branch's UUID 4. If no default branch exists, create one """ import asyncio import logging import sys from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession from app.database import AsyncSessionLocal from app.models.vehicle.asset import Asset from app.models.marketplace.organization import Branch, Organization from app.models.identity import User logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) async def get_or_create_default_branch(session: AsyncSession, organization_id: int) -> Branch: """Get the default branch (is_main = True) for an organization, create if doesn't exist.""" # Try to find existing default branch stmt = select(Branch).where( Branch.organization_id == organization_id, Branch.is_main == True, Branch.is_deleted == False ) result = await session.execute(stmt) branch = result.scalar_one_or_none() if branch: logger.info(f"Found default branch {branch.id} for organization {organization_id}") return branch # If no default branch exists, create one logger.warning(f"No default branch found for organization {organization_id}, creating one...") # Get organization name for branch naming org_stmt = select(Organization).where(Organization.id == organization_id) org_result = await session.execute(org_stmt) organization = org_result.scalar_one_or_none() org_name = organization.name if organization else f"Organization {organization_id}" # Create default branch import uuid from datetime import datetime from sqlalchemy.sql import func new_branch = Branch( id=uuid.uuid4(), organization_id=organization_id, name=f"{org_name} - Main Garage", is_main=True, status="active", is_deleted=False, created_at=datetime.utcnow() ) session.add(new_branch) await session.flush() logger.info(f"Created default branch {new_branch.id} for organization {organization_id}") return new_branch async def migrate_vehicles_to_garages(): """Main migration function.""" async with AsyncSessionLocal() as session: try: # Get all assets that have organization ownership but no branch_id stmt = select(Asset).where( (Asset.owner_org_id.is_not(None) | Asset.operator_org_id.is_not(None)), Asset.branch_id.is_(None) ) result = await session.execute(stmt) assets = result.scalars().all() logger.info(f"Found {len(assets)} assets without branch assignment") updated_count = 0 skipped_count = 0 for asset in assets: # Determine which organization to use (prefer owner, fallback to operator) org_id = asset.owner_org_id or asset.operator_org_id if not org_id: logger.warning(f"Asset {asset.id} has no organization reference, skipping") skipped_count += 1 continue # Get or create default branch for the organization branch = await get_or_create_default_branch(session, org_id) # Update the asset update_stmt = ( update(Asset) .where(Asset.id == asset.id) .values(branch_id=branch.id, relocation_performed=True) ) await session.execute(update_stmt) logger.info(f"Updated asset {asset.id} with branch {branch.id} (org {org_id})") updated_count += 1 # Commit all changes await session.commit() logger.info(f"Migration completed: {updated_count} assets updated, {skipped_count} skipped") # Also update assets that already have branch_id but need relocation_performed flag if updated_count > 0: stmt = select(Asset).where( Asset.branch_id.is_not(None), Asset.relocation_performed == False ) result = await session.execute(stmt) assets_without_flag = result.scalars().all() for asset in assets_without_flag: update_stmt = ( update(Asset) .where(Asset.id == asset.id) .values(relocation_performed=True) ) await session.execute(update_stmt) await session.commit() logger.info(f"Updated relocation_performed flag for {len(assets_without_flag)} assets") return updated_count except Exception as e: await session.rollback() logger.error(f"Migration failed: {e}") raise async def verify_migration(): """Verify the migration results.""" async with AsyncSessionLocal() as session: # Count assets with branch_id stmt = select(Asset).where(Asset.branch_id.is_not(None)) result = await session.execute(stmt) assets_with_branch = result.scalars().all() # Count assets without branch_id but with organizations stmt = select(Asset).where( (Asset.owner_org_id.is_not(None) | Asset.operator_org_id.is_not(None)), Asset.branch_id.is_(None) ) result = await session.execute(stmt) assets_still_missing = result.scalars().all() logger.info(f"Verification:") logger.info(f" - Assets with branch_id: {len(assets_with_branch)}") logger.info(f" - Assets still missing branch_id: {len(assets_still_missing)}") if assets_still_missing: logger.warning("Some assets still missing branch_id:") for asset in assets_still_missing[:5]: # Show first 5 logger.warning(f" Asset {asset.id}: owner_org={asset.owner_org_id}, operator_org={asset.operator_org_id}") return len(assets_with_branch), len(assets_still_missing) if __name__ == "__main__": logger.info("Starting vehicle-to-garage migration...") try: # Run migration updated = asyncio.run(migrate_vehicles_to_garages()) # Verify with_branch, missing = asyncio.run(verify_migration()) if missing == 0: logger.info("✅ Migration successful! All organizational vehicles now have branch assignments.") else: logger.warning(f"⚠️ Migration incomplete: {missing} assets still lack branch_id") sys.exit(1) except Exception as e: logger.error(f"❌ Migration failed: {e}") sys.exit(1)