MongoDB: Various fixes and improvements

CHANGES.md

@@ -3,6 +3,18 @@
 
 ## Unreleased
 - Processor: Updated Kinesis Lambda processor to understand AWS DMS
+- MongoDB: Fix missing output on STDOUT for `migr8 export`
+- MongoDB: Improve timestamp parsing by using `python-dateutil`
+- MongoDB: Converge `_id` input field to `id` column instead of dropping it
+- MongoDB: Make user interface use stderr, so stdout is for data only
+- MongoDB: Make `migr8 extract` write to stdout by default
+- MongoDB: Make `migr8 translate` read from stdin by default
+- MongoDB: Improve user interface messages
+- MongoDB: Strip single leading underscore character from all top-level fields
+- MongoDB: Map OID types to CrateDB TEXT columns
+- MongoDB: Make `migr8 extract` and `migr8 export` accept the `--limit` option
+- MongoDB: Fix indentation in prettified SQL output of `migr8 translate`
+- MongoDB: Add capability to give type hints and add transformations
 
 ## 2024/07/25 v0.0.16
 - `ctk load table`: Added support for MongoDB Change Streams

cratedb_toolkit/api/main.py

@@ -4,6 +4,7 @@
 import logging
 import typing as t
 from abc import abstractmethod
+from pathlib import Path
 
 from yarl import URL
 
@@ -18,7 +19,7 @@
 
 class ClusterBase(abc.ABC):
     @abstractmethod
-    def load_table(self, resource: InputOutputResource, target: TableAddress):
+    def load_table(self, resource: InputOutputResource, target: TableAddress, transformation: Path):
         raise NotImplementedError("Child class needs to implement this method")
 
 
@@ -35,7 +36,9 @@ class ManagedCluster(ClusterBase):
     def __post_init__(self):
         logger.info(f"Connecting to CrateDB Cloud Cluster: {self.cloud_id}")
 
-    def load_table(self, resource: InputOutputResource, target: t.Optional[TableAddress] = None):
+    def load_table(
+        self, resource: InputOutputResource, target: t.Optional[TableAddress] = None, transformation: Path = None
+    ):
         """
         Load data into a database table on CrateDB Cloud.
 
@@ -96,7 +99,7 @@ class StandaloneCluster(ClusterBase):
     address: DatabaseAddress
     info: t.Optional[ClusterInformation] = None
 
-    def load_table(self, resource: InputOutputResource, target: TableAddress):
+    def load_table(self, resource: InputOutputResource, target: TableAddress, transformation: Path = None):
         """
         Load data into a database table on a standalone CrateDB Server.
 
@@ -109,11 +112,11 @@ def load_table(self, resource: InputOutputResource, target: TableAddress):
         """
         source_url = resource.url
         target_url = self.address.dburi
+        source_url_obj = URL(source_url)
         if source_url.startswith("influxdb"):
             from cratedb_toolkit.io.influxdb import influxdb_copy
 
             http_scheme = "http://"
-            source_url_obj = URL(source_url)
             if asbool(source_url_obj.query.get("ssl")):
                 http_scheme = "https://"
             source_url = source_url.replace("influxdb2://", http_scheme)
@@ -130,7 +133,13 @@ def load_table(self, resource: InputOutputResource, target: TableAddress):
             else:
                 from cratedb_toolkit.io.mongodb.api import mongodb_copy
 
-                if not mongodb_copy(source_url, target_url, progress=True):
+                if not mongodb_copy(
+                    source_url,
+                    target_url,
+                    transformation=transformation,
+                    limit=int(source_url_obj.query.get("limit", 0)),
+                    progress=True,
+                ):
                     msg = "Data loading failed"
                     logger.error(msg)
                     raise OperationFailed(msg)

cratedb_toolkit/io/cli.py

@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 
 import click
 from click_aliases import ClickAliasedGroup
@@ -35,6 +36,7 @@ def cli(ctx: click.Context, verbose: bool, debug: bool):
 @click.option("--table", envvar="CRATEDB_TABLE", type=str, required=False, help="Table where to import the data")
 @click.option("--format", "format_", type=str, required=False, help="File format of the import resource")
 @click.option("--compression", type=str, required=False, help="Compression format of the import resource")
+@click.option("--transformation", type=Path, required=False, help="Path to Zyp transformation file")
 @click.pass_context
 def load_table(
     ctx: click.Context,
@@ -46,6 +48,7 @@ def load_table(
     table: str,
     format_: str,
     compression: str,
+    transformation: Path,
 ):
     """
     Import data into CrateDB and CrateDB Cloud clusters.
