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.
744 lines
28 KiB
Python
744 lines
28 KiB
Python
"""Generate PostgreSQL table schemas from MongoDB document structures."""
|
|
|
|
__author__ = 'Wendell Jones'
|
|
|
|
from typing import Any, Dict, List, Optional, Set
|
|
from pymongo import MongoClient
|
|
from sqlalchemy import text, MetaData, insert, Table
|
|
import logging
|
|
import json
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MongoToPostgresSchemaGenerator:
|
|
"""Infer PostgreSQL schema from MongoDB collection and create tables."""
|
|
|
|
# Map MongoDB/Python types to PostgreSQL types
|
|
TYPE_MAPPING = {
|
|
'null': 'TEXT',
|
|
'boolean': 'BOOLEAN',
|
|
'integer': 'BIGINT',
|
|
'float': 'DOUBLE PRECISION',
|
|
'string': 'TEXT',
|
|
'array': 'JSONB',
|
|
'object': 'JSONB',
|
|
'date': 'TIMESTAMP',
|
|
'objectid': 'TEXT',
|
|
}
|
|
|
|
def __init__(self, sample_size: int = 100):
|
|
"""
|
|
Initialize schema generator.
|
|
|
|
Args:
|
|
sample_size: Number of documents to sample for type inference.
|
|
"""
|
|
self.sample_size = sample_size
|
|
self.field_types: Dict[str, Set[str]] = {}
|
|
self.final_schema: Dict[str, str] = {}
|
|
|
|
def infer_type(self, value: Any) -> str:
|
|
"""
|
|
Infer the MongoDB type of a value.
|
|
|
|
Args:
|
|
value: The value to inspect.
|
|
|
|
Returns:
|
|
A type string suitable for TYPE_MAPPING.
|
|
"""
|
|
if value is None:
|
|
return 'null'
|
|
if isinstance(value, bool):
|
|
return 'boolean'
|
|
if isinstance(value, int) and not isinstance(value, bool):
|
|
return 'integer'
|
|
if isinstance(value, float):
|
|
return 'float'
|
|
if isinstance(value, str):
|
|
return 'string'
|
|
if isinstance(value, list):
|
|
return 'array'
|
|
if isinstance(value, dict):
|
|
return 'object'
|
|
# Handle MongoDB-specific types
|
|
if hasattr(value, '__class__'):
|
|
type_name = value.__class__.__name__.lower()
|
|
if 'objectid' in type_name:
|
|
return 'objectid'
|
|
if 'datetime' in type_name:
|
|
return 'date'
|
|
return 'string'
|
|
|
|
def merge_types(self, types: Set[str]) -> str:
|
|
"""
|
|
Merge multiple observed types for a field into a single PostgreSQL type.
|
|
Priority: object/array > string > float > integer > boolean > null
|
|
|
|
Args:
|
|
types: Set of type strings.
|
|
|
|
Returns:
|
|
The merged PostgreSQL type.
|
|
"""
|
|
if not types:
|
|
return self.TYPE_MAPPING['null']
|
|
|
|
# Priority order
|
|
priority = ['array', 'object', 'string', 'float', 'integer', 'boolean', 'null', 'date']
|
|
for ptype in priority:
|
|
if ptype in types:
|
|
return self.TYPE_MAPPING.get(ptype, self.TYPE_MAPPING['string'])
|
|
|
|
return self.TYPE_MAPPING['string']
|
|
|
|
def analyze_collection(
|
|
self,
|
|
mongo_db: Any,
|
|
collection_name: Any
|
|
) -> Dict[str, str]:
|
|
"""
|
|
Analyze a MongoDB collection and infer field types.
|
|
|
|
Args:
|
|
mongo_db: A pymongo database object or None if collection_name is a collection.
|
|
collection_name: Name of the collection (str) or a pymongo Collection object directly.
|
|
|
|
Returns:
|
|
Dictionary mapping field names to PostgreSQL types.
|
|
"""
|
|
self.field_types = {}
|
|
# Handle both cases: collection name string or collection object
|
|
if isinstance(collection_name, str):
|
|
collection = mongo_db[collection_name]
|
|
else:
|
|
# Assume it's already a pymongo Collection object
|
|
collection = collection_name
|
|
|
|
# Sample documents from the collection
|
|
sample = list(collection.find({}).limit(self.sample_size))
|
|
|
|
if not sample:
|
|
logger.warning(f"Collection '{collection_name}' is empty; no schema inferred.")
|
|
return {}
|
|
|
|
logger.info(f"Analyzing {len(sample)} documents from '{collection_name}'")
|
|
|
|
# Gather field types
|
|
for doc in sample:
|
|
for field_name, value in doc.items():
|
|
inferred = self.infer_type(value)
|
|
self.field_types.setdefault(field_name, set()).add(inferred)
|
|
|
|
# Merge types
|
|
self.final_schema = {}
|
|
for field_name, types in self.field_types.items():
|
|
merged_type = self.merge_types(types)
|
|
self.final_schema[field_name] = merged_type
|
|
logger.debug(f" {field_name}: {merged_type} (observed: {types})")
|
|
|
|
return self.final_schema
|
|
|
|
def generate_create_table_sql(
|
|
self,
|
|
table_name: str,
|
|
schema_name: str = 'public',
|
|
pk_field: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
) -> str:
|
|
"""
|
|
Generate a CREATE TABLE statement from the inferred schema.
|
|
|
|
Args:
|
|
table_name: Name of the table to create.
|
|
schema_name: PostgreSQL schema (default: 'public').
|
|
pk_field: Field to use as primary key (e.g., 'id', '_id', 'seriesid').
|
|
exclude_fields: List of field names to exclude.
|
|
|
|
Returns:
|
|
A CREATE TABLE SQL statement.
|
|
"""
|
|
if not self.final_schema:
|
|
raise ValueError("No schema has been analyzed. Call analyze_collection first.")
|
|
|
|
exclude_fields = exclude_fields or []
|
|
columns = []
|
|
|
|
for field_name, pg_type in self.final_schema.items():
|
|
if field_name in exclude_fields or field_name == '_id':
|
|
continue
|
|
columns.append(f" {field_name} {pg_type}")
|
|
|
|
# Add primary key constraint if specified
|
|
constraints = []
|
|
if pk_field and pk_field in self.final_schema:
|
|
constraints.append(f" PRIMARY KEY ({pk_field})")
|
|
|
|
all_lines = columns + constraints
|
|
columns_str = ',\n'.join(all_lines)
|
|
|
|
full_table_name = f"{schema_name}.{table_name}"
|
|
sql = f"""CREATE TABLE IF NOT EXISTS {full_table_name} (
|
|
{columns_str}
|
|
);"""
|
|
|
|
return sql
|
|
|
|
def create_table_in_postgres(
|
|
self,
|
|
engine: Any,
|
|
table_name: str,
|
|
schema_name: str = 'public',
|
|
pk_field: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
drop_existing: bool = False,
|
|
) -> None:
|
|
"""
|
|
Create a table in PostgreSQL based on the inferred schema.
|
|
|
|
Args:
|
|
engine: SQLAlchemy engine connected to PostgreSQL, or a dict with 'engine' key.
|
|
table_name: Name of the table to create.
|
|
schema_name: PostgreSQL schema (default: 'public').
|
|
pk_field: Field to use as primary key.
|
|
exclude_fields: List of field names to exclude.
|
|
drop_existing: If True, drop the table before creating.
|
|
"""
|
|
# Handle both dict (from db/functions.py) and engine object directly
|
|
if isinstance(engine, dict):
|
|
engine = engine['engine']
|
|
|
|
full_table_name = f"{schema_name}.{table_name}"
|
|
|
|
create_sql = self.generate_create_table_sql(
|
|
table_name,
|
|
schema_name=schema_name,
|
|
pk_field=pk_field,
|
|
exclude_fields=exclude_fields,
|
|
)
|
|
|
|
with engine.begin() as conn:
|
|
if drop_existing:
|
|
logger.info(f"Dropping existing table {full_table_name}...")
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {full_table_name};"))
|
|
|
|
logger.info(f"Creating table {full_table_name}...")
|
|
conn.execute(text(create_sql))
|
|
logger.info(f"Table {full_table_name} created successfully.")
|
|
|
|
def create_table_from_collection(
|
|
self,
|
|
mongo_db: Any,
|
|
collection_name: str,
|
|
engine: Any,
|
|
table_name: Optional[str] = None,
|
|
schema_name: str = 'dbo',
|
|
pk_field: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
drop_existing: bool = False,
|
|
) -> str:
|
|
"""
|
|
End-to-end: analyze a MongoDB collection and create a PostgreSQL table.
|
|
|
|
Args:
|
|
mongo_db: A pymongo database object.
|
|
collection_name: Name of the MongoDB collection.
|
|
engine: SQLAlchemy engine connected to PostgreSQL.
|
|
table_name: Name of the PostgreSQL table (defaults to collection_name).
|
|
schema_name: PostgreSQL schema (default: 'dbo').
|
|
pk_field: Field to use as primary key.
|
|
exclude_fields: List of field names to exclude.
|
|
drop_existing: If True, drop the table before creating.
|
|
|
|
Returns:
|
|
The generated SQL statement.
|
|
"""
|
|
if table_name is None:
|
|
table_name = collection_name
|
|
|
|
# Analyze the collection
|
|
self.analyze_collection(mongo_db, collection_name)
|
|
|
|
# Create the table
|
|
self.create_table_in_postgres(
|
|
engine,
|
|
table_name,
|
|
schema_name=schema_name,
|
|
pk_field=pk_field,
|
|
exclude_fields=exclude_fields,
|
|
drop_existing=drop_existing,
|
|
)
|
|
|
|
# Return the SQL for reference
|
|
return self.generate_create_table_sql(
|
|
table_name,
|
|
schema_name=schema_name,
|
|
pk_field=pk_field,
|
|
exclude_fields=exclude_fields,
|
|
)
|
|
|
|
|
|
# Convenience function
|
|
def create_postgres_table_from_mongo(
|
|
mongo_db: Any,
|
|
collection_name: str,
|
|
engine: Any,
|
|
table_name: Optional[str] = None,
|
|
schema_name: str = 'dbo',
|
|
pk_field: Optional[str] = None,
|
|
sample_size: int = 100,
|
|
drop_existing: bool = False,
|
|
) -> str:
|
|
"""
|
|
Convenience function to create a PostgreSQL table from a MongoDB collection.
|
|
|
|
Example:
|
|
from db.schema_generator import create_postgres_table_from_mongo
|
|
|
|
sql = create_postgres_table_from_mongo(
|
|
mongo_db=mgdb,
|
|
collection_name='series',
|
|
engine=pg_engine,
|
|
table_name='seriesdata',
|
|
schema_name='dbo',
|
|
pk_field='seriesid',
|
|
drop_existing=True,
|
|
)
|
|
print(f"Created table with schema:\\n{sql}")
|
|
|
|
Args:
|
|
mongo_db: A pymongo database object.
|
|
collection_name: Name of the MongoDB collection.
|
|
engine: SQLAlchemy engine connected to PostgreSQL.
|
|
table_name: Name of the PostgreSQL table (defaults to collection_name).
|
|
schema_name: PostgreSQL schema (default: 'dbo').
|
|
pk_field: Field to use as primary key (optional).
|
|
sample_size: Number of documents to sample (default: 100).
|
|
drop_existing: If True, drop the table before creating.
|
|
|
|
Returns:
|
|
The generated SQL statement.
|
|
"""
|
|
generator = MongoToPostgresSchemaGenerator(sample_size=sample_size)
|
|
return generator.create_table_from_collection(
|
|
mongo_db=mongo_db,
|
|
collection_name=collection_name,
|
|
engine=engine,
|
|
table_name=table_name,
|
|
schema_name=schema_name,
|
|
pk_field=pk_field,
|
|
drop_existing=drop_existing,
|
|
)
|
|
|
|
|
|
class MongoDocumentInserter:
|
|
"""Insert MongoDB documents into PostgreSQL tables with type conversion."""
|
|
|
|
def __init__(self, batch_size: int = 10000):
|
|
"""
|
|
Initialize the document inserter.
|
|
|
|
Args:
|
|
batch_size: Number of documents to insert per batch (default: 1000).
|
|
"""
|
|
self.batch_size = batch_size
|
|
|
|
def convert_value(self, value: Any) -> Any:
|
|
"""
|
|
Convert a MongoDB value to PostgreSQL-compatible type.
|
|
|
|
- MongoDB ObjectId → string
|
|
- datetime → string (ISO format)
|
|
- list/dict → JSON string (will be stored as JSONB)
|
|
- None → None
|
|
|
|
Args:
|
|
value: The value to convert.
|
|
|
|
Returns:
|
|
The converted value.
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
# Handle MongoDB ObjectId
|
|
if hasattr(value, '__class__') and 'ObjectId' in value.__class__.__name__:
|
|
return str(value)
|
|
|
|
# Handle datetime
|
|
if hasattr(value, 'isoformat'):
|
|
return value.isoformat()
|
|
|
|
# Handle lists and dicts (will be stored as JSONB)
|
|
if isinstance(value, (list, dict)):
|
|
return json.dumps(value)
|
|
|
|
return value
|
|
|
|
def prepare_documents(
|
|
self,
|
|
documents: List[Dict[str, Any]],
|
|
column_names: List[str],
|
|
exclude_fields: Optional[List[str]] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Prepare documents for insertion: convert types, exclude fields, etc.
|
|
|
|
Args:
|
|
documents: List of documents (dicts) from MongoDB.
|
|
column_names: Column names in the target table.
|
|
exclude_fields: Fields to exclude (e.g., ['_id']).
|
|
|
|
Returns:
|
|
List of prepared documents ready for insertion.
|
|
"""
|
|
exclude_fields = exclude_fields or ['_id']
|
|
prepared = []
|
|
|
|
for doc in documents:
|
|
row = {}
|
|
for col in column_names:
|
|
if col not in exclude_fields and col in doc:
|
|
row[col] = self.convert_value(doc[col])
|
|
elif col not in exclude_fields:
|
|
row[col] = None
|
|
prepared.append(row)
|
|
|
|
return prepared
|
|
|
|
def insert_documents(
|
|
self,
|
|
engine: Any,
|
|
table_name: str,
|
|
documents: List[Dict[str, Any]],
|
|
schema_name: str = 'dbo',
|
|
on_conflict: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
skip_null: bool = False,
|
|
) -> int:
|
|
"""
|
|
Insert documents into a PostgreSQL table in batches.
|
|
|
|
Args:
|
|
engine: SQLAlchemy engine (or dict with 'engine' key).
|
|
table_name: Target table name.
|
|
documents: List of documents (dicts) to insert.
|
|
schema_name: PostgreSQL schema (default: 'dbo').
|
|
on_conflict: ON CONFLICT clause (e.g., "DO NOTHING" or
|
|
"DO UPDATE SET field=EXCLUDED.field").
|
|
exclude_fields: Fields to exclude from insert (default: ['_id']).
|
|
skip_null: If True, skip None values in the INSERT (uses COALESCE).
|
|
|
|
Returns:
|
|
Total number of documents inserted.
|
|
"""
|
|
if isinstance(engine, dict):
|
|
engine = engine['engine']
|
|
|
|
if not documents:
|
|
logger.warning("No documents to insert.")
|
|
return 0
|
|
|
|
exclude_fields = exclude_fields or ['_id']
|
|
|
|
full_table_name = f"{schema_name}.{table_name}"
|
|
total_inserted = 0
|
|
|
|
# Reflect target table to get column definitions and types
|
|
metadata = MetaData()
|
|
try:
|
|
table = Table(table_name, metadata, autoload_with=engine, schema=schema_name)
|
|
column_objs = {c.name: c for c in table.columns}
|
|
column_names = list(column_objs.keys())
|
|
except Exception:
|
|
# Fallback: derive column names from first document
|
|
column_names = [k for k in documents[0].keys() if k not in exclude_fields]
|
|
column_objs = {}
|
|
|
|
# Prepare and convert documents according to target column types
|
|
prepared: List[Dict[str, Any]] = []
|
|
for doc in documents:
|
|
# create case-insensitive key map for doc
|
|
key_map = {k.lower(): k for k in doc.keys()}
|
|
row: Dict[str, Any] = {}
|
|
for col in column_names:
|
|
if exclude_fields and col in exclude_fields:
|
|
continue
|
|
# find matching key in document (case-insensitive)
|
|
val = None
|
|
if col in doc:
|
|
val = doc[col]
|
|
elif col.lower() in key_map:
|
|
val = doc[key_map[col.lower()]]
|
|
|
|
# Convert based on column type if available
|
|
colobj = column_objs.get(col)
|
|
if colobj is not None and val is not None:
|
|
try:
|
|
from datetime import datetime
|
|
import sqlalchemy as sa
|
|
|
|
ctype = colobj.type
|
|
# Integer target: convert ISO datetime strings to epoch
|
|
if isinstance(ctype, (sa.Integer, sa.BigInteger)):
|
|
if isinstance(val, str):
|
|
# try parse datetime string
|
|
try:
|
|
dt = datetime.fromisoformat(val)
|
|
val = int(dt.timestamp())
|
|
except Exception:
|
|
# try numeric parse
|
|
try:
|
|
val = int(val)
|
|
except Exception:
|
|
pass
|
|
elif isinstance(val, float):
|
|
val = int(val)
|
|
# Datetime target: convert epoch ints to datetime
|
|
elif isinstance(ctype, (sa.DateTime, sa.TIMESTAMP)):
|
|
if isinstance(val, (int, float)):
|
|
val = datetime.fromtimestamp(val)
|
|
elif isinstance(val, str):
|
|
try:
|
|
val = datetime.fromisoformat(val)
|
|
except Exception:
|
|
pass
|
|
# JSON target: ensure Python dict/list
|
|
elif 'JSON' in type(ctype).__name__.upper() or 'JSON' in str(ctype).upper():
|
|
if isinstance(val, str):
|
|
try:
|
|
val = json.loads(val)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
# ignore conversion errors and use original value
|
|
pass
|
|
|
|
# Final conversion for general values
|
|
if val is None:
|
|
row[col] = None
|
|
else:
|
|
row[col] = self.convert_value(val) if col not in (exclude_fields or []) else None
|
|
|
|
prepared.append(row)
|
|
|
|
# Build INSERT statement template using reflected column order
|
|
cols_str = ', '.join(column_names)
|
|
placeholders = ', '.join([':' + col for col in column_names])
|
|
insert_sql = f"INSERT INTO {full_table_name} ({cols_str}) VALUES ({placeholders})"
|
|
if on_conflict:
|
|
insert_sql += f" {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:
|
|
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:
|
|
conn.execute(text(insert_sql), batch)
|
|
batch_count = len(batch)
|
|
total_inserted += batch_count
|
|
logger.info(
|
|
f"Inserted {batch_count} documents into {full_table_name} (total: {total_inserted})"
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error inserting batch into {full_table_name}: {e}")
|
|
raise
|
|
|
|
return total_inserted
|
|
|
|
def insert_from_collection(
|
|
self,
|
|
engine: Any,
|
|
table_name: str,
|
|
collection: Any,
|
|
schema_name: str = 'dbo',
|
|
on_conflict: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
query_filter: Optional[Dict] = None,
|
|
limit: Optional[int] = None,
|
|
) -> int:
|
|
"""
|
|
Insert documents from a MongoDB collection into PostgreSQL.
|
|
|
|
Args:
|
|
engine: SQLAlchemy engine (or dict with 'engine' key).
|
|
table_name: Target PostgreSQL table name.
|
|
collection: PyMongo collection object.
|
|
schema_name: PostgreSQL schema (default: 'dbo').
|
|
on_conflict: ON CONFLICT clause.
|
|
exclude_fields: Fields to exclude (default: ['_id']).
|
|
query_filter: MongoDB query filter (default: {}).
|
|
limit: Maximum documents to insert (default: None for all).
|
|
|
|
Returns:
|
|
Total number of documents inserted.
|
|
"""
|
|
query_filter = query_filter or {}
|
|
|
|
logger.info(f"Fetching documents from MongoDB collection...")
|
|
cursor = collection.find(query_filter)
|
|
|
|
if limit:
|
|
cursor = cursor.limit(limit)
|
|
|
|
documents = list(cursor)
|
|
logger.info(f"Retrieved {len(documents)} documents.")
|
|
|
|
return self.insert_documents(
|
|
engine=engine,
|
|
table_name=table_name,
|
|
documents=documents,
|
|
schema_name=schema_name,
|
|
on_conflict=on_conflict,
|
|
exclude_fields=exclude_fields,
|
|
)
|
|
|
|
|
|
# Convenience function for inserting documents
|
|
def insert_mongo_documents_to_postgres(
|
|
engine: Any,
|
|
table_name: str,
|
|
documents: List[Dict[str, Any]],
|
|
schema_name: str = 'dbo',
|
|
on_conflict: Optional[str] = None,
|
|
exclude_fields: Optional[List[str]] = None,
|
|
batch_size: int = 1000,
|
|
) -> int:
|
|
"""
|
|
Convenience function to insert documents into PostgreSQL.
|
|
|
|
Example:
|
|
from db.schema_generator import insert_mongo_documents_to_postgres
|
|
|
|
count = insert_mongo_documents_to_postgres(
|
|
engine=dbengine,
|
|
table_name='seriesdata',
|
|
documents=series_docs,
|
|
schema_name='dbo',
|
|
on_conflict="DO NOTHING",
|
|
exclude_fields=['_id'],
|
|
)
|
|
print(f"Inserted {count} documents")
|
|
|
|
Args:
|
|
engine: SQLAlchemy engine (or dict with 'engine' key).
|
|
table_name: Target table name.
|
|
documents: List of documents to insert.
|
|
schema_name: PostgreSQL schema (default: 'dbo').
|
|
on_conflict: ON CONFLICT clause (optional).
|
|
exclude_fields: Fields to exclude (default: ['_id']).
|
|
batch_size: Number of documents per batch (default: 1000).
|
|
|
|
Returns:
|
|
Total number of documents inserted.
|
|
"""
|
|
inserter = MongoDocumentInserter(batch_size=batch_size)
|
|
return inserter.insert_documents(
|
|
engine=engine,
|
|
table_name=table_name,
|
|
documents=documents,
|
|
schema_name=schema_name,
|
|
on_conflict=on_conflict,
|
|
exclude_fields=exclude_fields,
|
|
)
|
|
|