From afe2dfa1c84fa448d9c4d0f4eb70425472ec5bd3 Mon Sep 17 00:00:00 2001 From: Wendell Jones Date: Tue, 10 Feb 2026 12:01:22 -0500 Subject: [PATCH] Update: Implement bulk update functionality using list of document IDs and updated data 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. --- .gitignore | 2 +- db/schema_generator.py | 107 ++++++++- mongodb2postgres.py | 498 +++++++++++++++++++++++++++++++---------- 3 files changed, 475 insertions(+), 132 deletions(-) diff --git a/.gitignore b/.gitignore index 96b66b3..6936573 100644 --- a/.gitignore +++ b/.gitignore @@ -24,4 +24,4 @@ dbsync.log.* /temp/tvmaze_updates.json /temp.skip_ids.txt /archived -/settings/tvsync_settings.cfg \ No newline at end of file +settings/tvsync_settings.cfg diff --git a/db/schema_generator.py b/db/schema_generator.py index 170d48b..7ccb9ec 100644 --- a/db/schema_generator.py +++ b/db/schema_generator.py @@ -530,17 +530,106 @@ class MongoDocumentInserter: if on_conflict: insert_sql += f" {on_conflict}" - with engine.begin() as conn: - if on_conflict: + # Use different strategies depending on whether ON CONFLICT is needed. + # For ON CONFLICT updates, per-row execution is slow; use psycopg2's + # execute_values to perform fast multi-row INSERT ... VALUES (...) ON CONFLICT ... + if on_conflict: + try: + # Prepare ordered tuples for insertion + values_list = [] for doc in prepared: + values_list.append(tuple(doc.get(col) for col in column_names)) + + # Build base INSERT with placeholder for execute_values + insert_base = f"INSERT INTO {full_table_name} ({cols_str}) VALUES %s {on_conflict}" + + # Use raw DB-API connection for execute_values + try: + import psycopg2 + except Exception: + # Fall back to SQLAlchemy executemany if psycopg2 not available + with engine.begin() as conn: + for doc in prepared: + try: + conn.execute(text(insert_sql), [doc]) + total_inserted += 1 + except Exception as e: + logger.error(f"Error inserting document into {full_table_name}: {e}") + raise + logger.info(f"Inserted {total_inserted} documents into {full_table_name}") + else: + # Use a single staging temporary table per insert_documents call. + # Stream each batch into the temp table via COPY, then run one + # INSERT ... SELECT ... ON CONFLICT ... to upsert all rows. + raw_conn = engine.raw_connection() try: - conn.execute(text(insert_sql), [doc]) - total_inserted += 1 - except Exception as e: - logger.error(f"Error inserting document into {full_table_name}: {e}") - raise - logger.info(f"Inserted {total_inserted} documents into {full_table_name}") - else: + cur = raw_conn.cursor() + import io + import csv + import time + import os + + temp_name = f"tmp_{table_name}_{int(time.time())}_{os.getpid()}" + try: + # Create temporary table once + cur.execute(f"CREATE TEMP TABLE {temp_name} (LIKE {full_table_name} INCLUDING ALL);") + # Stream batches into temp table + for i in range(0, len(prepared), self.batch_size): + batch = prepared[i:i + self.batch_size] + try: + sio = io.StringIO() + writer = csv.writer(sio) + for doc in batch: + row = [] + for col in column_names: + val = doc.get(col) + if val is None: + row.append('\\N') + else: + row.append(val) + writer.writerow(row) + sio.seek(0) + + copy_sql = f"COPY {temp_name} ({cols_str}) FROM STDIN WITH CSV NULL '\\N'" + cur.copy_expert(copy_sql, sio) + raw_conn.commit() + batch_count = len(batch) + total_inserted += batch_count + logger.info(f"Copied {batch_count} rows into staging {temp_name} (total staged: {total_inserted})") + except Exception as e: + raw_conn.rollback() + logger.error(f"Error copying batch into {temp_name}: {e}") + raise + + # Perform a single upsert from the staging table + try: + insert_from_temp = f"INSERT INTO {full_table_name} ({cols_str}) SELECT {cols_str} FROM {temp_name} {on_conflict}" + cur.execute(insert_from_temp) + raw_conn.commit() + logger.info(f"Upserted staged rows from {temp_name} into {full_table_name}") + except Exception as e: + raw_conn.rollback() + logger.error(f"Error upserting from {temp_name} into {full_table_name}: {e}") + raise + finally: + # Temp table will be dropped on commit/connection close, ensure commit + try: + raw_conn.commit() + except Exception: + pass + finally: + try: + cur.close() + except Exception: + pass + try: + raw_conn.close() + except Exception: + pass + except Exception: + raise + else: + with engine.begin() as conn: for i in range(0, len(prepared), self.batch_size): batch = prepared[i:i + self.batch_size] try: diff --git a/mongodb2postgres.py b/mongodb2postgres.py index 3cccc85..8c8bbb5 100644 --- a/mongodb2postgres.py +++ b/mongodb2postgres.py @@ -1,140 +1,394 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- import os -import time +import logging import sys -from tqdm import tqdm -import logs.Logger +from typing import Dict, Tuple, Any from db.functions import settype, dbmongo -from datetime import datetime -import json from db.schema_generator import MongoToPostgresSchemaGenerator, MongoDocumentInserter import settings.config - -def set_apienv(urls, uprocess, dbengine, dbExec, updatesBase, lprint): - """Populate updates table from the API and return available updates. - - Args: - urls: URL configuration with an "updatesurl" key. - uprocess: URL loader instance with `get_data`. - jget: JSON processor with `jconvert`. - dbengine: database engine dict (expects key "engine"). - dbExec: database execution helper with `update_tvupdates` and `rawsql_select`. - updatesBase: list-like containing tables (uses index 8 for updates table). - lprint: logger instance with `logprint`. - - Returns: - List of tuples (seriesid, timestamp) representing available updates. - """ - #updatetable = updatesBase[8] - #inputdata = uprocess.get_data(urls["updatesurl"]) - #lprint.logprint("info", f"Retrieved {len(availableupdates)} rows for processing......") - # - newupdates = dbExec.rawsql_select( - dbengine["engine"], - "select seriesid,timestamp from updates.tvupdates", - lprint - ) - return newupdates -ROOTDIR = os.getcwd() -config_options = settings.config.Config(ROOTDIR) -options = config_options.config_options -dbtype, apitype = options["dbtype"], options["apitype"] -dbs = settype(dbtype, apitype, options) -dbclass, dbengine = dbs.dbclass, dbs.dbengine -mongo_updater = dbmongo(options) -generator = MongoToPostgresSchemaGenerator(sample_size=100) +# Configuration constants +SAMPLE_SIZE = 10000 +BATCH_SIZE = 10000 +SCHEMA_NAME = 'updates' +PK_FIELD = 'id' -tableNames = ['seriesdata', 'episodesdata', 'actorsdata', 'charactersdata', 'crewdata'] - -for tableName in tableNames: - # Generate and print the schema for each collection - schema = generator.analyze_collection(mongo_updater, getattr(mongo_updater, f'mg{tableName[:-4]}')) - sql = generator.generate_create_table_sql( - tableName, - schema_name='updates', - pk_field='id' if tableName == 'seriesdata' else 'id', - ) - print(f"--- SQL Schema for {tableName} ---") - print(sql) - print("\n") - - # Or create it directly - generator.create_table_in_postgres( - dbengine, - tableName, - schema_name='updates', - #pk_field='seriesid' if tableName == 'seriesdata' else 'id', - pk_field='id', - drop_existing=True, - ) - - - -inserter = MongoDocumentInserter(batch_size=10000) - -# Map table names to their MongoDB collections and primary keys -collection_map = { - 'seriesdata': (mongo_updater.mgseries, 'id'), - 'episodesdata': (mongo_updater.mgepisodes, 'id'), - 'actorsdata': (mongo_updater.mgactors, 'id'), - 'charactersdata': (mongo_updater.mgcharacters, 'id'), - 'crewdata': (mongo_updater.mgcrew, 'id'), +# Table-to-collection mapping +TABLE_CONFIG = { + 'seriesdata': 'mgseries', + 'episodesdata': 'mgepisodes', + 'actorsdata': 'mgactors', + 'charactersdata': 'mgcharacters', + 'crewdata': 'mgcrew', } -# Insert documents from collections with UPSERT -for tableName in tableNames: - if tableName in collection_map: - collection, pk_field = collection_map[tableName] - # Build ON CONFLICT ... DO UPDATE clause - # Re-analyze the specific collection so we have the correct fields - schema = generator.analyze_collection(mongo_updater, collection) - # Use lowercased column names to match PostgreSQL unquoted identifiers - update_cols = [col.lower() for col in schema.keys() if col.lower() != pk_field.lower() and col != '_id'] - update_clause = ', '.join([f"{col}=EXCLUDED.{col}" for col in update_cols]) - on_conflict = f"ON CONFLICT ({pk_field}) DO UPDATE SET {update_clause}" + +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 - count = inserter.insert_from_collection( - engine=dbengine, - table_name=tableName, - collection=collection, - schema_name='updates', - on_conflict=on_conflict, - ) - print(f"Inserted/Updated {count} documents in {tableName}") - else: - print(f"Warning: No collection mapping for {tableName}") + 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 -# # Or directly from MongoDB collection -# count = inserter.insert_from_collection( -# engine=dbengine, -# table_name='seriesdata', -# collection=mongo_db['series'], -# schema_name='dbo', -# on_conflict="DO NOTHING", -# ) -# if __name__ == "__main__": -# sys.exit(main()) -# Below are the relevant methods from db/functions.py and db/query_utils.py that are used in the above code. They are included here for completeness and context. -# Needs to be incorporated into logic so that only releveant documents are updated. +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 -# from db.query_utils import MongoQueryUtils -# # With additional filters -# recent = MongoQueryUtils.get_recent_updates( -# collection=mongo_updater.mgseries, -# minutes=60, -# updated_field='updated', -# query_filter={'status': 'Ended'}, # Only ended shows -# ) +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}" -# # Since a specific timestamp -# docs = MongoQueryUtils.get_updates_since( -# collection=mongo_updater.mgseries, -# since_timestamp=None, # Epoch or ISO string -# ) -# for doc in docs: -# print(doc) \ No newline at end of file + +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())