feat(staging): use GCS/new entities for API

Makefile

@@ -70,7 +70,7 @@ run-api-server:
 run-api-server-test:
 	test -f $(HOME)/.config/gcloud/application_default_credentials.json || (echo "GCP Application Default Credentials not set, try 'gcloud auth login --update-adc'"; exit 1)
 	cd gcp/api && docker build -f Dockerfile.esp -t osv/esp:latest .
-	cd gcp/api && $(install-cmd) && GOOGLE_CLOUD_PROJECT=oss-vdb-test $(run-cmd) python test_server.py $(HOME)/.config/gcloud/application_default_credentials.json $(ARGS)
+	cd gcp/api && $(install-cmd) && GOOGLE_CLOUD_PROJECT=oss-vdb-test OSV_VULNERABILITIES_BUCKET=osv-test-vulnerabilities $(run-cmd) python test_server.py $(HOME)/.config/gcloud/application_default_credentials.json $(ARGS)
 
 # TODO: API integration tests.
 all-tests: lib-tests worker-tests importer-tests alias-tests website-tests vulnfeed-tests

deployment/clouddeploy/osv-api/run-staging.yaml

@@ -0,0 +1,32 @@
+apiVersion: serving.knative.dev/v1
+kind: Service
+metadata:
+  name: osv-grpc-backend # from-param: ${serviceName}
+spec:
+  template:
+    metadata:
+      annotations:
+        autoscaling.knative.dev/maxScale: '300'
+    spec:
+      containers:
+      - image: osv-server
+        env:
+        - name: OSV_VULNERABILITIES_BUCKET
+          value: osv-test-vulnerabilities
+        resources:
+          limits:
+            cpu: 2
+            memory: 8Gi
+        startupProbe:
+          grpc:
+            service: osv.v1.OSV
+          initialDelaySeconds: 5
+          timeoutSeconds: 5
+        livenessProbe:
+          grpc:
+            service: osv.v1.OSV
+          timeoutSeconds: 5
+          failureThreshold: 3
+          periodSeconds: 10
+      timeoutSeconds: 60
+      containerConcurrency: 5 # from-param: ${containerConcurrency}

deployment/clouddeploy/osv-api/skaffold.yaml

@@ -3,17 +3,16 @@ kind: Config
 metadata:
   name: osv-api
 
-# These profiles actually unnecessary since there is currently no difference between staging and production manifests.
 profiles:
 - name: staging
   manifests:
     rawYaml:
-    - run.yaml
+    - run-staging.yaml
 
 - name: prod
   manifests:
     rawYaml:
-    - run.yaml
+    - run-prod.yaml
 
 deploy:
   cloudrun: {}

gcp/api/cursor.py

@@ -14,6 +14,7 @@
 """OSV API server cursor implementation."""
 
 import base64
+import dataclasses
 from enum import Enum
 from typing import Self
 import typing
@@ -25,6 +26,16 @@
 # Use ':' as the separator as it doesn't appear in urlsafe_b64encode
 # (which is what is used for both _FIRST_PAGE_TOKEN, and ndb.Cursor.urlsafe())
 _METADATA_SEPARATOR = ':'
+CURSOR_LAST_ID = 'last_id'
+
+
+@dataclasses.dataclass
+class QueryCursorMetadata:
+  """Metadata for the query cursor."""
+  # The last vulnerability ID seen in a query.
+  # This is used to avoid reprocessing the same vulnerability across pages
+  # when multiple AffectedVersions entities exist for it.
+  last_id: str | None = None
 
 
 class _QueryCursorState(Enum):
@@ -58,11 +69,13 @@ class QueryCursor:
     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
+  def __init__(self):
+    self._ndb_cursor: ndb.Cursor | None = None
+    self._cursor_state: _QueryCursorState = _QueryCursorState.ENDED
+    self.metadata: QueryCursorMetadata = QueryCursorMetadata()
+    # The first query is numbered 1. This is because the query counter is
+    # incremented **before** the query and the query number being used.
+    self.query_number: int = 1
 
   @classmethod
   def from_page_token(cls, page_token: str | None) -> Self:
@@ -74,13 +87,24 @@ def from_page_token(cls, page_token: str | None) -> Self:
       qc._ndb_cursor = None
       return qc
 
-    split_values = page_token.split(_METADATA_SEPARATOR, 1)
-    if len(split_values) == 2:
+    split_values = page_token.split(_METADATA_SEPARATOR)
+    metadata_dict = {}
+    if len(split_values) >= 2:
       page_token = split_values[1]
       try:
         qc.query_number = int(split_values[0])
       except ValueError as e:
         raise ValueError('Invalid page token.') from e
