from fastapi import Request, Query, APIRouter from fastapi.responses import JSONResponse from pydantic import BaseModel from typing import Optional, List, Dict, Any import os import sys import logging from datetime import datetime import pytz from connectors.bigquery.bq import BigQuery from connectors.firebase.firebase import Firebase # Initialize logger first logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Add migration scripts to path (from system-api root) # system-api/api/migration/audience.py -> ../../migration/ # Migration scripts are now inside system-api/migration/ system_api_root = os.path.join(os.path.dirname(__file__), '../..') system_api_root = os.path.abspath(system_api_root) migration_path = os.path.join(system_api_root, 'migration') scripts_path = os.path.join(system_api_root, 'migration', 'scripts') # Add to path if not already there if migration_path not in sys.path: sys.path.insert(0, migration_path) if scripts_path not in sys.path: sys.path.insert(0, scripts_path) # Import functions from migration scripts # Note: These scripts add system-api to path, so we need to import after path setup try: import importlib.util # Load audience_migration module audience_migration_path = os.path.join(migration_path, 'audience_migration.py') spec = importlib.util.spec_from_file_location("audience_migration", audience_migration_path) audience_migration = importlib.util.module_from_spec(spec) sys.modules["audience_migration"] = audience_migration spec.loader.exec_module(audience_migration) # Load create_audience_members_table module create_table_path = os.path.join(scripts_path, 'create_audience_members_table.py') spec2 = importlib.util.spec_from_file_location("create_audience_members_table", create_table_path) create_table_module = importlib.util.module_from_spec(spec2) sys.modules["create_audience_members_table"] = create_table_module spec2.loader.exec_module(create_table_module) # Import functions migrate_audience_members = audience_migration.migrate_audience_members is_preset_audience = audience_migration.is_preset_audience create_table_for_property = create_table_module.create_table_for_property AUDIENCE_MEMBERS_SCHEMA = create_table_module.AUDIENCE_MEMBERS_SCHEMA PARTITION_FIELD = create_table_module.PARTITION_FIELD CLUSTER_FIELDS = create_table_module.CLUSTER_FIELDS except Exception as e: logger.error(f"Failed to import migration functions: {e}") raise ImportError(f"Cannot import migration functions: {e}") router = APIRouter() timezone_utc = pytz.utc PROJECT_ID = os.environ.get('GCP_PROJECT', 'customer-360-profile') fb = Firebase(host=os.environ.get("FIREBASE_HOST")) bq = BigQuery() # Request Models class CreateTableRequest(BaseModel): property_id: Optional[str] = None class MigrateRequest(BaseModel): property_id: Optional[str] = None dry_run: bool = False def get_all_properties() -> List[str]: """Get list of all property IDs from Firebase.""" try: accounts = fb.db.reference().child("account").get(shallow=True) if accounts: return list(accounts.keys()) return [] except Exception as e: logger.error(f"Failed to get properties from Firebase: {e}") return [] def check_table_creation_status(property_id: str) -> Dict[str, Any]: """Check if table was created for this property.""" try: table_creation_ref = fb.db.reference().child(f'account/{property_id}/migration/audience/table_creation') table_creation_logs = table_creation_ref.get() if not table_creation_logs: return {"table_created": False, "created_at": None} # Get the latest creation log latest_log = None latest_timestamp = None for timestamp, log_data in table_creation_logs.items(): if isinstance(log_data, dict) and log_data.get('status') == 'completed': if latest_timestamp is None or timestamp > latest_timestamp: latest_timestamp = timestamp latest_log = log_data if latest_log: return { "table_created": True, "created_at": latest_log.get('created_at'), "table_ref": latest_log.get('table_ref') } return {"table_created": False, "created_at": None} except Exception as e: logger.error(f"Error checking table creation status: {e}") return {"table_created": False, "created_at": None} def check_migration_status(property_id: str) -> Dict[str, Any]: """Check migration status for a property.""" try: migration_ref = fb.db.reference().child(f'account/{property_id}/migration/audience') migrations = migration_ref.get() if not migrations: return {"migrated": False, "last_migration": None} # Find completed migrations (exclude table_creation) completed_migrations = [] in_progress_migrations = [] for migration_id, migration_data in migrations.items(): if migration_id == 'table_creation': continue if isinstance(migration_data, dict): status = migration_data.get('status') if status == 'completed': completed_migrations.append({ 'migration_id': migration_id, **migration_data }) elif status == 'in_progress': in_progress_migrations.append({ 'migration_id': migration_id, **migration_data }) # Sort by completed_at completed_migrations.sort(key=lambda x: x.get('completed_at', ''), reverse=True) last_migration = completed_migrations[0] if completed_migrations else None return { "migrated": len(completed_migrations) > 0, "last_migration": last_migration, "in_progress": len(in_progress_migrations) > 0, "in_progress_migrations": in_progress_migrations, "migration_history": completed_migrations[:10] # Last 10 migrations } except Exception as e: logger.error(f"Error checking migration status: {e}") return {"migrated": False, "last_migration": None} @router.post("/table/create", description="Create audience_members table in BigQuery") async def create_table(request: Request, body: CreateTableRequest): """Create audience_members table for one or more properties.""" try: # Get property IDs if body.property_id: property_ids = [body.property_id] else: property_ids = get_all_properties() if not property_ids: return JSONResponse( status_code=400, content={"status": "error", "message": "No properties found"} ) results = [] for property_id in property_ids: try: # Check if table already created table_status = check_table_creation_status(property_id) if table_status.get('table_created'): logger.info(f"Table already exists for property {property_id}") results.append({ "property_id": property_id, "table_created": True, "already_exists": True, "table_ref": table_status.get('table_ref'), "created_at": table_status.get('created_at') }) continue # Create table success = create_table_for_property(property_id, bq) if success: # Log to Firebase now = datetime.now(timezone_utc) timestamp = now.strftime("%Y%m%d_%H%M%S") table_ref = f"{PROJECT_ID}.client_{property_id}.audience_members" log_data = { "status": "completed", "created_at": now.isoformat(), "table_ref": table_ref } fb.db.reference().child( f'account/{property_id}/migration/audience/table_creation/{timestamp}' ).set(log_data) results.append({ "property_id": property_id, "table_created": True, "already_exists": False, "table_ref": table_ref, "created_at": now.isoformat() }) else: results.append({ "property_id": property_id, "table_created": False, "error": "Failed to create table" }) except Exception as e: logger.error(f"Error creating table for property {property_id}: {e}") results.append({ "property_id": property_id, "table_created": False, "error": str(e) }) return JSONResponse( status_code=200, content={ "status": "ok", "message": "Table creation completed", "data": results } ) except Exception as e: logger.error(f"Error in create_table endpoint: {e}") return JSONResponse( status_code=500, content={"status": "error", "message": str(e)} ) @router.post("/migrate", description="Migrate Preset Audience data from Firebase to BigQuery") async def migrate_audience(request: Request, body: MigrateRequest): """Migrate Preset Audience data for one or more properties.""" try: # Get property IDs if body.property_id: property_ids = [body.property_id] else: property_ids = get_all_properties() if not property_ids: return JSONResponse( status_code=400, content={"status": "error", "message": "No properties found"} ) all_results = [] for property_id in property_ids: try: # Check if already migrated migration_status = check_migration_status(property_id) if migration_status.get('migrated') and not body.dry_run: all_results.append({ "property_id": property_id, "error": f"Migration already completed for property {property_id}", "skipped": True, "last_migration": migration_status.get('last_migration') }) logger.warning(f"Migration already completed for property {property_id}, skipping") continue # Check if migration in progress if migration_status.get('in_progress'): all_results.append({ "property_id": property_id, "error": f"Migration already in progress for property {property_id}", "skipped": True, "in_progress_migrations": migration_status.get('in_progress_migrations') }) logger.warning(f"Migration already in progress for property {property_id}, skipping") continue # Create migration_id now = datetime.now(timezone_utc) migration_id = f"migration_{now.strftime('%Y%m%d_%H%M%S')}" # Mark as in_progress if not body.dry_run: fb.db.reference().child( f'account/{property_id}/migration/audience/{migration_id}' ).set({ "status": "in_progress", "started_at": now.isoformat(), "dry_run": body.dry_run }) # Run migration result = migrate_audience_members(property_id, fb, bq, dry_run=body.dry_run) # Update migration log if not body.dry_run: completed_at = datetime.now(timezone_utc) fb.db.reference().child( f'account/{property_id}/migration/audience/{migration_id}' ).update({ "status": "completed", "completed_at": completed_at.isoformat(), "results": { "total_audiences": result.get('total_audiences', 0), "preset_audiences": result.get('preset_audiences', 0), "builder_audiences": result.get('builder_audiences', 0), "migrated": result.get('migrated', 0), "skipped": result.get('skipped', 0), "errors": result.get('errors', []) } }) all_results.append({ "property_id": property_id, "migration_id": migration_id, **result }) except Exception as e: logger.error(f"Error migrating property {property_id}: {e}") all_results.append({ "property_id": property_id, "error": str(e) }) return JSONResponse( status_code=200, content={ "status": "ok", "message": "Migration completed", "data": all_results } ) except Exception as e: logger.error(f"Error in migrate_audience endpoint: {e}") return JSONResponse( status_code=500, content={"status": "error", "message": str(e)} ) @router.get("/status", description="Check migration status for a property") async def get_migration_status( request: Request, property_id: str = Query(..., description="Property ID") ): """Get migration status for a property.""" try: # Check table creation status table_status = check_table_creation_status(property_id) # Check migration status migration_status = check_migration_status(property_id) return JSONResponse( status_code=200, content={ "status": "ok", "data": { "property_id": property_id, "table_created": table_status.get('table_created', False), "table_created_at": table_status.get('created_at'), "table_ref": table_status.get('table_ref'), "migrated": migration_status.get('migrated', False), "last_migration": migration_status.get('last_migration'), "in_progress": migration_status.get('in_progress', False), "migration_history": migration_status.get('migration_history', []) } } ) except Exception as e: logger.error(f"Error in get_migration_status endpoint: {e}") return JSONResponse( status_code=500, content={"status": "error", "message": str(e)} )