import { beforeAll, describe, expect, test, vi } from 'vitest'; process.env.DATABASE_URL ??= 'postgresql://user:pass@localhost:5432/umami?schema=public'; delete process.env.DATABASE_REPLICA_URL; vi.mock('@prisma/adapter-pg', () => ({ PrismaPg: class PrismaPg {}, })); vi.mock('@prisma/extension-read-replicas', () => ({ readReplicas: () => () => ({}), })); vi.mock('@/generated/prisma/client', () => ({ PrismaClient: class PrismaClient { $executeRawUnsafe = vi.fn(); $queryRawUnsafe = vi.fn(); $transaction = vi.fn(); $on = vi.fn(); $extends() { return this; } }, })); let getRawQueryClient!: typeof import('./prisma').getRawQueryClient; let prisma!: typeof import('./prisma').default; beforeAll(async () => { ({ getRawQueryClient, default: prisma } = await import('./prisma')); }); interface RawQueryClient { $executeRawUnsafe: (query: string, ...params: any[]) => unknown; $queryRawUnsafe: (query: string, ...params: any[]) => unknown; $primary?: () => unknown; $replica?: () => unknown; } function createClient(): RawQueryClient { return { $executeRawUnsafe: vi.fn(), $queryRawUnsafe: vi.fn(), }; } describe('getRawQueryClient', () => { test('uses a replica client for read queries when replicas are enabled', () => { const replica = createClient(); const client = { ...createClient(), $replica: vi.fn(() => replica), }; expect(getRawQueryClient(client, { useReplica: true })).toBe(replica); }); test('keeps read queries on the primary client when replicas are disabled', () => { const client = { ...createClient(), $replica: vi.fn(() => createClient()), }; expect(getRawQueryClient(client, { useReplica: false })).toBe(client); }); test('uses the primary client for raw writes when available', () => { const primary = createClient(); const client = { ...createClient(), $primary: vi.fn(() => primary), $replica: vi.fn(() => createClient()), }; expect(getRawQueryClient(client, { useReplica: true, write: true })).toBe(primary); }); test('falls back to the current client for raw writes without a primary helper', () => { const client = createClient(); expect(getRawQueryClient(client, { write: true })).toBe(client); }); }); describe('pagedRawQuery default ordering', () => { function mockQueries(count = '4') { const queryRaw = vi.mocked((prisma.client as any).$queryRawUnsafe); queryRaw.mockClear(); queryRaw.mockImplementation(async (sql: string) => sql.includes('count(*) as num') ? [{ num: count }] : [{ id: 1 }], ); return queryRaw; } test.each([undefined, 5])( 'applies trusted default order only to the page query with cap %s', async maxResults => { const queryRaw = mockQueries('5'); const result = await prisma.pagedRawQuery( 'select * from thing', {}, { page: 2, pageSize: 2, maxResults }, 'test', 'max(created_at) desc, session_id', ); expect(queryRaw).toHaveBeenCalledTimes(2); expect(queryRaw.mock.calls[0][0]).not.toContain('order by max(created_at) desc'); expect(queryRaw.mock.calls[1][0]).toContain('order by max(created_at) desc, session_id'); expect(queryRaw.mock.calls[1][0]).toContain('limit 2 offset 2'); expect(result).toEqual({ data: [{ id: 1 }], count: 5, page: 2, pageSize: 2, isCapped: !!maxResults, orderBy: undefined, }); }, ); test('keeps explicit ordering and legacy no-order behavior', async () => { const queryRaw = mockQueries(); const result = await prisma.pagedRawQuery( 'select * from thing', {}, { page: 1, pageSize: 2, orderBy: 'name', sortDescending: true }, undefined, 'max(created_at) desc, session_id', ); expect(queryRaw.mock.calls[1][0]).toContain('order by name desc'); expect(queryRaw.mock.calls[1][0]).not.toContain('order by max(created_at) desc, session_id'); expect(result.orderBy).toBe('name'); queryRaw.mockClear(); await prisma.pagedRawQuery('select * from thing', {}, { page: 1, pageSize: 2 }); expect(queryRaw.mock.calls[1][0]).not.toMatch(/order by/); }); });