190 lines
7.2 KiB
Python
190 lines
7.2 KiB
Python
#!/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) |