#!/usr/bin/env python3 # -*- coding: utf-8 -*- import os import logging import sys from typing import Dict, Tuple, Any from db.functions import settype, dbmongo from db.schema_generator import MongoToPostgresSchemaGenerator, MongoDocumentInserter import settings.config # Configuration constants SAMPLE_SIZE = 10000 BATCH_SIZE = 10000 SCHEMA_NAME = 'updates' PK_FIELD = 'id' # Table-to-collection mapping TABLE_CONFIG = { 'seriesdata': 'mgseries', 'episodesdata': 'mgepisodes', 'actorsdata': 'mgactors', 'charactersdata': 'mgcharacters', 'crewdata': 'mgcrew', } def setup_logger() -> logging.Logger: """Initialize and return configured logger.""" logger = logging.getLogger(__name__) if not logger.handlers: handler = logging.StreamHandler() formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') handler.setFormatter(formatter) logger.addHandler(handler) logger.setLevel(logging.INFO) return logger def initialize_environment(root_dir: str) -> Tuple[Dict[str, Any], Any, Any]: """Initialize config, database engine, and MongoDB updater. Args: root_dir: Root directory path Returns: Tuple of (options dict, db engine, mongo updater) """ logger = logging.getLogger(__name__) try: config_options = settings.config.Config(root_dir) options = config_options.config_options if not options: raise ValueError("Configuration is empty") dbtype = options.get('dbtype') apitype = options.get('apitype') if not dbtype or not apitype: raise ValueError("Missing required config options: dbtype, apitype") dbs = settype(dbtype, apitype, options) dbengine = dbs.dbengine mongo_updater = dbmongo(options) logger.info("Environment initialized successfully") return options, dbengine, mongo_updater except Exception as e: logger.error(f"Failed to initialize environment: {e}") raise def create_tables( generator: MongoToPostgresSchemaGenerator, mongo_updater: Any, dbengine: Dict[str, Any], table_config: Dict[str, str], logger: logging.Logger ) -> Dict[str, Dict[str, Any]]: """Create PostgreSQL tables and return cached schemas. Args: generator: Schema generator instance mongo_updater: MongoDB updater instance dbengine: Database engine configuration table_config: Mapping of table names to MongoDB collection attributes logger: Logger instance Returns: Dict mapping table names to their schemas """ schemas = {} for table_name, collection_attr in table_config.items(): try: # Get MongoDB collection collection = getattr(mongo_updater, collection_attr, None) if not collection: logger.warning(f"Collection '{collection_attr}' not found on mongo_updater") continue # Analyze schema schema = generator.analyze_collection(mongo_updater, collection) schemas[table_name] = schema # Generate and display SQL sql = generator.generate_create_table_sql( table_name, schema_name=SCHEMA_NAME, pk_field=PK_FIELD, ) logger.info(f"Schema for {table_name}:\n{sql}") # Create table in PostgreSQL generator.create_table_in_postgres( dbengine, table_name, schema_name=SCHEMA_NAME, pk_field=PK_FIELD, drop_existing=True, ) logger.info(f"Table '{table_name}' created/recreated successfully") except Exception as e: logger.error(f"Failed to create table '{table_name}': {e}") continue return schemas def build_upsert_clause(schema: Dict[str, Any], pk_field: str) -> str: """Build SQL ON CONFLICT clause for upsert. Args: schema: Column schema dictionary pk_field: Primary key field name Returns: ON CONFLICT clause string """ # Exclude primary key and MongoDB's _id field from updates update_cols = [ col.lower() for col in schema.keys() if col.lower() != pk_field.lower() and col != '_id' ] if not update_cols: return f"ON CONFLICT ({pk_field}) DO NOTHING" update_clause = ', '.join([f"{col}=EXCLUDED.{col}" for col in update_cols]) return f"ON CONFLICT ({pk_field}) DO UPDATE SET {update_clause}" def should_update_document( doc_id: Any, doc_updated: Any, dbengine: Dict[str, Any], table_name: str, schema_name: str, logger: logging.Logger ) -> bool: """Compare document's updated timestamp to the database row's updated timestamp. Args: doc_id: Document ID from MongoDB doc_updated: Updated timestamp from MongoDB document dbengine: Database engine configuration table_name: Target table name schema_name: Target schema name logger: Logger instance Returns: True if document should be updated (MongoDB is newer), False otherwise """ if doc_updated is None: logger.debug(f"Document {doc_id} has no 'updated' field, will update") return True try: from sqlalchemy import text # Query the database for the existing record's updated timestamp query = text(f""" SELECT updated FROM {schema_name}.{table_name} WHERE id = :id """) with dbengine['engine'].connect() as conn: result = conn.execute(query, {'id': doc_id}).fetchone() # If record doesn't exist, update it if result is None: logger.debug(f"Record {doc_id} not in DB, will insert") return True db_updated = result[0] # Compare timestamps if db_updated is None: logger.debug(f"Record {doc_id} has NULL updated in DB, will update") return True # Only update if MongoDB version is newer should_update = doc_updated > db_updated if not should_update: logger.debug(f"Record {doc_id}: DB updated={db_updated} >= doc updated={doc_updated}, skipping") else: logger.debug(f"Record {doc_id}: doc updated={doc_updated} > DB updated={db_updated}, updating") return should_update except Exception as e: logger.warning(f"Error comparing timestamps for {doc_id}: {e}. Proceeding with update.") return True def insert_documents( inserter: MongoDocumentInserter, mongo_updater: Any, dbengine: Dict[str, Any], table_config: Dict[str, str], schemas: Dict[str, Dict[str, Any]], logger: logging.Logger, compare_timestamps: bool = True ) -> None: """Insert/upsert documents from MongoDB collections to PostgreSQL. Args: inserter: Document inserter instance mongo_updater: MongoDB updater instance dbengine: Database engine configuration table_config: Mapping of table names to MongoDB collection attributes schemas: Cached schemas from table creation logger: Logger instance compare_timestamps: If True, only update if MongoDB doc is newer than DB record """ total_inserted = 0 total_skipped = 0 for table_name, collection_attr in table_config.items(): if table_name not in schemas: logger.warning(f"Skipping '{table_name}' - schema not available") continue try: # Get collection and schema collection = getattr(mongo_updater, collection_attr) schema = schemas[table_name] # Build upsert clause on_conflict = build_upsert_clause(schema, PK_FIELD) # If timestamp comparison is enabled, filter documents before inserting if compare_timestamps: logger.info(f"Filtering {table_name} documents by updated timestamp and batching every {BATCH_SIZE}...") batch = [] batch_num = 0 for doc in collection.find(): doc_id = doc.get(PK_FIELD) doc_updated = doc.get('updated') if should_update_document(doc_id, doc_updated, dbengine, table_name, SCHEMA_NAME, logger): batch.append(doc) else: total_skipped += 1 # When batch reaches size, insert and reset if len(batch) >= BATCH_SIZE: batch_num += 1 try: # Explicit transaction per batch with dbengine['engine'].begin(): count = inserter.insert_documents( engine=dbengine, table_name=table_name, documents=batch, schema_name=SCHEMA_NAME, on_conflict=on_conflict, ) total_inserted += count logger.info(f"Batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)") except Exception as e: logger.error(f"Batch {batch_num} failed for '{table_name}': {e}") finally: batch = [] # Insert any remaining docs in the final batch if batch: batch_num += 1 try: with dbengine['engine'].begin(): count = inserter.insert_documents( engine=dbengine, table_name=table_name, documents=batch, schema_name=SCHEMA_NAME, on_conflict=on_conflict, ) total_inserted += count logger.info(f"Final batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)") except Exception as e: logger.error(f"Final batch {batch_num} failed for '{table_name}': {e}") finally: batch = [] else: # Insert all documents without timestamp comparison # Stream and batch inserts to avoid memory spikes batch = [] batch_num = 0 for doc in collection.find(): batch.append(doc) if len(batch) >= BATCH_SIZE: batch_num += 1 try: with dbengine['engine'].begin(): count = inserter.insert_documents( engine=dbengine, table_name=table_name, documents=batch, schema_name=SCHEMA_NAME, on_conflict=on_conflict, ) total_inserted += count logger.info(f"Batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)") except Exception as e: logger.error(f"Batch {batch_num} failed for '{table_name}': {e}") finally: batch = [] if batch: batch_num += 1 try: with dbengine['engine'].begin(): count = inserter.insert_documents( engine=dbengine, table_name=table_name, documents=batch, schema_name=SCHEMA_NAME, on_conflict=on_conflict, ) total_inserted += count logger.info(f"Final batch {batch_num}: Inserted/Updated {count} documents into '{table_name}' (committed)") except Exception as e: logger.error(f"Final batch {batch_num} failed for '{table_name}': {e}") finally: batch = [] except Exception as e: logger.error(f"Failed to insert documents into '{table_name}': {e}") continue logger.info(f"Total documents inserted/updated: {total_inserted}, skipped: {total_skipped}") def main() -> int: """Main entry point for MongoDB to PostgreSQL sync.""" logger = setup_logger() try: root_dir = os.getcwd() # Initialize environment options, dbengine, mongo_updater = initialize_environment(root_dir) # Create schema generator and inserter generator = MongoToPostgresSchemaGenerator(sample_size=SAMPLE_SIZE) inserter = MongoDocumentInserter(batch_size=BATCH_SIZE) # Create tables and get cached schemas logger.info("Creating PostgreSQL tables...") schemas = create_tables(generator, mongo_updater, dbengine, TABLE_CONFIG, logger) if not schemas: logger.error("No schemas were successfully created. Aborting insert.") return 1 # Insert documents from MongoDB logger.info("Inserting documents from MongoDB...") insert_documents(inserter, mongo_updater, dbengine, TABLE_CONFIG, schemas, logger) logger.info("MongoDB to PostgreSQL sync completed successfully") return 0 except Exception as e: logger.error(f"Sync failed: {e}", exc_info=True) return 1 if __name__ == "__main__": sys.exit(main())