diff --git a/fixcore/fixcore/cli/command.py b/fixcore/fixcore/cli/command.py index 1df8e961f4..0f5a983add 100644 --- 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 index 4c1f1d1ca1..9fd104909d 100644 --- 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)