ktx/packages/cli/test/context/connections/read-only-sql.test.ts

import { describe, expect, it } from 'vitest';
import { KtxExpectedError, KtxQueryError } from '../../../src/errors.js';
import {
  assertReadOnlySql,
  hoistLeadingCte,
  limitSqlForExecution,
  stripTrailingSqlNoise,
} from '../../../src/context/connections/read-only-sql.js';

describe('assertReadOnlySql', () => {
  it('allows select and with queries', () => {
    expect(assertReadOnlySql('select * from orders')).toBe('select * from orders');
    expect(assertReadOnlySql('with paid as (select * from orders) select * from paid')).toContain('with paid');
  });

  it('rejects mutating statements before opening a database connection', () => {
    expect(() => assertReadOnlySql('delete from orders')).toThrow(
      'Only read-only SELECT/WITH queries can be executed locally',
    );
    expect(() => assertReadOnlySql('create table x(id int)')).toThrow(
      'Only read-only SELECT/WITH queries can be executed locally',
    );
  });

  // A guard refusing the agent's SQL is an expected outcome; classifying it as
  // KtxQueryError keeps reportException from filing it as a ktx fault.
  it('rejects with an expected KtxQueryError, not a bare Error', () => {
    expect(() => assertReadOnlySql('delete from orders')).toThrow(KtxQueryError);
    expect(() => assertReadOnlySql('describe orders')).toThrow(KtxQueryError);
    expect(() => assertReadOnlySql('select 1; drop table orders')).toThrow(KtxQueryError);
    try {
      assertReadOnlySql('describe orders');
      expect.unreachable('expected a throw');
    } catch (error) {
      expect(error).toBeInstanceOf(KtxExpectedError);
    }
  });

  it('accepts read-only queries that begin with leading comments', () => {
    expect(assertReadOnlySql('-- daily widget sales\nselect count(*) from public.widget_sales')).toBe(
      'select count(*) from public.widget_sales',
    );
    expect(assertReadOnlySql('/* block */\n  with paid as (select 1) select * from paid')).toContain('with paid');
  });

  it('still rejects mutating statements hidden behind leading comments', () => {
    expect(() => assertReadOnlySql('-- harmless\n  delete from orders')).toThrow(
      'Only read-only SELECT/WITH queries can be executed locally',
    );
  });

  it('rejects a second statement smuggled after a semicolon', () => {
    expect(() => assertReadOnlySql('select 1; drop table orders')).toThrow(
      'Only one SQL statement can be executed.',
    );
    expect(() => assertReadOnlySql('select 1;\n-- pad\ndelete from orders')).toThrow(
      'Only one SQL statement can be executed.',
    );
    expect(() => assertReadOnlySql('select 1; /* pad */ truncate orders;')).toThrow(
      'Only one SQL statement can be executed.',
    );
  });

  it('accepts trailing semicolons, including repeated ones followed by comments', () => {
    expect(assertReadOnlySql('select 1;')).toBe('select 1;');
    expect(assertReadOnlySql('select 1 ;; \n')).toBe('select 1 ;;');
    expect(assertReadOnlySql('select 1; -- done')).toBe('select 1; -- done');
  });

  it('ignores semicolons inside string literals, quoted identifiers, and comments', () => {
    expect(assertReadOnlySql("select string_agg(name, '; ') from t")).toBe("select string_agg(name, '; ') from t");
    expect(assertReadOnlySql("select 'it''s; quoted' from t")).toBe("select 'it''s; quoted' from t");
    expect(assertReadOnlySql('select ";" from "t;u"')).toBe('select ";" from "t;u"');
    expect(assertReadOnlySql('select 1 -- tail; comment')).toBe('select 1 -- tail; comment');
    expect(assertReadOnlySql('select 1 /* a;b */ + 2')).toBe('select 1 /* a;b */ + 2');
  });

  it('rejects statements smuggled after a string literal that closes a semicolon early', () => {
    expect(() => assertReadOnlySql("select 'a'; delete from orders")).toThrow(
      'Only one SQL statement can be executed.',
    );
  });
});