@@ -82,4 +85,4 @@ def load_table(
         cluster = StandaloneCluster(address=address)
     else:
         raise NotImplementedError("Unable to select backend")
-    return cluster.load_table(resource=resource, target=target)
+    return cluster.load_table(resource=resource, target=target, transformation=transformation)

cratedb_toolkit/io/mongodb/api.py

@@ -1,5 +1,6 @@
 import argparse
 import logging
+from pathlib import Path
 
 from cratedb_toolkit.io.mongodb.cdc import MongoDBCDCRelayCrateDB
 from cratedb_toolkit.io.mongodb.core import export, extract, translate
@@ -10,7 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-def mongodb_copy(source_url, target_url, progress: bool = False):
+def mongodb_copy(source_url, target_url, transformation: Path = None, limit: int = 0, progress: bool = False):
     """
     Synopsis
     --------
@@ -36,7 +37,12 @@ def mongodb_copy(source_url, target_url, progress: bool = False):
     # 1. Extract schema from MongoDB collection.
     logger.info(f"Extracting schema from MongoDB: {mongodb_database}.{mongodb_collection}")
     extract_args = argparse.Namespace(
-        url=str(mongodb_uri), database=mongodb_database, collection=mongodb_collection, scan="full"
+        url=str(mongodb_uri),
+        database=mongodb_database,
+        collection=mongodb_collection,
+        scan="full",
+        transformation=transformation,
+        limit=limit,
     )
     mongodb_schema = extract(extract_args)
     count = mongodb_schema[mongodb_collection]["count"]
@@ -64,7 +70,13 @@ def mongodb_copy(source_url, target_url, progress: bool = False):
         f"Transferring data from MongoDB to CrateDB: "
         f"source={mongodb_collection_address.fullname}, target={cratedb_table_address.fullname}"
     )
-    export_args = argparse.Namespace(url=str(mongodb_uri), database=mongodb_database, collection=mongodb_collection)
+    export_args = argparse.Namespace(
+        url=str(mongodb_uri),
+        database=mongodb_database,
+        collection=mongodb_collection,
+        transformation=transformation,
+        limit=limit,
+    )
     buffer = export(export_args)
     cr8_insert_json(infile=buffer, hosts=cratedb_address.httpuri, table=cratedb_table_address.fullname)
 

cratedb_toolkit/io/mongodb/cli.py

@@ -1,10 +1,17 @@
 import argparse
 import json
+import sys
+import typing as t
+from pathlib import Path
 
-import rich
+from rich.console import Console
 
 from cratedb_toolkit import __version__
 from cratedb_toolkit.io.mongodb.core import export, extract, translate
+from cratedb_toolkit.util.common import setup_logging
+
+console = Console(stderr=True)
+rich = console
 
 
 def extract_parser(subargs):
@@ -19,7 +26,9 @@ def extract_parser(subargs):
         choices=["full", "partial"],
         help="Whether to fully scan the MongoDB collections or only partially.",
     )
-    parser.add_argument("-o", "--out", default="mongodb_schema.json")
+    parser.add_argument("--limit", type=int, default=0, required=False, help="Limit export to N documents")
+    parser.add_argument("--transformation", type=Path, required=False, help="Zyp transformation file")
+    parser.add_argument("-o", "--out", required=False)
 
 
 def translate_parser(subargs):
@@ -37,6 +46,8 @@ def export_parser(subargs):
     parser.add_argument("--host", default="localhost", help="MongoDB host")
     parser.add_argument("--port", default=27017, help="MongoDB port")
     parser.add_argument("--database", required=True, help="MongoDB database")
+    parser.add_argument("--limit", type=int, default=0, required=False, help="Limit export to N documents")
+    parser.add_argument("--transformation", type=Path, required=False, help="Zyp transformation file")
 
 
 def get_args():
@@ -61,32 +72,45 @@ def extract_to_file(args):
     """
 
     schema = extract(args)
-    rich.print(f"\nWriting resulting schema to {args.out}...")
-    with open(args.out, "w") as out:
-        json.dump(schema, out, indent=4)
-    rich.print("[green bold]Done![/green bold]")
+
+    out_label = args.out or "stdout"
+    rich.print(f"Writing resulting schema to {out_label}")
+    fp: t.TextIO
+    if args.out:
+        fp = open(args.out, "w")
+    else:
+        fp = sys.stdout
+    json.dump(schema, fp=fp, indent=4)
+    fp.flush()
 
 
 def translate_from_file(args):
     """
     Read in a JSON file and extract the schema from it.
     """
-
-    with open(args.infile) as f:
-        schema = json.load(f)
-        translate(schema)
+    fp: t.TextIO
+    if args.infile:
+        fp = open(args.infile, "r")
+    else:
+        fp = sys.stdin
+    schema = json.load(fp)
+    translate(schema)
 
 
 def export_to_stdout(args):
-    export(args)
+    sys.stdout.buffer.write(export(args).read())
 
 
 def main():
-    rich.print("\n[green bold]MongoDB[/green bold] -> [blue bold]CrateDB[/blue bold] Exporter :: Schema Extractor\n\n")
+    setup_logging()
     args = get_args()
+    headline_prefix = "[green bold]MongoDB[/green bold] -> [blue bold]CrateDB[/blue bold] Exporter"
     if args.command == "extract":
+        rich.print(f"{headline_prefix} -> Schema Extractor")
         extract_to_file(args)
     elif args.command == "translate":
+        rich.print(f"{headline_prefix} -> Schema Translator")
         translate_from_file(args)
     elif args.command == "export":
+        rich.print(f"{headline_prefix} -> Data Exporter")
         export_to_stdout(args)

cratedb_toolkit/io/mongodb/core.py

@@ -1,20 +1,27 @@
 import io
 import logging
+import sys
 import typing as t
+from collections import OrderedDict
 
 import pymongo
 import pymongo.database
-import rich
 from bson.raw_bson import RawBSONDocument
+from rich.console import Console
 from rich.syntax import Syntax
+from rich.table import Table
 
 from .export import collection_to_json
 from .extract import extract_schema_from_collection
+from .transform import TransformationManager
 from .translate import translate as translate_schema
 from .util import parse_input_numbers
 
 logger = logging.getLogger(__name__)
 
+console = Console(stderr=True)
+rich = console
+
 
 def gather_collections(database) -> t.List[str]:
     """
@@ -23,7 +30,7 @@ def gather_collections(database) -> t.List[str]:
 
     collections = database.list_collection_names()
 
-    tbl = rich.table.Table(show_header=True, header_style="bold blue")
+    tbl = Table(show_header=True, header_style="bold blue")
     tbl.add_column("Id", width=3)
     tbl.add_column("Collection Name")
     tbl.add_column("Estimated Size")
@@ -33,9 +40,10 @@ def gather_collections(database) -> t.List[str]:
 
     rich.print(tbl)
 
-    rich.print("\nCollections to exclude: (eg: '0 1 2', '0, 1, 2', '0-2')")
+    rich.print("\nCollections to exclude: (eg: '0 1 2', '0, 1, 2', '0-2'). Leave empty to use all collections.")
 
-    collections_to_ignore = parse_input_numbers(input("> "))
+    sys.stderr.write("> ")
+    collections_to_ignore = parse_input_numbers(input())
     filtered_collections = []
     for i, c in enumerate(collections):
         if i not in collections_to_ignore:
@@ -80,15 +88,19 @@ def extract(args) -> t.Dict[str, t.Any]:
     else:
         rich.print("\nDo a [red bold]full[/red bold] collection scan?")
         rich.print("A full scan will iterate over all documents in the collection, a partial only one document. (Y/n)")
-        full = input(">  ").strip().lower()
+        sys.stderr.write("> ")
+        full = input().strip().lower()
 
         partial = full != "y"
 
-        rich.print(f"\nExecuting a [red bold]{'partial' if partial else 'full'}[/red bold] scan...")
+        rich.print(f"\nExecuting a [red bold]{'partial' if partial else 'full'}[/red bold] scan")
 
-    schemas = {}
-    for collection in filtered_collections:
-        schemas[collection] = extract_schema_from_collection(db[collection], partial)
+    tm = TransformationManager(path=args.transformation)
+    schemas = OrderedDict()
+    for collection_name in filtered_collections:
+        collection_schema = extract_schema_from_collection(db[collection_name], partial, limit=args.limit)
+        tm.apply_type_overrides(db.name, collection_name, collection_schema)
+        schemas[collection_name] = collection_schema
     return schemas
 
 
@@ -103,7 +115,6 @@ def translate(schemas, schemaname: str = None) -> t.Dict[str, str]:
         syntax = Syntax(query, "sql")
         rich.print(f"Collection [blue bold]'{collection}'[/blue bold]:")
         rich.print(syntax)
-        rich.print()
     return result
 
 
@@ -113,8 +124,9 @@ def export(args) -> t.IO[bytes]:
 
     TODO: Run on multiple collections, like `extract`.
     """
+    tm = TransformationManager(path=args.transformation)
     buffer = io.BytesIO()
     client, db = get_mongodb_client_database(args, document_class=RawBSONDocument)
-    collection_to_json(db[args.collection], file=buffer)
+    collection_to_json(db[args.collection], fp=buffer, tm=tm, limit=args.limit)
     buffer.seek(0)
     return buffer