diff --git a/CHANGELOG.md b/CHANGELOG.md index 21fd05b6a..170b18860 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Silenced a false-positive `MaxListenersExceededWarning` logged on every request proxied through an external rewrite. [#1697](https://github.com/sourcebot-dev/sourcebot/pull/1697) +- Fixed duplicate repository metadata lookups within search result chunks. [#1684](https://github.com/sourcebot-dev/sourcebot/pull/1684) ## [5.1.15] - 2026-09-29 diff --git a/packages/web/src/features/search/zoektSearcher.test.ts b/packages/web/src/features/search/zoektSearcher.test.ts index 0bdb72818..e1aca6260 100644 --- a/packages/web/src/features/search/zoektSearcher.test.ts +++ b/packages/web/src/features/search/zoektSearcher.test.ts @@ -1,13 +1,16 @@ import type { PrismaClient } from '@sourcebot/db'; import type { SearchRequest as ZoektGrpcSearchRequest } from '@/proto/zoekt/webserver/v1/SearchRequest'; import { beforeEach, describe, expect, test, vi } from 'vitest'; +import { EventEmitter } from 'node:events'; const mocks = vi.hoisted(() => { const close = vi.fn(); const search = vi.fn(); + const streamSearch = vi.fn(); class WebserverService { Search = search; + StreamSearch = streamSearch; close = close; } @@ -15,6 +18,7 @@ const mocks = vi.hoisted(() => { close, loadSync: vi.fn(() => ({})), search, + streamSearch, WebserverService, }; }); @@ -60,10 +64,28 @@ vi.mock('@/lib/posthog', () => ({ captureEvent: vi.fn(), })); -import { zoektSearch } from './zoektSearcher'; +import { zoektSearch, zoektStreamSearch } from './zoektSearcher'; +import type { SearchResultFile, StreamedSearchResponse } from './types'; const searchRequest = {} as ZoektGrpcSearchRequest; +const createFile = (id: number | undefined, repository = 'github.com/org/repo') => ({ + repository_id: id, + repository, + file_name: Buffer.from('src/index.ts'), + chunk_matches: [], + branches: ['main'], + language: 'TypeScript', +}); + +const createRepo = (id: number, name = 'github.com/org/repo') => ({ + id, + name, + displayName: name, + webUrl: null, + external_codeHostType: 'github', +}); + describe('zoektSearch', () => { beforeEach(() => { vi.clearAllMocks(); @@ -97,11 +119,141 @@ describe('zoektSearch', () => { }); const prisma = { repo: { - findUnique: vi.fn().mockRejectedValue(new Error('database unavailable')), + findMany: vi.fn().mockRejectedValue(new Error('database unavailable')), }, } as unknown as PrismaClient; await expect(zoektSearch(searchRequest, prisma)).rejects.toThrow('database unavailable'); expect(mocks.close).toHaveBeenCalledOnce(); }); + + test.each([1, 2])('batches %i repositories into one lookup for 100 files', async (repoCount) => { + const files = Array.from({ length: 100 }, (_, index) => createFile(index % repoCount + 1)); + mocks.search.mockImplementation((_request, _metadata, callback) => { + callback(null, { files }); + }); + const ids = Array.from({ length: repoCount }, (_, index) => index + 1); + const findMany = vi.fn().mockResolvedValue(ids.toReversed().map(id => createRepo(id))); + const prisma = { repo: { findMany } } as unknown as PrismaClient; + + const response = await zoektSearch(searchRequest, prisma); + + expect(findMany).toHaveBeenCalledExactlyOnceWith({ where: { id: { in: ids } } }); + expect(response.files).toHaveLength(100); + expect(response.files.map(file => file.repositoryId)).toEqual(files.map(file => file.repository_id)); + expect(response.repositoryInfo.map(repo => repo.id)).toEqual(ids); + }); + + test('deduplicates lookups by name for legacy shards without repository IDs', async () => { + mocks.search.mockImplementation((_request, _metadata, callback) => { + callback(null, { files: Array.from({ length: 100 }, () => createFile(undefined)) }); + }); + const findFirst = vi.fn().mockResolvedValue(createRepo(1)); + const prisma = { repo: { findFirst } } as unknown as PrismaClient; + + const response = await zoektSearch(searchRequest, prisma); + + expect(findFirst).toHaveBeenCalledExactlyOnceWith({ where: { name: 'github.com/org/repo' } }); + expect(response.files).toHaveLength(100); + }); + + test('looks up a missing repository once and omits its files', async () => { + mocks.search.mockImplementation((_request, _metadata, callback) => { + callback(null, { files: Array.from({ length: 100 }, () => createFile(1)) }); + }); + const findMany = vi.fn().mockResolvedValue([]); + const prisma = { repo: { findMany } } as unknown as PrismaClient; + + const response = await zoektSearch(searchRequest, prisma); + + expect(findMany).toHaveBeenCalledOnce(); + expect(response.files).toEqual([]); + expect(response.repositoryInfo).toEqual([]); + }); + + test('deduplicates streaming chunks and reuses repository metadata across chunks', async () => { + const grpcStream = Object.assign(new EventEmitter(), { + pause: vi.fn(), + resume: vi.fn(), + cancel: vi.fn(), + }); + mocks.streamSearch.mockReturnValue(grpcStream); + const findMany = vi.fn() + .mockResolvedValueOnce([createRepo(2, 'repo-2'), createRepo(1, 'repo-1')]) + .mockResolvedValueOnce([createRepo(3, 'repo-3')]); + const prisma = { repo: { findMany } } as unknown as PrismaClient; + const stream = await zoektStreamSearch(searchRequest, prisma); + const reader = stream.getReader(); + + for (const ids of [[1, 1, 2, 2], [1, 2, 3, 3]]) { + grpcStream.emit('data', { response_chunk: { files: ids.map(id => createFile(id)) } }); + const chunk = await reader.read(); + const response = JSON.parse(new TextDecoder().decode(chunk.value).slice('data: '.length)) as Extract; + expect(response.files.map(file => file.repositoryId)).toEqual(ids); + expect(response.repositoryInfo.map(repo => repo.id)).toEqual([...new Set(ids)]); + for (const file of response.files) { + expect(response.repositoryInfo.find(repo => repo.id === file.repositoryId)?.name).toBe(file.repository); + expect(file.repository).toBe(`repo-${file.repositoryId}`); + } + } + + expect(findMany).toHaveBeenCalledTimes(2); + expect(findMany).toHaveBeenNthCalledWith(1, { where: { id: { in: [1, 2] } } }); + expect(findMany).toHaveBeenNthCalledWith(2, { where: { id: { in: [3] } } }); + grpcStream.emit('end'); + while (!(await reader.read()).done) { + // Drain the final statistics and completion marker. + } + expect(mocks.close).toHaveBeenCalledOnce(); + }); + + test('keeps missing IDs and legacy names cached only for the current stream', async () => { + const findMany = vi.fn().mockResolvedValue([createRepo(2, 'visible-repo')]); + const findFirst = vi.fn().mockResolvedValue(null); + const prisma = { repo: { findMany, findFirst } } as unknown as PrismaClient; + const files = [createFile(1), createFile(undefined, 'missing-repo'), createFile(2)]; + + for (let request = 0; request < 2; request++) { + const grpcStream = Object.assign(new EventEmitter(), { + pause: vi.fn(), + resume: vi.fn(), + cancel: vi.fn(), + }); + mocks.streamSearch.mockReturnValue(grpcStream); + const reader = (await zoektStreamSearch(searchRequest, prisma)).getReader(); + for (let chunk = 0; chunk < 2; chunk++) { + grpcStream.emit('data', { response_chunk: { files } }); + const result = await reader.read(); + const response = JSON.parse(new TextDecoder().decode(result.value).slice('data: '.length)); + expect(response.files.map((file: SearchResultFile) => file.repositoryId)).toEqual([2]); + expect(response.repositoryInfo).toEqual([expect.objectContaining({ id: 2, name: 'visible-repo' })]); + } + grpcStream.emit('end'); + while (!(await reader.read()).done) { + // Drain the final statistics and completion marker. + } + } + + expect(findMany).toHaveBeenCalledTimes(2); + expect(findMany).toHaveBeenCalledWith({ where: { id: { in: [1, 2] } } }); + expect(findFirst).toHaveBeenCalledTimes(2); + expect(findFirst).toHaveBeenCalledWith({ where: { name: 'missing-repo' } }); + }); + + test('keeps numeric IDs separate from legacy names in mixed shards', async () => { + const files = [createFile(1), createFile(undefined, '1'), createFile(1), createFile(undefined, '1')]; + mocks.search.mockImplementation((_request, _metadata, callback) => { + callback(null, { files }); + }); + const findMany = vi.fn().mockResolvedValue([createRepo(1, 'numeric-repo')]); + const findFirst = vi.fn().mockResolvedValue(createRepo(2, '1')); + const prisma = { repo: { findMany, findFirst } } as unknown as PrismaClient; + + const response = await zoektSearch(searchRequest, prisma); + + expect(response.files.map(file => file.repositoryId)).toEqual([1, 2, 1, 2]); + expect(response.repositoryInfo.map(repo => repo.name)).toEqual(['numeric-repo', '1']); + expect(findMany).toHaveBeenCalledExactlyOnceWith({ where: { id: { in: [1] } } }); + expect(findFirst).toHaveBeenCalledExactlyOnceWith({ where: { name: '1' } }); + }); }); diff --git a/packages/web/src/features/search/zoektSearcher.ts b/packages/web/src/features/search/zoektSearcher.ts index 087e44043..4f450455c 100644 --- a/packages/web/src/features/search/zoektSearcher.ts +++ b/packages/web/src/features/search/zoektSearcher.ts @@ -136,7 +136,7 @@ export const zoektSearch = async (searchRequest: ZoektGrpcSearchRequest, prisma: }); }); - const reposMapCache = await createReposMapForChunk(response, new Map(), prisma); + const reposMapCache = await createReposMapForChunk(response, new Map(), prisma); const { stats, files, repositoryInfo } = await transformZoektSearchResponse(response, reposMapCache); return { @@ -210,7 +210,7 @@ export const zoektStreamSearch = async (searchRequest: ZoektGrpcSearchRequest, p // `_reposMapCache` is used to cache repository metadata across all chunks. // This reduces the number of database queries required to transform file matches. - const _reposMapCache = new Map(); + const _reposMapCache = new Map(); // Handle incoming data chunks grpcStream.on('data', async (chunk: ZoektGrpcStreamSearchResponse) => { @@ -339,43 +339,42 @@ const encodeSSEREsponseChunk = (response: object | string) => { // Creates a mapping between all repository ids in a given response // chunk. The mapping allows us to efficiently lookup repository metadata. -const createReposMapForChunk = async (chunk: ZoektGrpcSearchResponse, reposMapCache: Map, prisma: PrismaClient): Promise> => { +const createReposMapForChunk = async (chunk: ZoektGrpcSearchResponse, reposMapCache: Map, prisma: PrismaClient): Promise> => { const reposMap = new Map(); - await Promise.all(chunk.files.map(async (file) => { - const id = getRepoIdForFile(file); - - const repo = await (async () => { - // If it's in the cache, return the cached value. - if (reposMapCache.has(id)) { - return reposMapCache.get(id); + const repoIds = [...new Set(chunk.files.map(getRepoIdForFile))]; + const uncachedIds = repoIds.filter(id => !reposMapCache.has(id)); + const numericIds = uncachedIds.filter((id): id is number => typeof id === 'number'); + const names = uncachedIds.filter((id): id is string => typeof id === 'string'); + + await Promise.all([ + (async () => { + if (numericIds.length === 0) { + return; } - // Otherwise, query the database for the record. - const repo = typeof id === 'number' ? - await prisma.repo.findUnique({ - where: { - id: id, - }, - }) : - await prisma.repo.findFirst({ - where: { - name: id, - }, - }); - - // If a repository is found, cache it for future lookups. - if (repo) { - reposMapCache.set(id, repo); + const repos = await prisma.repo.findMany({ where: { id: { in: numericIds } } }); + // Cache missing rows too, so later chunks do not repeat those lookups. + for (const id of numericIds) { + reposMapCache.set(id, null); } + for (const repo of repos) { + reposMapCache.set(repo.id, repo); + } + })(), + ...names.map(async name => { + // Legacy shards use names, which are not guaranteed to be unique. + const repo = await prisma.repo.findFirst({ where: { name } }); + reposMapCache.set(name, repo); + }), + ]); - return repo; - })(); - + for (const id of repoIds) { + const repo = reposMapCache.get(id); // Only add the repository to the map if it was found. if (repo) { reposMap.set(id, repo); } - })); + } return reposMap; }