diff --git a/docker/importer/importer_test.py b/docker/importer/importer_test.py
index 868ff943c0f..e19be076562 100644
--- a/docker/importer/importer_test.py
+++ b/docker/importer/importer_test.py
@@ -1116,6 +1116,7 @@ def test_add_finding(self):
     """Test that creating an import finding works."""
     expected = osv.ImportFinding(
         bug_id='CVE-2024-1234',
+        source='cve-osv',
         findings=[
             osv.ImportFindings.INVALID_VERSION,
         ],
@@ -1124,9 +1125,8 @@ def test_add_finding(self):
     )
     expected.put()
 
-    for actual in osv.ImportFinding.query(
-        osv.ImportFinding.bug_id == expected.bug_id):
-      self.assertEqual(expected, actual)
+    actual = osv.ImportFinding.get_by_id(expected.bug_id)
+    self.assertEqual(expected, actual)
 
 
 if __name__ == '__main__':
diff --git a/osv/models.py b/osv/models.py
index e6ca98c95d9..6eb997eaa31 100644
--- a/osv/models.py
+++ b/osv/models.py
@@ -917,6 +917,7 @@ class ImportFindings(enum.IntEnum):
 class ImportFinding(ndb.Model):
   """Quality findings about an individual record."""
   bug_id: str = ndb.StringProperty()
+  source: str = ndb.StringProperty()
   findings: list[ImportFindings] = ndb.IntegerProperty(repeated=True)
   first_seen: datetime = ndb.DateTimeProperty()
   last_attempt: datetime = ndb.DateTimeProperty()