fix _count aggregation queries (#74)

m-Bilal · web-flow · commit 9bb3b77d3d09 · 2025-04-18T19:54:22.000+05:30
* set track_total_hits to true in aggregation queries

* update tests
diff --git a/connector/query.go b/connector/query.go
@@ -190,13 +190,17 @@ func prepareElasticsearchQuery(ctx context.Context, request *schema.QueryRequest
 		}
 		if len(aggs) != 0 {
 			query["aggs"] = aggs
-
-			// set query size to 0 if aggregation is present
-			// this is because, by default, an aggregation query returns both the aggregation result and the hit documents
-			// we only want the aggregation result
-			// more reading: https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations.html#return-only-agg-results
-			query["size"] = 0
 		}
+
+		// set query size to 0 if aggregation is present
+		// this is because, by default, an aggregation query returns both the aggregation result and the hit documents
+		// we only want the aggregation result
+		// more reading: https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations.html#return-only-agg-results
+		query["size"] = 0
+
+		// set track_total_hits to true if aggregation is present
+		// this is because, by default, a count query will only return 10,000 hits
+		query["track_total_hits"] = true
 	}
 
 	span.AddEvent("prepare_filter_query")
diff --git a/testdata/unit_tests/query_tests/payments/cardinality_aggregation/want.json b/testdata/unit_tests/query_tests/payments/cardinality_aggregation/want.json
@@ -11,5 +11,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/count_aggregation/want.json b/testdata/unit_tests/query_tests/payments/count_aggregation/want.json
@@ -4,5 +4,6 @@
       "*"
     ]
   },
-  "size": 10000
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/count_distinct_aggregation/want.json b/testdata/unit_tests/query_tests/payments/count_distinct_aggregation/want.json
@@ -11,5 +11,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/float_aggregations/want.json b/testdata/unit_tests/query_tests/payments/float_aggregations/want.json
@@ -53,5 +53,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/float_aggregations_with_range/want.json b/testdata/unit_tests/query_tests/payments/float_aggregations_with_range/want.json
@@ -61,5 +61,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/multiple_aggregations/want.json b/testdata/unit_tests/query_tests/payments/multiple_aggregations/want.json
@@ -60,5 +60,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/multiple_aggregations_with_where_clause/want.json b/testdata/unit_tests/query_tests/payments/multiple_aggregations_with_where_clause/want.json
@@ -65,5 +65,6 @@
       "customer_id": "CUST"
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/nested_aggregations/want.json b/testdata/unit_tests/query_tests/payments/nested_aggregations/want.json
@@ -31,5 +31,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }
diff --git a/testdata/unit_tests/query_tests/payments/subfield_cardinality_aggregation/want.json b/testdata/unit_tests/query_tests/payments/subfield_cardinality_aggregation/want.json
@@ -11,5 +11,6 @@
       }
     }
   },
-  "size": 0
+  "size": 0,
+  "track_total_hits": true
 }

Original file line number	Diff line number	Diff line change
`@@ -11,5 +11,6 @@`
`11`	`11`	`}`
`12`	`12`	`}`
`13`	`13`	`},`
`14`		`- "size": 0`
	`14`	`+ "size": 0,`
	`15`	`+ "track_total_hits": true`
`15`	`16`	`}`
Original file line number	Diff line number	Diff line change
`@@ -4,5 +4,6 @@`
`4`	`4`	`"*"`
`5`	`5`	`]`
`6`	`6`	`},`
`7`		`- "size": 10000`
	`7`	`+ "size": 0,`
	`8`	`+ "track_total_hits": true`
`8`	`9`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,5 +53,6 @@`
`53`	`53`	`}`
`54`	`54`	`}`
`55`	`55`	`},`
`56`		`- "size": 0`
	`56`	`+ "size": 0,`
	`57`	`+ "track_total_hits": true`
`57`	`58`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,5 +61,6 @@`
`61`	`61`	`}`
`62`	`62`	`}`
`63`	`63`	`},`
`64`		`- "size": 0`
	`64`	`+ "size": 0,`
	`65`	`+ "track_total_hits": true`
`65`	`66`	`}`
Original file line number	Diff line number	Diff line change
`@@ -60,5 +60,6 @@`
`60`	`60`	`}`
`61`	`61`	`}`
`62`	`62`	`},`
`63`		`- "size": 0`
	`63`	`+ "size": 0,`
	`64`	`+ "track_total_hits": true`
`64`	`65`	`}`
Original file line number	Diff line number	Diff line change
`@@ -65,5 +65,6 @@`
`65`	`65`	`"customer_id": "CUST"`
`66`	`66`	`}`
`67`	`67`	`},`
`68`		`- "size": 0`
	`68`	`+ "size": 0,`
	`69`	`+ "track_total_hits": true`
`69`	`70`	`}`
Original file line number	Diff line number	Diff line change
`@@ -31,5 +31,6 @@`
`31`	`31`	`}`
`32`	`32`	`}`
`33`	`33`	`},`
`34`		`- "size": 0`
	`34`	`+ "size": 0,`
	`35`	`+ "track_total_hits": true`
`35`	`36`	`}`