From 1ae4f861f31fc636d4e08806f2e57da9f176b103 Mon Sep 17 00:00:00 2001
From: Michael Kedar <michaelkedar@google.com>
Date: Fri, 16 Jan 2026 14:45:19 +1100
Subject: [PATCH] feat(api): filter query by coarse version

---
 gcp/api/server.py        | 80 ++++++++--------------------------------
 gcp/datastore/index.yaml |  8 ++++
 2 files changed, 23 insertions(+), 65 deletions(-)

diff --git a/gcp/api/server.py b/gcp/api/server.py
index 47d1fbbf695..1f25d0c9f59 100644
--- a/gcp/api/server.py
+++ b/gcp/api/server.py
@@ -25,7 +25,6 @@
 import threading
 import time
 import concurrent.futures
-from typing import Callable
 
 from collections import defaultdict
 
@@ -100,11 +99,6 @@
 
 _TEST_INSTANCE = 'oss-vdb-test'
 
-# ----
-# Type Aliases:
-
-ToResponseCallable = Callable[[osv.Bug], ndb.Future]
-
 # ----
 
 _ndb_client = ndb.Client()
@@ -392,9 +386,9 @@ def Check(self, request, context: grpc.ServicerContext):
     del request  # Unused.
     del context  # Unused.
 
-    # Read up to a single Bug entity from the DB. This should not cause an
-    # exception or time out.
-    osv.Bug.query().fetch(1)
+    # Read up to a single Vulnerability entity from the DB.
+    # This should not cause an exception or time out.
+    osv.Vulnerability.query().fetch(1)
     return health_pb2.HealthCheckResponse(
         status=health_pb2.HealthCheckResponse.ServingStatus.SERVING)
 
@@ -965,12 +959,6 @@ def query_package(context: QueryContext,
   if not package_name:
     return []
 
-  if (package_name.startswith('linux') and ecosystem.startswith('Ubuntu') and
-      version):
-    result = yield query_ubuntu_linux(context, package_name, ecosystem, version,
-                                      include_details)
-    return result
-
   # Ideally, we'd check both unnormalized and normalized named at once, if there
   # is no provided ecosystem (even though that is not explicitly supported), but
   # Datastore cannot give cursors for 'OR' queries, so just only normalize if
@@ -981,6 +969,18 @@ def query_package(context: QueryContext,
   query = osv.AffectedVersions.query(osv.AffectedVersions.name == package_name)
   if ecosystem:
     query = query.filter(osv.AffectedVersions.ecosystem == ecosystem)
+    if version:
+      helper = ecosystems.get(ecosystem)
+      if helper is not None:
+        try:
+          coarse = helper.coarse_version(version)
+          query = query.filter(osv.AffectedVersions.coarse_min <= coarse)
+          query = query.filter(osv.AffectedVersions.coarse_max >= coarse)
+        except (ValueError, NotImplementedError):
+          # We want to avoid logging user requests.
+          pass
+          # TODO(michaelkedar): I think a ValueError would mean the version
+          # would not match any affected versions, so we may be able to return.
   query = query.order(osv.AffectedVersions.vuln_id)
 
   bugs = []
@@ -1147,56 +1147,6 @@ def cleanup(_: ndb.Future):
   return future
 
 
-@ndb.tasklet
-def query_ubuntu_linux(context: QueryContext,
-                       package_name: str,
-                       ecosystem: str,
-                       version: str,
-                       include_details: bool = True) -> list[ndb.Future]:
-  """Workaround query for linux kernel vulns in Ubuntu, because there's heaps of
-  them and it takes a while to check all the version ranges.
-  
-  The logic here is copied from the original implementation of server.py,
-  using the Bug entity to perform matching on only the version strings, and not
-  doing range-based matching.
-  """
-
-  query = osv.Bug.query(
-      osv.Bug.status == osv.BugStatus.PROCESSED,
-      osv.Bug.project == package_name,
-      # pylint: disable=singleton-comparison
-      osv.Bug.public == True,
-      osv.Bug.ecosystem == ecosystem,
-      osv.Bug.affected_fuzzy == version)
-
-  bugs = []
-  it: ndb.QueryIterator = query.iter(start_cursor=context.cursor_at_current())
-  while (yield it.has_next_async()):
-    if context.should_break_page(len(bugs)):
-      context.save_cursor_at_page_break(it)
-      break
-
-    bug: osv.Bug = it.next()
-    for affected in bug.affected_packages:
-      eco = affected.package.ecosystem
-      if ecosystem not in (eco, ecosystems.normalize(eco),
-                           ecosystems.remove_variants(eco)):
-        continue
-      if package_name != affected.package.name:
-        continue
-      if version not in affected.versions:
-        continue
-      # Found a match: retrieve the proto from the bucket / Vulnerability entity
-      if include_details:
-        bugs.append(get_vuln_async(bug.db_id))
-      else:
-        bugs.append(get_minimal_async(bug.db_id))
-      context.total_responses.add(1)
-      break
-
-  return bugs
-
-
 def serve(port: int, local: bool):
   """Configures and runs the OSV API server."""
   server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=5))
diff --git a/gcp/datastore/index.yaml b/gcp/datastore/index.yaml
index a1504ca13ca..9211b025c94 100644
--- a/gcp/datastore/index.yaml
+++ b/gcp/datastore/index.yaml
@@ -123,6 +123,14 @@ indexes:
     properties:
       - name: name
       - name: vuln_id
+  
+  - kind: AffectedVersions
+    properties:
+      - name: ecosystem
+      - name: name
+      - name: vuln_id
+      - name: coarse_min
+      - name: coarse_max
 
   # Indexes for website /list search
   - kind: ListedVulnerability