dbmole-mcp/test/unit/db/sqlGuard.test.ts

import { describe, expect, it } from 'vitest'
import { assertSafeStatement, SqlGuardError } from '../../../src/db/sqlGuard.js'

describe('assertSafeStatement', () => {
    describe('single-statement guard', () => {
        it('allows a single statement', () => {
            expect(() => assertSafeStatement('select 1', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('select 1', 'mysql')).not.toThrow()
        })

        it('allows a single trailing semicolon', () => {
            expect(() => assertSafeStatement('select 1;', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('select 1;   ', 'postgres')).not.toThrow()
        })

        it('rejects two statements', () => {
            expect(() => assertSafeStatement('select 1; select 2', 'postgres')).toThrow(
                /one SQL statement/
            )
            expect(() => assertSafeStatement('select 1; select 2', 'mysql')).toThrow(
                /one SQL statement/
            )
        })

        it('throws SqlGuardError specifically', () => {
            expect(() => assertSafeStatement('select 1; select 2', 'postgres')).toThrow(
                SqlGuardError
            )
        })

        it('rejects two statements with a comment after the trailing semicolon line', () => {
            expect(() => assertSafeStatement('select 1;\nselect 2', 'postgres')).toThrow(
                /one SQL statement/
            )
        })

        it('allows a trailing semicolon followed only by a comment', () => {
            expect(() => assertSafeStatement('select 1; -- done', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('select 1; /* done */', 'postgres')).not.toThrow()
        })

        it('ignores semicolons inside single-quoted strings', () => {
            expect(() => assertSafeStatement("select 'a;b'", 'postgres')).not.toThrow()
            expect(() => assertSafeStatement("select 'a;b'", 'mysql')).not.toThrow()
        })

        it('handles single-quote escape by doubling', () => {
            expect(() => assertSafeStatement("select 'a''b;c'", 'postgres')).not.toThrow()
        })

        it('ignores semicolons inside double-quoted identifiers', () => {
            expect(() => assertSafeStatement('select "a;b"', 'postgres')).not.toThrow()
        })

        it('ignores semicolons inside mysql backtick identifiers', () => {
            expect(() => assertSafeStatement('select `a;b`', 'mysql')).not.toThrow()
        })

        it('ignores semicolons inside pg dollar-quoted strings', () => {
            expect(() => assertSafeStatement('select $$ a ; b $$', 'postgres')).not.toThrow()
        })

        it('ignores semicolons inside tagged pg dollar-quoted strings', () => {
            expect(() =>
                assertSafeStatement('select $fn$ a ; select 2 $fn$', 'postgres')
            ).not.toThrow()
        })

        it('ignores semicolons inside line comments', () => {
            expect(() => assertSafeStatement('select 1 -- ; select 2', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('select 1 # ; select 2', 'mysql')).not.toThrow()
        })

        it('ignores semicolons inside block comments', () => {
            expect(() => assertSafeStatement('select /* ; */ 1', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('select /* ; */ 1', 'mysql')).not.toThrow()
        })

        it('treats # as a comment only for mysql', () => {
            // For postgres '#' is not a comment, so the ; after it is top-level
            expect(() => assertSafeStatement('select 1 # ; select 2', 'postgres')).toThrow(
                /one SQL statement/
            )
        })

        it('handles nested block comments for postgres', () => {
            expect(() =>
                assertSafeStatement('/* outer /* inner */ still comment */ select 1', 'postgres')
            ).not.toThrow()
        })

        it('handles mysql backslash escape inside string', () => {
            // \' does not end the string for mysql, so the ; stays inside
            expect(() => assertSafeStatement("select 'a\\'; select 2'", 'mysql')).not.toThrow()
        })
    })

    describe('session-statement guard', () => {
        it('rejects BEGIN', () => {
            expect(() => assertSafeStatement('BEGIN', 'postgres')).toThrow(/session-level/)
        })

        it('rejects set session transaction read write', () => {
            expect(() =>
                assertSafeStatement('set session transaction read write', 'mysql')
            ).toThrow(/session-level/)
        })

        it('rejects USE other', () => {
            expect(() => assertSafeStatement('USE other', 'mysql')).toThrow(/session-level/)
        })

        it('is case-insensitive', () => {
            expect(() => assertSafeStatement('CoMmIt', 'postgres')).toThrow(/session-level/)
        })

        it('skips leading whitespace', () => {
            expect(() => assertSafeStatement('   \n  begin', 'postgres')).toThrow(/session-level/)
        })

        it('skips a leading comment', () => {
            expect(() => assertSafeStatement('/* comment */ begin', 'postgres')).toThrow(
                /session-level/
            )
        })

        it('rejects start, rollback, savepoint, release, prepare, deallocate, discard, reset', () => {
            for (const kw of [
                'start transaction',
                'rollback',
                'savepoint s1',
                'release savepoint s1',
                'prepare p as select 1',
                'deallocate p',
                'discard all',
                'reset all'
            ]) {
                expect(() => assertSafeStatement(kw, 'postgres')).toThrow(/session-level/)
            }
        })

        it('does not false-positive on identifiers that start with a keyword', () => {
            expect(() =>
                assertSafeStatement('select * from SETTINGS_TABLE', 'postgres')
            ).not.toThrow()
            expect(() => assertSafeStatement('select * from usefulness', 'mysql')).not.toThrow()
            expect(() => assertSafeStatement('select * from beginning', 'postgres')).not.toThrow()
        })

        it('allows a leading paren before a real query', () => {
            expect(() =>
                assertSafeStatement('(select 1) union (select 2)', 'postgres')
            ).not.toThrow()
        })

        it('skips a leading line comment before the keyword', () => {
            expect(() => assertSafeStatement('-- note\nbegin', 'postgres')).toThrow(/session-level/)
        })

        it('skips a leading mysql hash comment before the keyword', () => {
            expect(() => assertSafeStatement('# note\nuse other', 'mysql')).toThrow(/session-level/)
        })

        it('does not treat a leading hash comment as a comment for postgres', () => {
            // For postgres '#' is not a comment, so the first token is '#', not a keyword.
            expect(() => assertSafeStatement('# note\nselect 1', 'postgres')).not.toThrow()
        })

        it('treats empty or comment-only sql as having no keyword', () => {
            expect(() => assertSafeStatement('', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('   ', 'postgres')).not.toThrow()
            expect(() => assertSafeStatement('-- just a comment', 'postgres')).not.toThrow()
        })

        it('does not treat a statement starting with a non-word char as a keyword', () => {
            expect(() => assertSafeStatement('123', 'postgres')).not.toThrow()
        })
    })
})