+      for enc_meta in split_values[2:]:
+        meta = base64.urlsafe_b64decode(enc_meta).decode(errors='ignore')
+        k, _, v = meta.partition('=')
+        metadata_dict[k] = v
+
+    try:
+      qc.metadata = QueryCursorMetadata(**metadata_dict)
+    except TypeError as e:
+      # Unexpected metadata field.
+      raise ValueError('Invalid page token.') from e
 
     if not page_token or page_token == _FIRST_PAGE_TOKEN:
       qc._cursor_state = _QueryCursorState.STARTED
@@ -144,4 +168,9 @@ def url_safe_encode(self) -> str | None:
         # a token in the response
         return None
 
-    return str(self.query_number) + _METADATA_SEPARATOR + cursor_part
+    meta_dict = dataclasses.asdict(self.metadata)
+    meta_items = (f'{k}={v}' for k, v in meta_dict.items() if v is not None)
+    meta = (
+        base64.urlsafe_b64encode(item.encode()).decode() for item in meta_items)
+    return _METADATA_SEPARATOR.join(
+        [str(self.query_number), cursor_part, *meta])

gcp/api/server.py

@@ -29,6 +29,7 @@
 
 from collections import defaultdict
 
+from google.cloud import exceptions
 from google.cloud import ndb
 from google.api_core.exceptions import InvalidArgument
 import google.cloud.ndb.exceptions as ndb_exceptions
@@ -47,7 +48,8 @@
 import osv_service_v1_pb2
 import osv_service_v1_pb2_grpc
 
-from cursor import QueryCursor
+from cursor import QueryCursor, QueryCursorMetadata
+from server_new import query_package
 
 import googlecloudprofiler
 
@@ -165,6 +167,26 @@ class OSVServicer(osv_service_v1_pb2_grpc.OSVServicer,
   @ndb.synctasklet
   def GetVulnById(self, request, context: grpc.ServicerContext):
     """Return a `Vulnerability` object for a given OSV ID."""
+
+    if get_gcp_project() in ('oss-vdb-test', 'test-osv'):
+      # Get vuln from GCS
+      try:
+        return osv.gcs.get_by_id(request.id)
+      except exceptions.NotFound:
+        # Check for aliases
+        alias_group = yield osv.AliasGroup.query(
+            osv.AliasGroup.bug_ids == request.id).get_async()
+        if alias_group:
+          alias_string = ' '.join([
+              f'{alias}' for alias in alias_group.bug_ids if alias != request.id
+          ])
+          context.abort(
+              grpc.StatusCode.NOT_FOUND,
+              f'Bug not found, but the following aliases were: {alias_string}')
+          return None
+        context.abort(grpc.StatusCode.NOT_FOUND, 'Bug not found.')
+        return None
+
     bug = yield osv.Bug.query(osv.Bug.db_id == request.id).get_async()
 
     if not bug:
@@ -535,12 +557,18 @@ def cursor_at_current(self) -> ndb.Cursor | None:
 
     return None
 
-  def save_cursor_at_page_break(self, it: ndb.QueryIterator):
+  def save_cursor_at_page_break(self,
+                                it: ndb.QueryIterator,
+                                meta: QueryCursorMetadata | None = None):
     """
     Saves the cursor at the current page break position
     """
     self.output_cursor.update_from_iterator(it)
     self.output_cursor.query_number = self.query_counter
+    if meta:
+      self.output_cursor.metadata = meta
+    else:
+      self.output_cursor.metadata = QueryCursorMetadata()
 
 
 def should_skip_bucket(path: str) -> bool:
@@ -843,6 +871,10 @@ def to_response(b: osv.Bug):
       return None
 
     bugs = yield query_by_commit(context, commit_bytes, to_response=to_response)
+  elif package_name and get_gcp_project() in ('oss-vdb-test', 'test-osv'):
+    # New Database table & GCS querying
+    bugs = yield query_package(context, package_name, ecosystem, version,
+                               include_details)
   # Version query needs to include a package.
   elif package_name and version:
     bugs = yield query_by_version(
@@ -874,7 +906,7 @@ def to_response(b: osv.Bug):
   # Wait on all the bug futures
   bugs = yield bugs
 
-  return list(bugs), next_page_token_str
+  return [b for b in bugs if b is not None], next_page_token_str
 
 
 def bug_to_response(bug: osv.Bug, include_details=True) -> ndb.Future: