206 lines
7.0 KiB
Python
206 lines
7.0 KiB
Python
#!/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()
|