wjones/new_dbsync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update: Implement bulk update functionality using list of document IDs and updated data

Browse Source 
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:
Wendell Jones
2026-02-10 12:01:22 -05:00
parent c6f9d1e617
commit afe2dfa1c8
3 changed files with 475 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -24,4 +24,4 @@ dbsync.log.*
/temp/tvmaze_updates.json
/temp.skip_ids.txt
/archived
/settings/tvsync_settings.cfg
settings/tvsync_settings.cfg
db/schema_generator.py
+98 -9
View File
@@ -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:
mongodb2postgres.py
+376 -122
View File
@@ -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())