supabase/apps/studio/components/interfaces/Settings/Logs/logs-sql-rewrite.test.ts

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

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

describe('buildClickhouseRewritePrompt', () => {
  it('includes the query, the schema, and a reply-with-only-SQL instruction', () => {
    const prompt = buildClickhouseRewritePrompt('select count(*) from edge_logs')
    expect(prompt).toContain('select count(*) from edge_logs')
    expect(prompt).toContain('log_attributes')
    expect(prompt).toContain("source = 'edge_logs'")
    expect(prompt.toLowerCase()).toContain('reply with only')
  })

  it('spells out the FROM-to-logs conversion and shows a worked example', () => {
    const prompt = buildClickhouseRewritePrompt('select 1 from postgres_logs')
    expect(prompt).toContain("from logs where source = 'postgres_logs'")
    expect(prompt.toLowerCase()).toContain('remove every')
    expect(prompt).toContain('cross join unnest')
    expect(prompt).toContain('BigQuery:')
    expect(prompt).toContain('ClickHouse:')
    expect(prompt).toContain("log_attributes['parsed.error_severity']")
  })

  it('lists the real log_attributes keys when provided and demands exact paths', () => {
    const prompt = buildClickhouseRewritePrompt('select 1 from edge_logs', [
      'request.headers.x_real_ip',
      'request.cf.country',
    ])
    expect(prompt).toContain("log_attributes['request.headers.x_real_ip']")
    expect(prompt).toContain("log_attributes['request.cf.country']")
    expect(prompt.toLowerCase()).toContain('exact')
  })

  it('omits the keys section when none are provided', () => {
    const prompt = buildClickhouseRewritePrompt('select 1 from edge_logs')
    expect(prompt).not.toContain('actual log_attributes keys present')
  })
})

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

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('posts to the completion endpoint and returns the cleaned query', 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',
    })

    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.completionMetadata.selection).toBe('select 1 from edge_logs')
    expect(body.completionMetadata.prompt.toLowerCase()).toContain('reply with only')
  })

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