Do the Right Thing™ when a query has a LIMIT inline

Author: jmountifield

jupiterone/client.py

@@ -6,6 +6,7 @@
 from warnings import warn
 from typing import Dict, List
 
+import re
 import requests
 from retrying import retry
 
@@ -149,6 +150,15 @@ def _cursor_query(
             include_deleted (bool): Include recently deleted entities in query/search
         """
 
+        # If the query itself includes a LIMIT then we must parse that and check if we've reached
+        # or exceeded the required number of results.
+        limit_match = re.search(r"(?i)LIMIT\s+(?P<inline_limit>\d+)", query)
+
+        if limit_match:
+            result_limit = int(limit_match.group("inline_limit"))
+        else:
+            result_limit = False
+
         results: List = []
         while True:
             variables = {"query": query, "includeDeleted": include_deleted}
@@ -165,14 +175,24 @@ def _cursor_query(
 
             results.extend(data)
 
-            if (
+            if result_limit and len(results) >= result_limit:
+                # We can stop paginating if we've collected enough results based on the requested limit
+                break
+            elif (
                 "cursor" in response["data"]["queryV1"]
                 and response["data"]["queryV1"]["cursor"] is not None
             ):
+                # We got a cursor and haven't collected enough results
                 cursor = response["data"]["queryV1"]["cursor"]
             else:
+                # No cursor returned so we're done
                 break
 
+        # If we detected an inline LIMIT make sure we only return that many results
+        if result_limit:
+            return {"data": results[:result_limit]}
+
+        # Return everything
         return {"data": results}
 
     def _limit_and_skip_query(

setup.py

@@ -1,30 +1,28 @@
 # Copyright (c) 2020-2025 JupiterOne
 from setuptools import setup, find_packages
 
-install_reqs = [
-    'requests',
-    'retrying'
-]
+install_reqs = ["requests", "retrying"]
 
-setup(name='jupiterone',
-      version='1.0.0',
-      description='A Python client for the JupiterOne API',
-      license='MIT License',
-      author='JupiterOne',
-      author_email='[email protected]',
-      maintainer='JupiterOne',
-      url='https://github.com/JupiterOne/jupiterone-api-client-python',
-      install_requires=install_reqs,
-      classifiers=[
-          'Development Status :: 4 - Beta',
-          'Intended Audience :: Developers',
-          'Intended Audience :: Information Technology',
-          'Intended Audience :: System Administrators',
-          'License :: OSI Approved :: MIT License',
-          'Natural Language :: English',
-          'Operating System :: POSIX :: Linux',
-          'Programming Language :: Python',
-          'Topic :: Security',
-      ],
-      packages=find_packages()
-      )
+setup(
+    name="jupiterone",
+    version="1.0.1",
+    description="A Python client for the JupiterOne API",
+    license="MIT License",
+    author="JupiterOne",
+    author_email="[email protected]",
+    maintainer="JupiterOne",
+    url="https://github.com/JupiterOne/jupiterone-api-client-python",
+    install_requires=install_reqs,
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Information Technology",
+        "Intended Audience :: System Administrators",
+        "License :: OSI Approved :: MIT License",
+        "Natural Language :: English",
+        "Operating System :: POSIX :: Linux",
+        "Programming Language :: Python",
+        "Topic :: Security",
+    ],
+    packages=find_packages(),
+)