From 94d051510da3b01001c03ced1bb77cb445b0853b Mon Sep 17 00:00:00 2001 From: Rex P Date: Tue, 27 Aug 2024 15:13:40 +1000 Subject: [PATCH] Add cursor at current function --- gcp/api/cursor.py | 7 +++---- gcp/api/integration_tests.py | 6 +++--- gcp/api/server.py | 22 ++++++++++++++++------ 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/gcp/api/cursor.py b/gcp/api/cursor.py index 7c4810242d4..2a6e3c52d8e 100644 --- a/gcp/api/cursor.py +++ b/gcp/api/cursor.py @@ -53,13 +53,15 @@ class QueryCursor: Attributes: query_number: This cursor is specifically for the Nth ndb datastore - query in the current query request. + query in the current query request. (Starts from 1) ndb_cursor: Get the internal ndb_cursor. This could be None. ended: Whether this cursor is for a query that has finished returning data. """ _ndb_cursor: ndb.Cursor | None = None _cursor_state: _QueryCursorState = _QueryCursorState.ENDED + # The first query is numbered 1. This is because the query counter is + # incremented **before** the query and the query number being used. query_number: int = 1 @classmethod @@ -142,7 +144,4 @@ def url_safe_encode(self) -> str | None: # a token in the response return None - if self.query_number == 0: - return cursor_part - return str(self.query_number) + _METADATA_SEPARATOR + cursor_part diff --git a/gcp/api/integration_tests.py b/gcp/api/integration_tests.py index 465a848868c..997706caedf 100644 --- a/gcp/api/integration_tests.py +++ b/gcp/api/integration_tests.py @@ -775,14 +775,14 @@ def test_query_pagination_no_ecosystem(self): 'package': { 'name': 'django', }, - 'version': '5.0', + 'version': '5.0.1', }), timeout=_TIMEOUT) result = response.json() vulns_first = set(v['id'] for v in result['vulns']) self.assertIn('next_page_token', result) - self.assertTrue(str.startswith(result['next_page_token'], '1:')) + self.assertTrue(str.startswith(result['next_page_token'], '2:')) response = requests.post( _api() + _BASE_QUERY, @@ -790,7 +790,7 @@ def test_query_pagination_no_ecosystem(self): 'package': { 'name': 'django', }, - 'version': '5.0', + 'version': '5.0.1', 'page_token': result['next_page_token'], }), timeout=_TIMEOUT) diff --git a/gcp/api/server.py b/gcp/api/server.py index 5347c5a6087..471a6565bd0 100644 --- a/gcp/api/server.py +++ b/gcp/api/server.py @@ -458,6 +458,16 @@ def should_skip_query(self): return (self.query_counter < self.input_cursor.query_number or not self.output_cursor.ended) + def cursor_at_current(self) -> ndb.Cursor | None: + """ + Return the cursor if the stored cursor is for the current query. + """ + if self.input_cursor.query_number == self.query_counter: + return self.input_cursor.ndb_cursor + else: + return None + + def save_cursor_at_page_break(self, it: ndb.QueryIterator): """ Saves the cursor at the current page break position @@ -853,7 +863,7 @@ def query_by_commit(context: QueryContext, bug_ids = [] it: ndb.QueryIterator = query.iter( - keys_only=True, start_cursor=context.input_cursor.ndb_cursor) + keys_only=True, start_cursor=context.cursor_at_current()) while (yield it.has_next_async()): if context.should_break_page(len(bug_ids)): @@ -1025,7 +1035,7 @@ def _query_by_semver(context: QueryContext, query: ndb.Query, return [] it: ndb.QueryIterator = query.iter( - start_cursor=context.input_cursor.ndb_cursor) + start_cursor=context.cursor_at_current()) while (yield it.has_next_async()): if context.should_break_page(len(results)): @@ -1107,7 +1117,7 @@ def query_by_generic_helper(context: QueryContext, base_query: ndb.Query, return [] it: ndb.QueryIterator = query.iter( - start_cursor=context.input_cursor.ndb_cursor) + start_cursor=context.cursor_at_current()) while (yield it.has_next_async()): if context.should_break_page(len(results)): @@ -1235,7 +1245,7 @@ def _query_by_comparing_versions(context: QueryContext, query: ndb.Query, context: QueryContext for the current query. query: A partially completed ndb.Query object which only needs version filters to be added before query is performed. - ecosystem: Optional ecosystem of the package to query. + ecosystem: Required ecosystem of the package to query. version: The version str to query by. Returns: @@ -1248,7 +1258,7 @@ def _query_by_comparing_versions(context: QueryContext, query: ndb.Query, return [] it: ndb.QueryIterator = query.iter( - start_cursor=context.input_cursor.ndb_cursor) + start_cursor=context.cursor_at_current()) # Checks if the query specifies a release (e.g., "Debian:12") has_release = ':' in ecosystem @@ -1331,7 +1341,7 @@ def query_by_package(context: QueryContext, package_name: str | None, return [] it: ndb.QueryIterator = query.iter( - start_cursor=context.input_cursor.ndb_cursor) + start_cursor=context.cursor_at_current()) while (yield it.has_next_async()): if context.should_break_page(len(bugs)):