[core][fix] Apply limit to aggregate result not query (#2018)

someengineering · Apr 8, 2024 · a85ff90 · a85ff90
1 parent 7d8accf
commit a85ff90
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/fixcore/fixcore/cli/command.py b/fixcore/fixcore/cli/command.py
@@ -1489,7 +1489,7 @@ async def prepare() -> Tuple[CLISourceContext, AsyncIterator[Json]]:
                     query_model, changes, before, after, with_count=count, timeout=timeout
                 )
             elif query.aggregate:
-                context = await db.search_aggregation(query_model)
+                context = await db.search_aggregation(query_model, with_count=count, timeout=timeout)
             elif with_edges:
                 context = await db.search_graph_gen(query_model, with_count=count, timeout=timeout)
             else:

diff --git a/fixcore/fixcore/db/graphdb.py b/fixcore/fixcore/db/graphdb.py
@@ -203,7 +203,9 @@ async def search_graph(self, query: QueryModel) -> MultiDiGraph:
         pass
 
     @abstractmethod
-    async def search_aggregation(self, query: QueryModel) -> AsyncCursorContext:
+    async def search_aggregation(
+        self, query: QueryModel, with_count: bool = False, timeout: Optional[timedelta] = None
+    ) -> AsyncCursorContext:
         pass
 
     @abstractmethod
@@ -781,10 +783,18 @@ async def search_graph(self, query: QueryModel) -> MultiDiGraph:
                     graph.add_node(item["id"], **item)
             return graph
 
-    async def search_aggregation(self, query: QueryModel) -> AsyncCursorContext:
+    async def search_aggregation(
+        self, query: QueryModel, with_count: bool = False, timeout: Optional[timedelta] = None
+    ) -> AsyncCursorContext:
         q_string, bind = await self.to_query(query)
         assert query.query.aggregate is not None, "Given query has no aggregation section"
-        return await self.db.aql_cursor(query=q_string, bind_vars=bind)
+        return await self.db.aql_cursor(
+            query=q_string,
+            bind_vars=bind,
+            count=with_count,
+            full_count=with_count,
+            ttl=cast(Number, int(timeout.total_seconds())) if timeout else None,
+        )
 
     async def explain(self, query: QueryModel, with_edges: bool = False) -> EstimatedSearchCost:
         return await arango_query.query_cost(self, query, with_edges)
@@ -1779,8 +1789,10 @@ async def search_graph_gen(
     ) -> AsyncCursorContext:
         return await self.real.search_graph_gen(query, with_count, timeout)
 
-    async def search_aggregation(self, query: QueryModel) -> AsyncCursorContext:
-        return await self.real.search_aggregation(query)
+    async def search_aggregation(
+        self, query: QueryModel, with_count: bool = False, timeout: Optional[timedelta] = None
+    ) -> AsyncCursorContext:
+        return await self.real.search_aggregation(query, with_count, timeout)
 
     async def search_graph(self, query: QueryModel) -> MultiDiGraph:
         return await self.real.search_graph(query)