describe('hoistLeadingCte', () => {
  it('leaves non-CTE SQL untouched', () => {
    expect(hoistLeadingCte('select * from orders')).toEqual({
      withPrefix: '',
      body: 'select * from orders',
    });
  });

  it('splits a single leading CTE from the main query', () => {
    expect(hoistLeadingCte('WITH paid AS (SELECT * FROM orders) SELECT * FROM paid')).toEqual({
      withPrefix: 'WITH paid AS (SELECT * FROM orders) ',
      body: 'SELECT * FROM paid',
    });
  });

  it('splits multiple CTEs, recursive CTEs, column lists, and UNION bodies', () => {
    expect(
      hoistLeadingCte(
        'WITH RECURSIVE nodes(id, parent_id) AS (SELECT id, parent_id FROM roots UNION ALL SELECT child.id, child.parent_id FROM child JOIN nodes ON child.parent_id = nodes.id), totals AS (SELECT count(*) AS total FROM nodes) SELECT total FROM totals',
      ),
    ).toEqual({
      withPrefix:
        'WITH RECURSIVE nodes(id, parent_id) AS (SELECT id, parent_id FROM roots UNION ALL SELECT child.id, child.parent_id FROM child JOIN nodes ON child.parent_id = nodes.id), totals AS (SELECT count(*) AS total FROM nodes) ',
      body: 'SELECT total FROM totals',
    });
  });

  it('ignores WITH, commas, and closing parens inside literals, identifiers, and comments', () => {
    expect(
      hoistLeadingCte(
        `WITH tricky AS (
          SELECT 'WITH x AS (SELECT 1), nope' AS "value, )"
          FROM orders
          WHERE note = ')'
          /* comment ), next AS (SELECT 1) */
        ) SELECT * FROM tricky`,
      ),
    ).toEqual({
      withPrefix: `WITH tricky AS (
          SELECT 'WITH x AS (SELECT 1), nope' AS "value, )"
          FROM orders
          WHERE note = ')'
          /* comment ), next AS (SELECT 1) */
        ) `,
      body: 'SELECT * FROM tricky',
    });
  });

  it('falls back to the legacy whole-query body when the CTE is malformed', () => {
    const malformed = 'WITH broken AS (SELECT * FROM orders SELECT * FROM broken';
    expect(hoistLeadingCte(malformed)).toEqual({ withPrefix: '', body: malformed });
  });
});

describe('limitSqlForExecution', () => {
  it('wraps compiled SQL and strips trailing semicolons', () => {
    expect(limitSqlForExecution('select * from public.orders; ', 25)).toBe(
      'select * from (select * from public.orders) as ktx_query_result limit 25',
    );
  });

  it('returns the trimmed SQL when no maxRows value is provided', () => {
    expect(limitSqlForExecution('select * from orders; ', undefined)).toBe('select * from orders');
  });

  it('strips leading comments before wrapping with a row limit', () => {
    expect(limitSqlForExecution('-- top customers\nselect * from public.orders', 25)).toBe(
      'select * from (select * from public.orders) as ktx_query_result limit 25',
    );
  });

  it('hoists leading CTEs before applying the generic LIMIT wrapper', () => {
    expect(limitSqlForExecution('WITH paid AS (SELECT * FROM orders) SELECT * FROM paid', 25)).toBe(
      'WITH paid AS (SELECT * FROM orders) select * from (SELECT * FROM paid) as ktx_query_result limit 25',
    );
  });

  it('keeps the generic wrapper byte-identical for non-CTE SQL', () => {
    expect(limitSqlForExecution('select id, status from public.orders', 25)).toBe(
      'select * from (select id, status from public.orders) as ktx_query_result limit 25',
    );
  });

  it('drops a trailing semicolon followed by a comment so the subquery stays valid', () => {
    // The single-statement gate accepts `select 1; -- done`; without stripping
    // the terminator the wrapper would embed `select 1; -- done` and comment out
    // the closing paren and limit clause.
    expect(limitSqlForExecution('select 1; -- done', 5)).toBe(
      'select * from (select 1) as ktx_query_result limit 5',
    );
    expect(limitSqlForExecution('select 1; /* note */', 5)).toBe(
      'select * from (select 1) as ktx_query_result limit 5',
    );
  });

  it('drops a trailing line comment with no semicolon before wrapping', () => {
    expect(limitSqlForExecution('select 1 -- done', 5)).toBe('select * from (select 1) as ktx_query_result limit 5');
  });
});

describe('stripTrailingSqlNoise', () => {
  it('removes trailing semicolons, comments, and whitespace', () => {
    expect(stripTrailingSqlNoise('select 1;')).toBe('select 1');
    expect(stripTrailingSqlNoise('select 1 ;; ')).toBe('select 1');
    expect(stripTrailingSqlNoise('select 1; -- done')).toBe('select 1');
    expect(stripTrailingSqlNoise('select 1 -- done')).toBe('select 1');
    expect(stripTrailingSqlNoise('select 1; /* trailing */')).toBe('select 1');
  });

  it('preserves semicolons and comment markers inside literals and mid-statement', () => {
    expect(stripTrailingSqlNoise("select 'a; -- b'")).toBe("select 'a; -- b'");
    expect(stripTrailingSqlNoise('select 1 /* a;b */ + 2')).toBe('select 1 /* a;b */ + 2');
    expect(stripTrailingSqlNoise('select ";" from "t;u"')).toBe('select ";" from "t;u"');
  });
});