supabase/apps/studio/data/sql/__tests__/utils.test.ts

import { safeSql } from '@supabase/pg-meta'
import { describe, expect, it, test } from 'vitest'

import { applyAutoLimit, getSqlErrorLines, trimTrailingSemicolons } from '../utils'

describe('getSqlErrorLines', () => {
  it('returns formattedError lines when present', () => {
    const lines = getSqlErrorLines({
      message: 'permission denied for table users',
      formattedError:
        'ERROR:  42501: permission denied for table users\n' +
        'HINT:  To grant access to anon on a specific table:\n' +
        '  GRANT SELECT ON TABLE public.users TO anon;',
    })

    expect(lines).toEqual([
      'ERROR:  42501: permission denied for table users',
      'HINT:  To grant access to anon on a specific table:',
      '  GRANT SELECT ON TABLE public.users TO anon;',
    ])
  })

  it('strips empty lines from formattedError', () => {
    const lines = getSqlErrorLines({
      formattedError: 'ERROR: boom\n\nHINT: retry\n',
    })

    expect(lines).toEqual(['ERROR: boom', 'HINT: retry'])
  })

  it('falls back to message lines when formattedError is missing and message is multi-line', () => {
    const lines = getSqlErrorLines({
      message:
        'ERROR:  42501: permission denied for table users\n' +
        'HINT:  To grant access to anon on a specific table:\n' +
        '  GRANT SELECT ON TABLE public.users TO anon;',
    })

    expect(lines).toEqual([
      'ERROR:  42501: permission denied for table users',
      'HINT:  To grant access to anon on a specific table:',
      '  GRANT SELECT ON TABLE public.users TO anon;',
    ])
  })

  it('returns empty array for a single-line message so callers render the fallback', () => {
    const lines = getSqlErrorLines({ message: 'permission denied for table users' })
    expect(lines).toEqual([])
  })

  it('returns empty array when both fields are missing', () => {
    expect(getSqlErrorLines({})).toEqual([])
  })

  it('returns empty array when message is an empty string', () => {
    expect(getSqlErrorLines({ message: '' })).toEqual([])
  })

  it('returns empty array when message only contains whitespace newlines', () => {
    // Only empty segments after filtering — treated as single-line
    expect(getSqlErrorLines({ message: '\n\n' })).toEqual([])
  })

  it('prefers formattedError even when message is also multi-line', () => {
    const lines = getSqlErrorLines({
      message: 'message line 1\nmessage line 2',
      formattedError: 'formatted line 1\nformatted line 2',
    })

    expect(lines).toEqual(['formatted line 1', 'formatted line 2'])
  })

  it('falls through to message when formattedError is empty string', () => {
    const lines = getSqlErrorLines({
      message: 'ERROR: line 1\nHINT: line 2',
      formattedError: '',
    })

    expect(lines).toEqual(['ERROR: line 1', 'HINT: line 2'])
  })
})

describe('trimTrailingSemicolons', () => {
  test('removes a single trailing semicolon', () => {
    const sql = safeSql`select * from countries;`
    expect(trimTrailingSemicolons(sql)).toBe('select * from countries')
  })
  test('removes multiple trailing semicolons', () => {
    const sql = safeSql`select * from countries;;;;;;;`
    expect(trimTrailingSemicolons(sql)).toBe('select * from countries')
  })
  test('leaves a fragment with no trailing semicolon unchanged', () => {
    const sql = safeSql`select * from countries`
    expect(trimTrailingSemicolons(sql)).toBe('select * from countries')
  })
  test('does not touch semicolons that are not trailing', () => {
    const sql = safeSql`select 1; select 2`
    expect(trimTrailingSemicolons(sql)).toBe('select 1; select 2')
  })
})

