import { describe, expect, it } from 'vitest'; import { applyQueryAutoLimit } from './queryAutoLimit'; describe('applyQueryAutoLimit', () => { const limitDialects = [ 'mysql', 'goldendb', 'mariadb', 'oceanbase', 'diros', 'doris', 'starrocks', 'sphinx', 'postgres', 'postgresql', 'kingbase', 'kingbase8', 'highgo', 'vastbase', 'opengauss', 'gaussdb', 'iris', 'intersystemsiris', 'sqlite', 'sqlite3', 'duckdb', 'clickhouse', 'tdengine', 'iotdb', ]; it.each(limitDialects)('adds generic LIMIT for %s connections', (dbType) => { expect(applyQueryAutoLimit('SELECT * FROM users', dbType, 500).sql) .toBe('SELECT * FROM users LIMIT 500'); }); it.each([ ['oracle'], ['dameng'], ['dm'], ['dm8'], ])('adds FETCH FIRST limit for %s connections', (dbType) => { expect(applyQueryAutoLimit('SELECT * FROM MYCIMLED.EDC_LOG', dbType, 500).sql) .toBe('SELECT * FROM MYCIMLED.EDC_LOG FETCH FIRST 500 ROWS ONLY'); }); it.each([ ['sqlserver'], ['mssql'], ['sql_server'], ['sql-server'], ])('adds TOP limit for %s connections', (dbType) => { expect(applyQueryAutoLimit('SELECT * FROM users', dbType, 500).sql) .toBe('SELECT TOP 500 * FROM users'); }); it('adds SQL Server TOP after DISTINCT', () => { expect(applyQueryAutoLimit('SELECT DISTINCT name FROM users', 'sqlserver', 500).sql) .toBe('SELECT DISTINCT TOP 500 name FROM users'); }); it.each([ ['oracle', 'SELECT * FROM users FETCH FIRST 500 ROWS ONLY'], ['dm8', 'SELECT * FROM users FETCH FIRST 500 ROWS ONLY'], ['mssql', 'SELECT TOP 500 * FROM users'], ['postgresql', 'SELECT * FROM users LIMIT 500'], ['gauss-db', 'SELECT * FROM users LIMIT 500'], ['doris', 'SELECT * FROM users LIMIT 500'], ['starrocks', 'SELECT * FROM users LIMIT 500'], ['sqlite3', 'SELECT * FROM users LIMIT 500'], ])('uses custom driver dialect %s', (driver, expected) => { expect(applyQueryAutoLimit('SELECT * FROM users', 'custom', 500, driver).sql) .toBe(expected); }); it('keeps trailing semicolon and comments after injected Oracle limit', () => { expect(applyQueryAutoLimit('SELECT * FROM MYCIMLED.EDC_LOG; -- preview', 'oracle', 500).sql) .toBe('SELECT * FROM MYCIMLED.EDC_LOG FETCH FIRST 500 ROWS ONLY; -- preview'); }); it('uses Oracle FETCH FIRST limit for simple table queries', () => { expect(applyQueryAutoLimit('select 1 from xxx', 'oracle', 500).sql) .toBe('select 1 from xxx FETCH FIRST 500 ROWS ONLY'); }); it('keeps ORDER BY semantics with Oracle FETCH FIRST', () => { expect(applyQueryAutoLimit('SELECT * FROM users ORDER BY created_at DESC', 'oracle', 100).sql) .toBe('SELECT * FROM users ORDER BY created_at DESC FETCH FIRST 100 ROWS ONLY'); }); it('does not add another generic limit when SQL already limits rows', () => { expect(applyQueryAutoLimit('SELECT * FROM users LIMIT 10', 'mysql', 500).applied) .toBe(false); expect(applyQueryAutoLimit('SELECT * FROM users OFFSET 10 LIMIT 10', 'postgres', 500).applied) .toBe(false); }); it('does not treat nested LIMIT as the outer query limit', () => { expect(applyQueryAutoLimit('SELECT * FROM (SELECT * FROM users LIMIT 10) t', 'postgres', 500).sql) .toBe('SELECT * FROM (SELECT * FROM users LIMIT 10) t LIMIT 500'); }); it('does not add another Oracle limit when Oracle SQL already limits rows', () => { expect(applyQueryAutoLimit('SELECT * FROM users WHERE ROWNUM <= 10', 'oracle', 500).applied) .toBe(false); expect(applyQueryAutoLimit('SELECT * FROM users FETCH FIRST 10 ROWS ONLY', 'oracle', 500).applied) .toBe(false); }); it('does not wrap Oracle FOR UPDATE queries', () => { expect(applyQueryAutoLimit('SELECT * FROM users FOR UPDATE', 'oracle', 500).applied) .toBe(false); }); it('does not add another SQL Server limit when SQL already uses TOP', () => { expect(applyQueryAutoLimit('SELECT TOP 10 * FROM users', 'sqlserver', 500).applied) .toBe(false); }); it('adds generic LIMIT before locking clauses', () => { expect(applyQueryAutoLimit('SELECT * FROM users FOR UPDATE', 'mysql', 500).sql) .toBe('SELECT * FROM users LIMIT 500 FOR UPDATE'); }); it('adds generic LIMIT before OFFSET clauses', () => { expect(applyQueryAutoLimit('SELECT * FROM users OFFSET 10', 'postgres', 500).sql) .toBe('SELECT * FROM users LIMIT 500 OFFSET 10'); }); it('does not limit non-select statements', () => { expect(applyQueryAutoLimit('UPDATE users SET name = \'a\'', 'mysql', 500).applied) .toBe(false); }); });