655 lines
23 KiB
Python
655 lines
23 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 = 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,
|
|
)
|
|
|