supabase/apps/studio/data/logs/logs-sql-rewrite.test.ts

import { afterEach, describe, expect, it, vi } from 'vitest'

import {
  detectLogSource,
  looksLikeLegacyLogsQuery,
  rewriteLogsSqlWithAI,
  shouldOfferLegacyLogsRewrite,
  stripSqlCodeFences,
} from './logs-sql-rewrite'

describe('shouldOfferLegacyLogsRewrite', () => {
  const legacySql = 'select 1 from edge_logs cross join unnest(metadata) as m'

  it('offers the rewrite for BigQuery-dialect SQL once logs run on ClickHouse', () => {
    expect(shouldOfferLegacyLogsRewrite({ sql: legacySql, isClickhouseLogsEnabled: true })).toBe(
      true
    )
  })

  it('never offers it on a non-migrated org, where the BigQuery SQL is still correct', () => {
    expect(shouldOfferLegacyLogsRewrite({ sql: legacySql, isClickhouseLogsEnabled: false })).toBe(
      false
    )
  })

  it('does not offer it for SQL that is already ClickHouse, or for empty SQL', () => {
    expect(
      shouldOfferLegacyLogsRewrite({
        sql: "select timestamp from logs where source = 'edge_logs' limit 5",
        isClickhouseLogsEnabled: true,
      })
    ).toBe(false)
    expect(shouldOfferLegacyLogsRewrite({ sql: '', isClickhouseLogsEnabled: true })).toBe(false)
  })
})

describe('detectLogSource', () => {
  it('reads an explicit source filter', () => {
    expect(detectLogSource("select 1 from logs where source = 'edge_logs'")).toBe('edge_logs')
  })

  it('falls back to the legacy FROM table name', () => {
    expect(detectLogSource('select 1 from edge_logs as t')).toBe('edge_logs')
  })

  it('maps pg_cron_logs to postgres_logs from either the FROM table or the source filter', () => {
    expect(detectLogSource('select 1 from pg_cron_logs')).toBe('postgres_logs')
    expect(detectLogSource("select 1 from logs where source = 'pg_cron_logs'")).toBe(
      'postgres_logs'
    )
  })

  it('returns undefined for the bare logs table with no source', () => {
    expect(detectLogSource('select 1 from logs limit 5')).toBeUndefined()
  })

  it('returns undefined when nothing matches', () => {
    expect(detectLogSource('select 1')).toBeUndefined()
  })

  it('ignores a column that merely ends in "source"', () => {
    expect(detectLogSource("select 1 from logs where resource = 'nope'")).toBeUndefined()
    expect(detectLogSource("select 1 from logs where datasource = 'nope'")).toBeUndefined()
  })

  it('still reads a qualified source column', () => {
    expect(detectLogSource("select 1 from logs t where t.source = 'auth_logs'")).toBe('auth_logs')
  })

  it('prefers the real source column over a lookalike earlier in the query', () => {
    expect(detectLogSource("select resource = 'nope' from logs where source = 'edge_logs'")).toBe(
      'edge_logs'
    )
  })
})

describe('looksLikeLegacyLogsQuery', () => {
  it('flags per-service FROM tables', () => {
    expect(looksLikeLegacyLogsQuery('select 1 from edge_logs')).toBe(true)
  })

  it('flags unnest joins and the cast-timestamp idiom', () => {
    expect(looksLikeLegacyLogsQuery('select 1 from logs cross join unnest(metadata) as m')).toBe(
      true
    )
    expect(looksLikeLegacyLogsQuery('select cast(timestamp as datetime) from logs')).toBe(true)
  })

  it('does not flag a ClickHouse query against the logs table', () => {
    expect(
      looksLikeLegacyLogsQuery("select timestamp from logs where source = 'edge_logs' limit 5")
    ).toBe(false)
  })
})

describe('stripSqlCodeFences', () => {
  it('removes a ```sql fenced block', () => {
    expect(stripSqlCodeFences('```sql\nselect 1 from logs\n```')).toBe('select 1 from logs')
  })

  it('removes a plain ``` fenced block', () => {
    expect(stripSqlCodeFences('```\nselect 1\n```')).toBe('select 1')
  })

  it('leaves unfenced SQL untouched (trimmed)', () => {
    expect(stripSqlCodeFences('  select 1 from logs  ')).toBe('select 1 from logs')
  })

  it('extracts the fenced block when wrapped in prose', () => {
    expect(stripSqlCodeFences('Here is the rewrite:\n```sql\nselect 1 from logs\n```')).toBe(
      'select 1 from logs'
    )
  })
})

describe('rewriteLogsSqlWithAI', () => {
  afterEach(() => {
    vi.unstubAllGlobals()
  })

  it('declares the rewrite intent and sends the query as the selection', async () => {
    const fetchMock = vi.fn().mockResolvedValue({
      ok: true,
      json: async () => '```sql\nselect 1 from logs\n```',
    })
    vi.stubGlobal('fetch', fetchMock)

    const result = await rewriteLogsSqlWithAI({
      sql: 'select 1 from edge_logs',
      projectRef: 'abc',
      availableKeys: ['request.method'],
    })

    expect(result).toBe('select 1 from logs')
    const [url, init] = fetchMock.mock.calls[0]
    expect(url).toContain('/api/ai/code/complete')
    const body = JSON.parse(init.body)
    expect(body.dialect).toBe('clickhouse')
    expect(body.intent).toBe('rewrite')
    // The whole query is the selection, so the rewrite replaces all of it.
    expect(body.completionMetadata.selection).toBe('select 1 from edge_logs')
    expect(body.completionMetadata.textBeforeCursor).toBe('')
    expect(body.completionMetadata.textAfterCursor).toBe('')
    expect(body.completionMetadata.availableKeys).toEqual(['request.method'])
    // No prompt text is carried client-side — the route owns the instruction.
    expect(body.completionMetadata.prompt).toBe('')
  })

  it('throws when the request fails', async () => {
    vi.stubGlobal('fetch', vi.fn().mockResolvedValue({ ok: false, text: async () => 'boom' }))
    await expect(rewriteLogsSqlWithAI({ sql: 'select 1', projectRef: 'abc' })).rejects.toThrow(
      'boom'
    )
  })

  it('throws when the model returns an empty query', async () => {
    vi.stubGlobal('fetch', vi.fn().mockResolvedValue({ ok: true, json: async () => '   ' }))
    await expect(rewriteLogsSqlWithAI({ sql: 'select 1', projectRef: 'abc' })).rejects.toThrow(
      'empty'
    )
  })
})