wjones
afe2dfa1c8
This commit introduces a new function, `insert_mongo_documents_to_postgres`, which allows for the bulk updating of documents in a PostgreSQL table. It takes a list of document IDs and a dictionary containing the updated data as input. This method simplifies the process by reducing the number of individual insert statements required. The changes include: 1. Adding a new function `insert_mongo_documents_to_postgres` to handle the bulk update. 2. Updating the existing code to call this new function when required. 3. Refactoring the `MongoDocumentInserter` class to encapsulate common functionality for document insertion and updating. This refactoring enhances maintainability and efficiency by reducing redundancy in the codebase.
395 lines
14 KiB
Python
395 lines
14 KiB
Python
#!/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())
|