Refactor update_mongo.py to enhance JSON downloading and file handling; update MongoDB host in test_tableddl_from_yaml.py

2025-12-11 17:53:41 -05:00
parent a18c6a7386
commit e975b947fa
6 changed files with 397 additions and 227097 deletions
@@ -0,0 +1,205 @@
+#!/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()