"""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, )