describe('applyAutoLimit', () => {
  test('Should return false if limit passed is <= 0', () => {
    const sql = safeSql`select * from countries;`
    const limit = -1
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return true if limit passed is > 0', () => {
    const sql = safeSql`select * from countries;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(true)
  })
  test('Should return false if query already has a limit', () => {
    const sql = safeSql`select * from countries limit 10;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit (check for case-insensitiveness)', () => {
    const sql = safeSql`SELECT * FROM countries LIMIT 10;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit with whitespace before the semi colon', () => {
    const sql = safeSql`select * from countries limit 10 ;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit and offset', () => {
    const sql = safeSql`select * from countries limit 10 offset 0;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit and offset with whitespace before the semi colon', () => {
    const sql = safeSql`select * from countries limit 10 offset 0 ;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit and offset (flip order of limit and offset)', () => {
    const sql = safeSql`select * from countries offset 0 limit 1;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query already has a limit, even if no value provided for limit', () => {
    const sql = safeSql`select * from countries limit`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query uses `FETCH FIRST` instead of limit ', () => {
    const sql = safeSql`select * from countries FETCH FIRST 5 rows only`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query uses `fetch first` instead of limit ', () => {
    const sql = safeSql`select * from countries fetch first 5 rows only`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query uses `fetch   first` (with random spaces) instead of limit ', () => {
    const sql = safeSql`select * from countries FETCH FIRST 5 rows only`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query uses `FETCH NEXT` instead of limit ', () => {
    const sql = safeSql`select * from countries offset 5 rows FETCH NEXT 5 rows only`
    const { sql: result, appendAutoLimit } = applyAutoLimit(sql, 100)
    expect(appendAutoLimit).toBe(false)
    expect(result).toBe(sql)
  })
  test('Should return false if query is not a select statement', () => {
    const sql = safeSql`create table test (id int8 primary key, name varchar);`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if there are multiple queries I', () => {
    const sql1 = safeSql`select * from countries;
select * from cities;`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql1, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if there are multiple queries II', () => {
    const sql1 = safeSql`select * from countries;
select * from cities`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql1, limit)
    expect(appendAutoLimit).toBe(false)
  })
  // [Joshen] Opting to just avoid appending in this case to prevent making the logic overly complex atm
  test('Should return false if query has with a comment I', () => {
    const sql = safeSql`-- This is a comment
select * from cities`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })
  test('Should return false if query has with a comment II', () => {
    const sql = safeSql`select * from cities
-- This is a comment`
    const limit = 100
    const { appendAutoLimit } = applyAutoLimit(sql, limit)
    expect(appendAutoLimit).toBe(false)
  })

  // [Joshen] These will just need to test the cases when appendAutoLimit returns true then
  test('Should add the limit param properly if query ends without a semi colon', () => {
    const sql = safeSql`select * from countries`
    const limit = 100
    const { sql: formattedSql } = applyAutoLimit(sql, limit)
    expect(formattedSql).toBe('select * from countries limit 100;')
  })
  test('Should add the limit param properly if query ends with a semi colon', () => {
    const sql = safeSql`select * from countries;`
    const limit = 100
    const { sql: formattedSql } = applyAutoLimit(sql, limit)
    expect(formattedSql).toBe('select * from countries limit 100;')
  })
  test('Should add the limit param properly if query ends with multiple semi colon', () => {
    const sql = safeSql`select * from countries;;;;;;;`
    const limit = 100
    const { sql: formattedSql } = applyAutoLimit(sql, limit)
    expect(formattedSql).toBe('select * from countries limit 100;')
  })
  test('Should not append a limit if query already has one with whitespace before the semi colon', () => {
    const sql = safeSql`select * from countries limit 10 ;`
    const limit = 100
    const { sql: formattedSql } = applyAutoLimit(sql, limit)
    expect(formattedSql).toBe('select * from countries limit 10 ;')
  })
  test('returns the SafeSqlFragment result unchanged when no limit is appended', () => {
    const sql = safeSql`select * from countries limit 10;`
    const { sql: formattedSql } = applyAutoLimit(sql, 100)
    expect(formattedSql).toBe(sql)
  })
})