MongoDB: Filter server collection using MongoDB query expression

CHANGES.md

@@ -5,6 +5,7 @@
 - MongoDB: Unlock processing multiple collections, either from server database,
   or from filesystem directory
 - MongoDB: Unlock processing JSON files from HTTP resource, using `https+bson://`
+- MongoDB: Optionally filter server collection using MongoDB query expression
 
 ## 2024/09/10 v0.0.22
 - MongoDB: Rename columns with leading underscores to use double leading underscores

cratedb_toolkit/io/mongodb/adapter.py

@@ -1,5 +1,6 @@
 import glob
 import itertools
+import json
 import logging
 import typing as t
 from abc import abstractmethod
@@ -28,7 +29,7 @@ class MongoDBAdapterBase:
     database_name: str
     collection_name: str
 
-    _custom_query_parameters = ["batch-size", "limit", "offset"]
+    _custom_query_parameters = ["batch-size", "filter", "limit", "offset"]
 
     @classmethod
     def from_url(cls, url: t.Union[str, boltons.urlutils.URL, yarl.URL]):
@@ -54,6 +55,10 @@ def __attrs_post_init__(self):
     def batch_size(self) -> int:
         return int(self.address.uri.query_params.get("batch-size", 500))
 
+    @cached_property
+    def filter(self) -> t.Union[str, None]:
+        return json.loads(self.address.uri.query_params.get("filter", "null"))
+
     @cached_property
     def limit(self) -> int:
         return int(self.address.uri.query_params.get("limit", 0))
@@ -100,6 +105,8 @@ def record_count(self, filter_=None) -> int:
     def query(self):
         if not self._path.exists():
             raise FileNotFoundError(f"Resource not found: {self._path}")
+        if self.filter:
+            raise NotImplementedError("Using MongoDB filter expressions is not supported by Polars' NDJSON reader")
         if self.offset:
             raise NotImplementedError("Using offsets is not supported by Polars' NDJSON reader")
         if self._path.suffix in [".json", ".jsonl", ".ndjson"]:
@@ -129,6 +136,8 @@ def record_count(self, filter_=None) -> int:
         return -1
 
     def query(self):
+        if self.filter:
+            raise NotImplementedError("Using MongoDB filter expressions is not supported by Polars' NDJSON reader")
         if self.offset:
             raise NotImplementedError("Using offsets is not supported by Polars' NDJSON reader")
         if self._url.path.endswith(".json") or self._url.path.endswith(".jsonl") or self._url.path.endswith(".ndjson"):
@@ -165,7 +174,12 @@ def record_count(self, filter_=None) -> int:
         return self._mongodb_collection.count_documents(filter=filter_)
 
     def query(self):
-        data = self._mongodb_collection.find().batch_size(self.batch_size).skip(self.offset).limit(self.limit)
+        data = (
+            self._mongodb_collection.find(filter=self.filter)
+            .batch_size(self.batch_size)
+            .skip(self.offset)
+            .limit(self.limit)
+        )
         return batches(data, self.batch_size)
 
 

doc/io/mongodb/loader.md

@@ -91,16 +91,24 @@ The default batch size is 500. You can adjust the value by appending the HTTP
 URL query parameter `batch-size` to the source URL, like
 `mongodb+srv://managed.mongodb.net/ticker/stocks?batch-size=5000`.
 
-### Offset
-Use the HTTP URL query parameter `offset` on the source URL, like
-`&offset=42`, in order to start processing at this record from the
-beginning.
+### Filter
+Use the HTTP URL query parameter `filter` on the source URL, like
+`&filter={"exchange": {"$eq": "NASDAQ"}}`, in order to provide a MongoDB
+query filter as a JSON string.
+It works in the same way like `mongoexport`'s `--query` option.
+On more complex query expressions, make sure to properly encode the right
+value using URL/Percent Encoding.
 
 ### Limit
 Use the HTTP URL query parameter `limit` on the source URL, like
 `&limit=100`, in order to limit processing to a total number of
 records.
 
+### Offset
+Use the HTTP URL query parameter `offset` on the source URL, like
+`&offset=42`, in order to start processing at this record from the
+beginning.
+
 ## Zyp Transformations
 You can use [Zyp Transformations] to change the shape of the data while being
 transferred. In order to add it to the pipeline, use the `--transformation`