import { beforeEach, describe, expect, it, vi } from 'vitest'; import { MySQLAdapter } from '../../../src/adapters/mysql'; import { PostgreSQLAdapter } from '../../../src/adapters/postgresql'; import { SnowflakeAdapter } from '../../../src/adapters/snowflake'; import { SQLServerAdapter } from '../../../src/adapters/sqlserver'; describe('MaxRows Limiting Tests', () => { describe('PostgreSQL Adapter', () => { let adapter: PostgreSQLAdapter; let mockClient: { connect: ReturnType; query: ReturnType; end: ReturnType; }; let mockCursor: { read: ReturnType; close: ReturnType; _result: { fields: Array<{ name: string; dataTypeID: number; dataTypeSize: number }> }; }; beforeEach(() => { adapter = new PostgreSQLAdapter(); mockClient = { connect: vi.fn(), query: vi.fn(), end: vi.fn(), }; mockCursor = { read: vi.fn(), close: vi.fn(), _result: { fields: [ { name: 'id', dataTypeID: 23, dataTypeSize: 4 }, { name: 'name', dataTypeID: 25, dataTypeSize: -1 }, ], }, }; (adapter as PostgreSQLAdapter & { client: typeof mockClient; connected: boolean }).client = mockClient; (adapter as PostgreSQLAdapter & { client: typeof mockClient; connected: boolean }).connected = true; }); it('should limit results to exactly 1 row when maxRows=1', async () => { // Mock cursor behavior for 1 row limit mockClient.query.mockReturnValue(mockCursor); mockCursor.read.mockImplementationOnce( (count: number, callback: (err: unknown, rows: unknown[]) => void) => { // First read returns 2 rows (more than requested) callback(null, [ { id: 1, name: 'User 1' }, { id: 2, name: 'User 2' }, ]); } ); mockCursor.close.mockImplementation((callback: (err: unknown) => void) => callback(null)); const result = await adapter.query('SELECT * FROM users', undefined, 1); expect(result.rows).toHaveLength(1); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.hasMoreRows).toBe(true); expect(mockCursor.read).toHaveBeenCalledWith(2, expect.any(Function)); // When maxRows=1, readSize = min(1, 1) + 1 = 2 }); it('should limit results to exactly 5 rows when maxRows=5', async () => { mockClient.query.mockReturnValue(mockCursor); const allRows = Array.from({ length: 10 }, (_, i) => ({ id: i + 1, name: `User ${i + 1}`, })); mockCursor.read.mockImplementationOnce( (count: number, callback: (err: unknown, rows: unknown[]) => void) => { // Return 6 rows (more than requested 5) callback(null, allRows.slice(0, 6)); } ); mockCursor.close.mockImplementation((callback: (err: unknown) => void) => callback(null)); const result = await adapter.query('SELECT * FROM users', undefined, 5); expect(result.rows).toHaveLength(5); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.rows[4]).toEqual({ id: 5, name: 'User 5' }); expect(result.hasMoreRows).toBe(true); }); it('should return all rows when result set is smaller than maxRows', async () => { mockClient.query.mockReturnValue(mockCursor); mockCursor.read .mockImplementationOnce( (count: number, callback: (err: unknown, rows: unknown[]) => void) => { // Return only 3 rows when asking for 10 callback(null, [ { id: 1, name: 'User 1' }, { id: 2, name: 'User 2' }, { id: 3, name: 'User 3' }, ]); } ) .mockImplementationOnce( (count: number, callback: (err: unknown, rows: unknown[]) => void) => { // No more rows callback(null, []); } ); mockCursor.close.mockImplementation((callback: (err: unknown) => void) => callback(null)); const result = await adapter.query('SELECT * FROM users', undefined, 10); expect(result.rows).toHaveLength(3); expect(result.hasMoreRows).toBe(false); }); }); describe('MySQL Adapter', () => { let adapter: MySQLAdapter; let mockConnection: { execute: ReturnType; end: ReturnType; }; beforeEach(() => { adapter = new MySQLAdapter(); mockConnection = { execute: vi.fn(), end: vi.fn(), }; ( adapter as MySQLAdapter & { connection: typeof mockConnection; connected: boolean } ).connection = mockConnection; ( adapter as MySQLAdapter & { connection: typeof mockConnection; connected: boolean } ).connected = true; }); it('should limit results to exactly 1 row when maxRows=1', async () => { // MySQL adapter limits in memory after fetching mockConnection.execute.mockResolvedValue([ [ { id: 1, name: 'User 1' }, { id: 2, name: 'User 2' }, { id: 3, name: 'User 3' }, ], [ { name: 'id', type: 3, flags: 0 }, { name: 'name', type: 253, flags: 0 }, ], ]); const result = await adapter.query('SELECT * FROM users', undefined, 1); expect(result.rows).toHaveLength(1); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.hasMoreRows).toBe(true); }); it('should limit results to exactly 5 rows when maxRows=5', async () => { const allRows = Array.from({ length: 10 }, (_, i) => ({ id: i + 1, name: `User ${i + 1}`, })); mockConnection.execute.mockResolvedValue([ allRows, [ { name: 'id', type: 3, flags: 0 }, { name: 'name', type: 253, flags: 0 }, ], ]); const result = await adapter.query('SELECT * FROM users', undefined, 5); expect(result.rows).toHaveLength(5); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.rows[4]).toEqual({ id: 5, name: 'User 5' }); expect(result.hasMoreRows).toBe(true); }); }); describe('Snowflake Adapter', () => { let adapter: SnowflakeAdapter; let mockConnection: { execute: ReturnType; destroy: ReturnType; }; let mockStatement: { getColumns: ReturnType; streamRows: ReturnType; }; let mockStream: { on: ReturnType; destroy: ReturnType; }; beforeEach(() => { adapter = new SnowflakeAdapter(); mockStatement = { getColumns: vi.fn().mockReturnValue([ { getName: () => 'id', getType: () => 'NUMBER', isNullable: () => false, getScale: () => 0, getPrecision: () => 10, }, { getName: () => 'name', getType: () => 'VARCHAR', isNullable: () => true, getScale: () => 0, getPrecision: () => 0, }, ]), streamRows: vi.fn(), }; mockStream = { on: vi.fn(), destroy: vi.fn(), destroyed: false, }; mockConnection = { execute: vi.fn(), }; mockStatement.streamRows.mockReturnValue(mockStream); ( adapter as SnowflakeAdapter & { connection: typeof mockConnection; connected: boolean } ).connection = mockConnection; ( adapter as SnowflakeAdapter & { connection: typeof mockConnection; connected: boolean } ).connected = true; }); it('should limit results to exactly 1 row when maxRows=1', async () => { let dataHandler: (data: unknown) => void; let endHandler: () => void; mockStream.on.mockImplementation((event: string, handler: (data?: unknown) => void) => { if (event === 'data') dataHandler = handler; if (event === 'end') endHandler = handler; return mockStream; }); mockConnection.execute.mockImplementation( (options: { streamResult: boolean; complete: (err?: unknown) => void }) => { expect(options.streamResult).toBe(true); // Defer the callback to allow the statement to be returned first setTimeout(() => { options.complete(undefined); // Now simulate streaming after complete callback dataHandler({ id: 1, name: 'User 1' }); dataHandler({ id: 2, name: 'User 2' }); // This should trigger destroy // Verify stream was destroyed after 1 row expect(mockStream.destroy).toHaveBeenCalled(); endHandler(); }, 0); return mockStatement; } ); const result = await adapter.query('SELECT * FROM users', undefined, 1); expect(result.rows).toHaveLength(1); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.hasMoreRows).toBe(true); }); }); describe('SQL Server Adapter', () => { let adapter: SQLServerAdapter; let mockPool: { request: ReturnType; close: ReturnType; }; let mockRequest: { input: ReturnType; query: ReturnType; on: ReturnType; pause: ReturnType; cancel: ReturnType; stream: boolean; }; beforeEach(() => { adapter = new SQLServerAdapter(); mockRequest = { input: vi.fn().mockReturnThis(), query: vi.fn(), on: vi.fn(), pause: vi.fn(), cancel: vi.fn(), stream: false, }; mockPool = { request: vi.fn().mockReturnValue(mockRequest), close: vi.fn(), }; (adapter as SQLServerAdapter & { pool: typeof mockPool; connected: boolean }).pool = mockPool; (adapter as SQLServerAdapter & { pool: typeof mockPool; connected: boolean }).connected = true; }); it('should limit results to exactly 1 row when maxRows=1', async () => { let rowHandler: (row: Record) => void; let doneHandler: () => void; let recordsetHandler: (columns: Record) => void; mockRequest.on.mockImplementation((event: string, handler: (arg?: unknown) => void) => { if (event === 'row') rowHandler = handler; if (event === 'done') doneHandler = handler; if (event === 'recordset') recordsetHandler = handler; }); const queryPromise = adapter.query('SELECT * FROM users', undefined, 1); // Verify streaming was enabled expect(mockRequest.stream).toBe(true); // Emit column metadata recordsetHandler!({ id: { type: () => ({ name: 'int' }), nullable: false }, name: { type: () => ({ name: 'varchar' }), nullable: true }, }); // Emit rows rowHandler!({ id: 1, name: 'User 1' }); rowHandler!({ id: 2, name: 'User 2' }); // This should trigger pause and cancel // Verify request was paused and cancelled expect(mockRequest.pause).toHaveBeenCalled(); expect(mockRequest.cancel).toHaveBeenCalled(); doneHandler!(); const result = await queryPromise; expect(result.rows).toHaveLength(1); expect(result.rows[0]).toEqual({ id: 1, name: 'User 1' }); expect(result.hasMoreRows).toBe(true); }); }); });