work around cb response 6.1 issue where total_results may not equal the number of results for a query

jgarman · jgarman · commit ec4cc7acc103 · 2017-05-30T23:57:39.000-04:00
diff --git a/examples/response/walk_children.py b/examples/response/walk_children.py
@@ -35,7 +35,7 @@ def main():
             print("Encountered unknown error retrieving process: {0:s}".format(str(e)))
             return 1
     elif args.query:
-        procs = c.select(Process).where(args.query)
+        procs = c.select(Process).where(args.query).group_by("id")
     else:
         print("Requires either a --process or --query argument")
         parser.print_usage()
diff --git a/src/cbapi/response/query.py b/src/cbapi/response/query.py
@@ -5,6 +5,10 @@
 from ..errors import ApiError
 from ..utils import convert_query_params
 
+import logging
+
+log = logging.getLogger(__name__)
+
 
 class Query(PaginatedQuery):
     """Represents a prepared query to the Carbon Black Enterprise Response server.
@@ -183,7 +187,9 @@ def _search(self, start=0, rows=0):
             self._total_results = result.get('total_results')
             self._count_valid = True
 
-            for item in result.get('results'):
+            results = result.get('results', [])
+
+            for item in results:
                 yield item
                 current += 1
                 numrows += 1
@@ -195,3 +201,9 @@ def _search(self, start=0, rows=0):
 
             if current >= self._total_results:
                 break
+            if not results:
+                log.debug("server reported total_results overestimated the number of results for this query by {0}"
+                          .format(self._total_results - current))
+                log.debug("resetting total_results for this query to {0}".format(current))
+                self._total_results = current
+                break
