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('does not treat a digit-leading dollar tag as a dollar quote', () => {
            // $1$ is a positional param $1 followed by a bare $, NOT a dollar
            // quote. The ; after it is top-level, so this is two statements.
            expect(() => assertSafeStatement('SELECT $1$ ; DROP TABLE x', 'postgres')).toThrow(
                /one SQL statement/
            )
        })

        it('allows digits later inside a dollar tag', () => {
            expect(() =>
                assertSafeStatement('select $tag1$ a ; b $tag1$', 'postgres')
            ).not.toThrow()
        })

        it('keeps plain positional params with values unaffected', () => {
            expect(() => assertSafeStatement('SELECT $1', '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()
        })

        it('handles mysql backslash escape inside double-quoted string', () => {
            // mysql treats "..." as a string; \" does not close it, so the ;
            // stays inside the literal and this is a single statement.
            expect(() => assertSafeStatement('SELECT "x\\";y" AS c', 'mysql')).not.toThrow()
        })

        it('handles double-quote doubling in both engines', () => {
            expect(() => assertSafeStatement('select "a""b;c"', 'mysql')).not.toThrow()
            expect(() => assertSafeStatement('select "a""b;c"', 'postgres')).not.toThrow()
        })

        it('does not treat backslash as escape in postgres double-quoted identifiers', () => {
            // For postgres "..." is an identifier with no backslash escapes, so
            // "a\" closes at the second quote and '; SELECT 1' is a second statement.
            expect(() => assertSafeStatement('SELECT "a\\";SELECT 1', 'postgres')).toThrow(
                /one SQL statement/
            )
        })

        it('treats postgres E-strings as backslash-escaped (single statement)', () => {
            // E'it\'s; fine' is one literal: \' is an escaped quote, so the ;
            // stays inside the string and this is a single statement.
            expect(() => assertSafeStatement("SELECT E'it\\'s; fine'", 'postgres')).not.toThrow()
            expect(() => assertSafeStatement("SELECT e'it\\'s; fine'", 'postgres')).not.toThrow()
        })

        it('does not treat plain postgres strings as backslash-escaped', () => {
            // In a plain '...' literal \ is a literal char, so 'a\' closes at the
            // second quote, leaving '; SELECT 1' as a second statement.
            expect(() => assertSafeStatement("SELECT 'a\\'; SELECT 1", 'postgres')).toThrow(
                /one SQL statement/
            )
        })

        it('does not treat a trailing-E identifier as an E-string prefix', () => {
            // The E here is the tail of the identifier `somE`, not a standalone
            // prefix, so backslash escapes stay off: 'a\' closes at the second
            // quote and '; SELECT 1' is a second statement.
            expect(() => assertSafeStatement("SELECT somE'a\\'; SELECT 1", 'postgres')).toThrow(
                /one SQL statement/
            )
        })
    })

    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()
        })
    })
})