Merge pull request #1417 from p-l-/enh-documentdb-top-net

DB/DocumentDB: support top net values
ivre · Aug 8, 2022 · 0d4aa2b · 0d4aa2b
2 parents 26d545b + bbf729e
commit 0d4aa2b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 10 deletions.
diff --git a/ivre/db/document.py b/ivre/db/document.py
@@ -34,17 +34,18 @@
 
 
 class DocumentDBNmap(MongoDBNmap):
-    pass
+    is_documentdb = True
 
 
 class DocumentDBView(MongoDBView):
+    is_documentdb = True
     # DocumentDB has no support for text indexes
     indexes = MongoDBActive.indexes
     schema_migrations_indexes = MongoDBActive.schema_migrations_indexes
 
 
 class DocumentDBPassive(MongoDBPassive):
-    pass
+    is_documentdb = True
 
 
 class DocumentDBAgent(MongoDBAgent):

diff --git a/ivre/db/mongo.py b/ivre/db/mongo.py
@@ -82,6 +82,7 @@ def log_pipeline(pipeline):
 
 class MongoDB(DB):
 
+    is_documentdb = False  # set to True for AWS DocumentDB sub-classes
     indexes: List[List[Tuple[List[IndexKey], Dict[str, Any]]]] = []
     schema_migrations_indexes: List[
         Dict[int, Dict[str, List[Tuple[List[IndexKey], Dict[str, Any]]]]]
@@ -3322,10 +3323,18 @@ def outputproc(x):
             field = "addr"
             # This should not overflow thanks to .searchipv4() filter
             addr = {"$add": ["$addr_1", 0x7FFF000100000000]}
-            specialproj = {
-                "_id": 0,
-                "addr": {"$floor": {"$divide": [addr, 2 ** (32 - mask)]}},
-            }
+            specialproj = {"_id": 0}
+            if self.is_documentdb:
+                # AWS DocumentDB lacks $floor aggregation operator,
+                # just like MongoDB < 3.2
+                specialproj["addr"] = {
+                    "$subtract": [
+                        {"$divide": [addr, 2 ** (32 - mask)]},
+                        {"$mod": [{"$divide": [addr, 2 ** (32 - mask)]}, 1]},
+                    ]
+                }
+            else:
+                specialproj["addr"] = {"$floor": {"$divide": [addr, 2 ** (32 - mask)]}}
             flt = self.flt_and(flt, self.searchipv4())
 
             def outputproc(x):
@@ -5025,10 +5034,18 @@ def outputproc(x):
             field = "addr"
             # This should not overflow thanks to .searchipv4() filter
             addr = {"$add": ["$addr_1", 0x7FFF000100000000]}
-            specialproj = {
-                "_id": 0,
-                "addr": {"$floor": {"$divide": [addr, 2 ** (32 - mask)]}},
-            }
+            specialproj = {"_id": 0}
+            if self.is_documentdb:
+                # AWS DocumentDB lacks $floor aggregation operator,
+                # just like MongoDB < 3.2
+                specialproj["addr"] = {
+                    "$subtract": [
+                        {"$divide": [addr, 2 ** (32 - mask)]},
+                        {"$mod": [{"$divide": [addr, 2 ** (32 - mask)]}, 1]},
+                    ]
+                }
+            else:
+                specialproj["addr"] = {"$floor": {"$divide": [addr, 2 ** (32 - mask)]}}
             flt = self.flt_and(flt, self.searchipv4())
 
             def outputproc(x):