#!/usr/bin/env python3 """Convert a MongoDB collection into a PostgreSQL table. Contains `convert_mongo_to_pg(mongo_uri, mongo_db, collection_name, pg_dsn, pg_table=None, ...)` and a small CLI. Notes: - Uses `pymongo` to read documents and `psycopg2` to write to Postgres. - Infers basic column types from a sample of documents; nested objects/arrays are stored as `JSONB`. """ from __future__ import annotations import argparse import json import logging from typing import Any, Dict, Iterable, List, Optional import pymongo import psycopg2 import psycopg2.extras as pgextras import psycopg2.sql as sql LOG = logging.getLogger("mongo2pgsql") def _detect_type(value: Any) -> str: 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 "real" if isinstance(value, (dict, list)): return "jsonb" return "text" def _merge_types(types: List[str]) -> str: # priority: jsonb > text > real > integer > boolean > null s = set(types) if "jsonb" in s: return "jsonb" if "text" in s: return "text" if "real" in s: return "real" if "integer" in s: return "integer" if "boolean" in s: return "boolean" return "text" def _pg_type_for(kind: str) -> str: return { "integer": "BIGINT", "real": "DOUBLE PRECISION", "boolean": "BOOLEAN", "jsonb": "JSONB", "text": "TEXT", }.get(kind, "TEXT") def convert_mongo_to_pg( mongo_uri: str, mongo_db: str, collection_name: str, pg_dsn: str, pg_table: Optional[str] = None, sample_size: int = 500, batch_size: int = 1000, create_table: bool = True, pk_field: str = "_id", ) -> None: """Copy a MongoDB collection to PostgreSQL. - `pg_table` defaults to `collection_name` if omitted. - Nested documents/arrays are stored as JSONB. - `_id` is stored as TEXT by default. """ if pg_table is None: pg_table = collection_name client = pymongo.MongoClient(mongo_uri) coll = client[mongo_db][collection_name] sample = list(coll.find({}, projection=None, limit=sample_size)) if not sample: raise ValueError("collection is empty or not found") # gather all field names and observed types field_types: Dict[str, List[str]] = {} for doc in sample: for k, v in doc.items(): t = _detect_type(v) field_types.setdefault(k, []).append(t) # finalize types final_types: Dict[str, str] = {} for k, types in field_types.items(): merged = _merge_types(types) final_types[k] = _pg_type_for(merged) # ensure primary key present if pk_field not in final_types: final_types[pk_field] = "TEXT" cols = list(final_types.items()) # create table if create_table: # Build CREATE TABLE statement using psycopg2.sql for safe identifiers col_defs = [] for name, ptype in cols: col_defs.append(sql.SQL("{} {}").format(sql.Identifier(name), sql.SQL(ptype))) create_body = sql.SQL(", ").join(col_defs) if pk_field in final_types: pk_clause = sql.SQL(", PRIMARY KEY ({})").format(sql.SQL(",").join(map(sql.Identifier, [pk_field]))) create_body = create_body + pk_clause create_stmt = sql.SQL("CREATE TABLE IF NOT EXISTS {} ({})").format(sql.Identifier(pg_table), create_body) with psycopg2.connect(pg_dsn) as pgconn: with pgconn.cursor() as cur: LOG.info("Creating table %s", pg_table) cur.execute(create_stmt) pgconn.commit() # stream documents and insert insert_cols = [n for n, _ in cols] with psycopg2.connect(pg_dsn) as pgconn: with pgconn.cursor() as cur: # build INSERT statement for execute_values: it expects a single %s base_insert = sql.SQL("INSERT INTO {} ({}) VALUES %s").format( sql.Identifier(pg_table), sql.SQL(', ').join(map(sql.Identifier, insert_cols)), ).as_string(pgconn) batch = [] # Use an explicit session when using no_cursor_timeout to avoid pymongo warning with client.start_session() as session: cursor = coll.find({}, no_cursor_timeout=True, session=session).batch_size(batch_size) try: for doc in cursor: row = [] for col_name in insert_cols: v = doc.get(col_name) if v is None: row.append(None) continue # convert ObjectId and other non-JSON types if isinstance(v, (dict, list)): row.append(pgextras.Json(v)) else: row.append(str(v) if col_name == pk_field else v) batch.append(tuple(row)) if len(batch) >= batch_size: pgextras.execute_values(cur, base_insert, batch, template=None, page_size=batch_size) pgconn.commit() batch = [] finally: try: cursor.close() except Exception: pass if batch: pgextras.execute_values(cur, base_insert, batch, template=None, page_size=len(batch)) pgconn.commit() LOG.info("Finished copying collection %s.%s to %s", mongo_db, collection_name, pg_table) def _cli() -> None: parser = argparse.ArgumentParser(description="Copy a MongoDB collection to PostgreSQL") parser.add_argument("mongo_uri", help="MongoDB connection URI, e.g. mongodb://user:pw@host:27017") parser.add_argument("mongo_db", help="MongoDB database name") parser.add_argument("collection", help="Collection name to copy") parser.add_argument("pg_dsn", help="Postgres DSN, e.g. \"host=.. dbname=.. user=.. password=..\"") parser.add_argument("--pg-table", help="Target Postgres table name (defaults to collection name)") parser.add_argument("--sample", type=int, default=500, help="Sample size to infer schema") parser.add_argument("--batch", type=int, default=1000, help="Insert batch size") parser.add_argument("--no-create", dest="create", action="store_false", help="Don't attempt to create table in PG") parser.add_argument("--verbose", "-v", action="store_true") args = parser.parse_args() logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) convert_mongo_to_pg( args.mongo_uri, args.mongo_db, args.collection, args.pg_dsn, pg_table=args.pg_table, sample_size=args.sample, batch_size=args.batch, create_table=args.create, ) if __name__ == "__main__": _cli()