Refactor MongoDB to PostgreSQL data synchronization: remove json2postgres.py, add mongodb2postgres.py for improved schema generation and data insertion
This commit is contained in:
+7
-7
@@ -114,13 +114,13 @@ class settype(object):
|
||||
)
|
||||
currentDir = os.path.dirname(os.path.abspath(__file__))
|
||||
currentDir = os.path.dirname(currentDir)
|
||||
ddl_file_path = f"{currentDir}/db/ddl/recreate_updates_schema.sql"
|
||||
with open(ddl_file_path, "r") as file:
|
||||
sql = file.read()
|
||||
file
|
||||
with engine.connect() as conn:
|
||||
with conn as cursor:
|
||||
cursor.execute(text(sql))
|
||||
# ddl_file_path = f"{currentDir}/db/ddl/recreate_updates_schema.sql"
|
||||
# with open(ddl_file_path, "r") as file:
|
||||
# sql = file.read()
|
||||
# file
|
||||
# with engine.connect() as conn:
|
||||
# with conn as cursor:
|
||||
# cursor.execute(text(sql))
|
||||
updatesBase.prepare(engine, reflect=True, schema="updates")
|
||||
dboBase.prepare(engine, reflect=True, schema="dbo")
|
||||
metadata = sqlalchemy.MetaData(schema="updates")
|
||||
|
||||
@@ -0,0 +1,654 @@
|
||||
"""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 = 1000):
|
||||
"""
|
||||
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}"
|
||||
|
||||
with engine.begin() as conn:
|
||||
if on_conflict:
|
||||
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:
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -1,287 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from pymongo import MongoClient
|
||||
import json
|
||||
import psycopg2
|
||||
import pandas as pd
|
||||
from collections.abc import MutableMapping
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def flatten_dict(d: MutableMapping, sep: str = '.') -> MutableMapping:
|
||||
[flat_dict] = pd.json_normalize(d, sep=sep).to_dict(orient='records')
|
||||
return flat_dict
|
||||
|
||||
|
||||
def delete_none(_dict):
|
||||
"""
|
||||
Delete None values recursively from all of the dictionaries, tuples,
|
||||
lists, sets
|
||||
"""
|
||||
if isinstance(_dict, dict):
|
||||
for key, value in list(_dict.items()):
|
||||
if isinstance(value, (list, dict, tuple, set)):
|
||||
_dict[key] = delete_none(value)
|
||||
elif value is None or key is None:
|
||||
del _dict[key]
|
||||
|
||||
elif isinstance(_dict, (list, set, tuple)):
|
||||
_dict = type(_dict)(
|
||||
delete_none(item) for item in _dict if item is not None
|
||||
)
|
||||
return _dict
|
||||
|
||||
|
||||
def process_data(data, tablename):
|
||||
inserts = []
|
||||
failed_count = 0
|
||||
seriesconflict = "ON CONFLICT (seriesid) DO UPDATE SET"
|
||||
epconflict = "ON CONFLICT (episodeid, seriesid) DO UPDATE SET"
|
||||
crewconflict = "ON CONFLICT (crewid) DO UPDATE SET"
|
||||
actorconflict = "ON CONFLICT (actorid) DO UPDATE SET"
|
||||
characterconflict = "ON CONFLICT (seriesid, characterid) DO UPDATE SET"
|
||||
for document in data:
|
||||
if document['name'].lower() == 'too many requests':
|
||||
print("Skipping document due to 'Too Many Requests' in "
|
||||
"document data")
|
||||
continue
|
||||
else:
|
||||
new_document = delete_none(document)
|
||||
if '_id' in new_document:
|
||||
del new_document['_id']
|
||||
try:
|
||||
flatten_data = flatten_dict(new_document)
|
||||
except TypeError as e:
|
||||
print(
|
||||
(
|
||||
f"Unable to flatten data for current document. {e}\n"
|
||||
"Exitting!"
|
||||
)
|
||||
)
|
||||
failed_count += 1
|
||||
else:
|
||||
column_list = []
|
||||
data_list = []
|
||||
for key, value in flatten_data.items():
|
||||
if isinstance(value, str):
|
||||
value = value.replace('"', "'").replace("'", "''")
|
||||
if key.lower() == '_links.self.href':
|
||||
key = 'apilink'
|
||||
if key.lower() == '_links.show.href':
|
||||
key = 'apishowlink'
|
||||
if key.lower() == '_links.show.name':
|
||||
key = 'apishowname'
|
||||
if '.' in key.lower():
|
||||
key = key.replace('.', '')
|
||||
if tablename == 'seriesdata' and (
|
||||
key.lower() == 'name' or key.lower() == 'type'
|
||||
):
|
||||
key = f"series_{key}"
|
||||
if tablename == 'seriesdata' and key.lower() == 'id':
|
||||
key = "seriesid"
|
||||
if tablename == 'crewdata' and key.lower() == 'id':
|
||||
key = "crewid"
|
||||
if (
|
||||
tablename == 'actordata' and (
|
||||
key.lower() == 'name' or key.lower() == 'number' or
|
||||
key.lower() == 'type'
|
||||
)
|
||||
):
|
||||
key = f"actor{key}"
|
||||
if tablename == 'actordata' and key.lower() == 'id':
|
||||
key = "actorid"
|
||||
if tablename == 'characterdata' and key.lower() == 'id':
|
||||
key = "characterid"
|
||||
if tablename == 'characterdata' and key.lower() == 'name':
|
||||
key = "charactername"
|
||||
if tablename == 'characterdata' and key.lower() == 'id':
|
||||
key = "seriesid, characterid"
|
||||
if (
|
||||
(
|
||||
tablename == 'epdata' or
|
||||
tablename == 'actordata'
|
||||
) and (
|
||||
key.lower() == 'name' or
|
||||
key.lower() == 'number' or
|
||||
key.lower() == 'type'
|
||||
)
|
||||
):
|
||||
key = f"episode_{key}"
|
||||
if key.lower() == 'language':
|
||||
key = "language_name"
|
||||
if tablename == 'epdata' and key.lower() == 'id':
|
||||
key = "episodeid"
|
||||
if tablename == 'actordata' and 'seriesid' in key.lower():
|
||||
continue
|
||||
else:
|
||||
column_list.append(key)
|
||||
if (
|
||||
'genres' in key.lower() or
|
||||
'schedule' in key.lower()
|
||||
):
|
||||
if isinstance(value, list):
|
||||
value = ', '.join(value)
|
||||
if isinstance(value, int) or isinstance(value, float):
|
||||
value = str(json.dumps(value))
|
||||
data_list.append(f"'{value}'")
|
||||
columns = ", ".join(column_list)
|
||||
column_data = ", ".join(data_list)
|
||||
conflict_list = []
|
||||
for indexer, column in enumerate(column_list):
|
||||
conflict_list.append(
|
||||
f"{column_list[indexer]} = {data_list[indexer]}"
|
||||
)
|
||||
conflict_data = ', '.join(conflict_list).replace("\\", "")
|
||||
conflict_clause = ''
|
||||
if tablename == 'seriesdata':
|
||||
conflict_clause = seriesconflict
|
||||
elif tablename == 'epdata':
|
||||
conflict_clause = epconflict
|
||||
elif tablename == 'crewdata':
|
||||
conflict_clause = crewconflict
|
||||
elif tablename == 'actordata':
|
||||
conflict_clause = actorconflict
|
||||
elif tablename == 'characterdata':
|
||||
conflict_clause = characterconflict
|
||||
insert_command = (
|
||||
(
|
||||
f"Insert Into updates.{tablename} ({columns}) "
|
||||
f"values ({column_data}) "
|
||||
f"{conflict_clause} {conflict_data};"
|
||||
)
|
||||
)
|
||||
inserts.append(insert_command)
|
||||
return inserts, failed_count
|
||||
|
||||
|
||||
try:
|
||||
connection_string = (
|
||||
(
|
||||
'postgres://postgres:Optimus0329@192.168.128.7:5432/media_dbsync'
|
||||
'?options=-csearch_path%3Ddbo,public,updates'
|
||||
)
|
||||
)
|
||||
pgclient = psycopg2.connect(connection_string)
|
||||
pgcursor = pgclient.cursor()
|
||||
# print(pgclient.get_dsn_parameters(), "\n")
|
||||
print("Connected to postgresql database server.")
|
||||
pgcursor.execute("SELECT version();")
|
||||
record = pgcursor.fetchone()
|
||||
print("You are connected to - ", record, "\n")
|
||||
pgclient.set_isolation_level(0)
|
||||
except Exception as e:
|
||||
print(f"Unable to connect to postgresql database server:{e}\nExitting!")
|
||||
exit()
|
||||
|
||||
|
||||
host = '192.168.128.8'
|
||||
port = 27017
|
||||
print(f"Connecting to mongodb at {host}:{port}")
|
||||
|
||||
try:
|
||||
connection_string = (
|
||||
f"mongodb://{host}:{port}"
|
||||
)
|
||||
mgclient = MongoClient(connection_string)
|
||||
mgdb = mgclient['test2']
|
||||
except BaseException as e:
|
||||
print(f"Unable to connect to mongodb.{e}\nExitting!")
|
||||
exit()
|
||||
else:
|
||||
mgseries = mgdb["series"]
|
||||
mgseries.create_index('id', unique=True)
|
||||
mgepisodes = mgdb["episodes"]
|
||||
mgepisodes.create_index('id', unique=True)
|
||||
mgactors = mgdb["actors"]
|
||||
mgactors.create_index('id', unique=True)
|
||||
mgcharacters = mgdb["characters"]
|
||||
mgcharacters.create_index('id', unique=True)
|
||||
mgcrew = mgdb["crew"]
|
||||
mgcrew.create_index('id', unique=True)
|
||||
print(
|
||||
f"Connection to mongodb at {host}:"
|
||||
f"{port} successful."
|
||||
)
|
||||
|
||||
print("Creating Insert commands from collection data.")
|
||||
|
||||
print("Collecting series data for series inserts.")
|
||||
series_list = mgseries.find({})
|
||||
inserts, failed_count = process_data(series_list, 'seriesdata')
|
||||
print(f"Total number of series inserts: {len(inserts)}")
|
||||
for i in tqdm(range(len(inserts))):
|
||||
try:
|
||||
insert = inserts[i]
|
||||
pgcursor.execute(insert)
|
||||
except Exception as e:
|
||||
print(f"Unable to insert data into seriesdata table. {e}")
|
||||
failed_count += 1
|
||||
print(f"Total number of failed series inserts: {failed_count}")
|
||||
pgclient.commit()
|
||||
|
||||
epinserts = []
|
||||
failures = 0
|
||||
episode_array = mgepisodes.find({})
|
||||
print("Collecting episode data for episode inserts.")
|
||||
inserts, failed_count = process_data(episode_array, 'epdata')
|
||||
print(f"Total number of episode inserts: {len(inserts)}")
|
||||
for i in tqdm(range(len(inserts))):
|
||||
try:
|
||||
insert = inserts[i]
|
||||
pgcursor.execute(insert)
|
||||
except Exception as e:
|
||||
print(f"Unable to insert data into epdata table. {e}")
|
||||
failed_count += 1
|
||||
print(f"Total number of failed episode inserts: {failed_count}")
|
||||
pgclient.commit()
|
||||
|
||||
failures = 0
|
||||
actor_array = mgactors.find({})
|
||||
print("Collecting actor data for actor inserts.")
|
||||
inserts, failed_count = process_data(actor_array, 'actordata')
|
||||
print(f"Total number of actor inserts: {len(inserts)}")
|
||||
for i in tqdm(range(len(inserts))):
|
||||
try:
|
||||
insert = inserts[i]
|
||||
pgcursor.execute(insert)
|
||||
except Exception as e:
|
||||
print(f"Unable to insert data into actors table. {e}")
|
||||
failed_count += 1
|
||||
print(f"Total number of failed actor inserts: {failed_count}")
|
||||
pgclient.commit()
|
||||
|
||||
failures = 0
|
||||
character_array = mgcharacters.find({})
|
||||
print("Collecting actor data for actor inserts.")
|
||||
inserts, failed_count = process_data(character_array, 'characterdata')
|
||||
print(f"Total number of character inserts: {len(inserts)}")
|
||||
for i in tqdm(range(len(inserts))):
|
||||
try:
|
||||
insert = inserts[i]
|
||||
pgcursor.execute(insert)
|
||||
except Exception as e:
|
||||
print(f"Unable to insert data into character table. {e}")
|
||||
failed_count += 1
|
||||
print(f"Total number of failed character inserts: {failed_count}")
|
||||
pgclient.commit()
|
||||
|
||||
failures = 0
|
||||
crew_array = mgcrew.find({})
|
||||
print("Collecting crew data for crew inserts.")
|
||||
inserts, failed_count = process_data(crew_array, 'crewdata')
|
||||
print(f"Total number of crew inserts: {len(inserts)}")
|
||||
for i in tqdm(range(len(inserts))):
|
||||
try:
|
||||
insert = inserts[i]
|
||||
pgcursor.execute(insert)
|
||||
except Exception as e:
|
||||
print(f"Unable to insert data into epdata table. {e}")
|
||||
failed_count += 1
|
||||
print(f"Total number of failed crew inserts: {failed_count}")
|
||||
pgclient.commit()
|
||||
|
||||
pgclient.close()
|
||||
print("PostgreSQL connection is closed")
|
||||
|
||||
|
||||
exit()
|
||||
@@ -0,0 +1,120 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
import time
|
||||
import sys
|
||||
from tqdm import tqdm
|
||||
import logs.Logger
|
||||
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)
|
||||
|
||||
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=1000)
|
||||
|
||||
# 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'),
|
||||
}
|
||||
|
||||
# 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}"
|
||||
|
||||
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}")
|
||||
|
||||
# # 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())
|
||||
|
||||
Reference in New Issue
Block a user