feat(ingestion/mongodb) re-order aggregation logic

metadata-ingestion/src/datahub/ingestion/source/mongodb.py

@@ -218,26 +218,25 @@
     """
 
     aggregations: List[Dict] = []
+
+    if sample_size:
+        if use_random_sampling:
+            aggregations.append({"$sample": {"size": sample_size}})
+        else:
+            aggregations.append({"$limit": sample_size})
+
     if should_add_document_size_filter:
         doc_size_field = "temporary_doc_size_field"
         # create a temporary field to store the size of the document. filter on it and then remove it.
-        aggregations = [
-            {"$addFields": {doc_size_field: {"$bsonSize": "$$ROOT"}}},
-            {"$match": {doc_size_field: {"$lt": max_document_size}}},
-            {"$project": {doc_size_field: 0}},
-        ]
-    if use_random_sampling:
-        # get sample documents in collection
-        if sample_size:
-            aggregations.append({"$sample": {"size": sample_size}})
-        documents = collection.aggregate(
-            aggregations,
-            allowDiskUse=True,
+        aggregations.extend(
+            [
+                {"$addFields": {doc_size_field: {"$bsonSize": "$$ROOT"}}},
+                {"$match": {doc_size_field: {"$lt": max_document_size}}},
+                {"$project": {doc_size_field: 0}},
+            ]
         )
-    else:
-        if sample_size:
-            aggregations.append({"$limit": sample_size})
-        documents = collection.aggregate(aggregations, allowDiskUse=True)
+
+    documents = collection.aggregate(aggregations, allowDiskUse=True)
 
     return construct_schema(list(documents), delimiter)