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.
This commit is contained in:
+376
-122
@@ -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)
|
||||
|
||||
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())
|
||||
|
||||
Reference in New Issue
Block a user