supabase__supabase/apps/studio/components/ui/CodeEditor/Providers/PgSQLCompletionProvider.test.ts

import type { Monaco } from '@monaco-editor/react'
import type { editor, languages } from 'monaco-editor'
import { describe, expect, it } from 'vitest'

import { getPgsqlCompletionProvider } from './PgSQLCompletionProvider'
import type { PgInfo } from './Providers.types'

// A minimal stand-in for the parts of `Monaco` this provider actually reads.
const monaco = {
  languages: {
    CompletionItemKind: {
      Keyword: 1,
      Class: 2,
      Interface: 3,
      Field: 4,
      Function: 5,
      Property: 6,
    },
  },
} as unknown as Monaco

function createModel(sql: string): editor.ITextModel {
  return {
    getValue: () => sql,
    getOffsetAt: () => sql.length,
    getLineContent: () => sql,
    getWordUntilPosition: () => ({
      word: '',
      startColumn: sql.length + 1,
      endColumn: sql.length + 1,
    }),
  } as unknown as editor.ITextModel
}

// Simulates the cursor sitting mid-identifier between an already-present pair of double quotes,
// e.g. `where "OrderD|"` (the closing quote auto-closed by the editor when `"` was typed).
function createQuotedIdentModel(line: string, word: string): editor.ITextModel {
  const wordStartColumn = line.indexOf(`"${word}`) + 2
  const wordEndColumn = wordStartColumn + word.length

  return {
    getValue: () => line,
    getOffsetAt: () => wordEndColumn - 1,
    getLineContent: () => line,
    getWordUntilPosition: () => ({
      word,
      startColumn: wordStartColumn,
      endColumn: wordEndColumn,
    }),
  } as unknown as editor.ITextModel
}

function createQuotedIdentPosition(line: string, word: string): ProvideCompletionItemsParams[1] {
  const wordEndColumn = line.indexOf(`"${word}`) + 2 + word.length
  return { column: wordEndColumn, lineNumber: 1 } as unknown as ProvideCompletionItemsParams[1]
}

type ProvideCompletionItemsParams = Parameters<
  languages.CompletionItemProvider['provideCompletionItems']
>

function createPosition(sql: string): ProvideCompletionItemsParams[1] {
  // Monaco columns are 1-indexed; column = length + 1 places the cursor at the end.
  return { column: sql.length + 1, lineNumber: 1 } as unknown as ProvideCompletionItemsParams[1]
}

// Minimal fixtures — only the fields this provider reads, cast to the real (much larger) zod types.
function createPgInfoRef(): { current: PgInfo } {
  return {
    current: {
      keywords: ['abort', 'absent', 'absolute'],
      schemas: [{ name: 'public' }] as unknown as PgInfo['schemas'],
      functions: [
        { name: '_crypto_aead_det_decrypt', return_type: 'bytea' },
      ] as unknown as PgInfo['functions'],
      tableColumns: [
        {
          schemaname: 'public',
          tablename: 'colors',
          quoted_name: 'colors',
          is_table: true,
          columns: [
            { attname: 'id', data_type: 'bigint' },
            { attname: 'hex', data_type: 'text' },
          ],
        },
        {
          schemaname: 'public',
          tablename: 'shapes',
          quoted_name: 'shapes',
          is_table: true,
          columns: [{ attname: 'sides', data_type: 'smallint' }],
        },
      ],
    },
  }
}

function getSuggestions(pgInfoRef: { current: PgInfo }, sql: string): languages.CompletionItem[] {
  return provideSuggestions(pgInfoRef, sql, ' ')
}

function getDotSuggestions(
  pgInfoRef: { current: PgInfo },
  sql: string
): languages.CompletionItem[] {
  return provideSuggestions(pgInfoRef, sql, '.')
}

function provideSuggestions(
  pgInfoRef: { current: PgInfo },
  sql: string,
  triggerCharacter: string
): languages.CompletionItem[] {
  const provider = getPgsqlCompletionProvider(monaco, pgInfoRef)
  const context = { triggerCharacter } as unknown as ProvideCompletionItemsParams[2]
  const token = {} as ProvideCompletionItemsParams[3]

  const result = provider.provideCompletionItems(
    createModel(sql),
    createPosition(sql),
    context,
    token
  ) as languages.CompletionList

  return result.suggestions
}

describe('getPgsqlCompletionProvider - default scenario', () => {
  it('suggests only the FROM-clause table columns, ranked above keywords/functions', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getSuggestions(pgInfoRef, 'select * from colors where ')

    const columnSuggestions = suggestions.filter(
      (s) => s.kind === monaco.languages.CompletionItemKind.Field
    )
    expect(columnSuggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])

    // Every in-scope column must be ranked (sortText) ahead of keywords/functions
    columnSuggestions.forEach((s) => expect(s.sortText).toMatch(/^0_/))
    const nonColumnSuggestions = suggestions.filter(
      (s) => s.kind !== monaco.languages.CompletionItemKind.Field
    )
    nonColumnSuggestions.forEach((s) => expect(s.sortText).toBeUndefined())
  })

  it('falls back to every table column when no FROM clause is resolved yet', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getSuggestions(pgInfoRef, 'select ')

    const columnSuggestions = suggestions.filter(
      (s) => s.kind === monaco.languages.CompletionItemKind.Field
    )
    expect(columnSuggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id', 'sides'])
    columnSuggestions.forEach((s) => expect(s.sortText).toBeUndefined())
  })

  it('narrows columns per-table across a join', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getSuggestions(
      pgInfoRef,
      'select * from colors c join shapes s on s.id = c.id where '
    )

    const columnSuggestions = suggestions.filter(
      (s) => s.kind === monaco.languages.CompletionItemKind.Field
    )
    expect(columnSuggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id', 'sides'])
  })

  it('scopes to a single alias when the statement ends with `alias.`', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getSuggestions(
      pgInfoRef,
      'select * from colors c join shapes s on s.id = c.id where c.'
    )

    const columnSuggestions = suggestions.filter(
      (s) => s.kind === monaco.languages.CompletionItemKind.Field
    )
    expect(columnSuggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])
  })
})

describe('getPgsqlCompletionProvider - dot scenario', () => {
  it('scopes to a table alias in a single-table query', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getDotSuggestions(pgInfoRef, 'select * from colors c where c.')
    expect(suggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])
  })

  it('scopes to the aliased table only, not every table joined in', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getDotSuggestions(
      pgInfoRef,
      'select * from colors c join shapes s on s.id = c.id where c.'
    )
    expect(suggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])
  })

  it('still resolves a plain (un-aliased) table name', () => {
    const pgInfoRef = createPgInfoRef()
    const suggestions = getDotSuggestions(
      pgInfoRef,
      'select * from colors join shapes s on s.id = colors.id where colors.'
    )
    expect(suggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])
  })

  it('resolves an alias when there is no `public` schema at all', () => {
    const pgInfoRef: { current: PgInfo } = {
      current: {
        keywords: [],
        schemas: [{ name: 'app' }] as unknown as PgInfo['schemas'],
        functions: [] as unknown as PgInfo['functions'],
        tableColumns: [
          {
            schemaname: 'app',
            tablename: 'colors',
            quoted_name: 'colors',
            is_table: true,
            columns: [
              { attname: 'id', data_type: 'bigint' },
              { attname: 'hex', data_type: 'text' },
            ],
          },
        ],
      },
    }
    const suggestions = getDotSuggestions(pgInfoRef, 'select * from app.colors c where c.')
    expect(suggestions.map((s) => s.label).sort()).toStrictEqual(['hex', 'id'])
  })
})

describe('getPgsqlCompletionProvider - quoted identifiers', () => {
  function createQuotedPgInfoRef(): { current: PgInfo } {
    return {
      current: {
        keywords: [],
        schemas: [{ name: 'public' }] as unknown as PgInfo['schemas'],
        functions: [] as unknown as PgInfo['functions'],
        tableColumns: [
          {
            schemaname: 'public',
            tablename: 'test_orders',
            quoted_name: 'test_orders',
            is_table: true,
            columns: [
              { attname: 'OrderDate', data_type: 'date' },
              { attname: 'id', data_type: 'bigint' },
            ],
          },
        ],
      },
    }
  }

  it('replaces the whole existing quote pair instead of doubling up quotes on a mixed-case column', () => {
    // Regression test for: typing `where "OrderD` with the editor auto-closing the quote to
    // `where "OrderD|"`, then accepting the `OrderDate` suggestion produced `""OrderDate""`.
    const pgInfoRef = createQuotedPgInfoRef()
    const line = 'select * from test_orders where "OrderDate"'
    const word = 'OrderDate'

    const provider = getPgsqlCompletionProvider(monaco, pgInfoRef)
    const context = { triggerCharacter: undefined } as unknown as ProvideCompletionItemsParams[2]
    const token = {} as ProvideCompletionItemsParams[3]

    const result = provider.provideCompletionItems(
      createQuotedIdentModel(line, word),
      createQuotedIdentPosition(line, word),
      context,
      token
    ) as languages.CompletionList

    const suggestion = result.suggestions.find((s) => s.label === 'OrderDate')
    expect(suggestion?.insertText).toBe('"OrderDate"')

    // The range must swallow both surrounding quotes, otherwise the quoted insertText lands
    // inside the pre-existing, untouched quote pair and doubles them up.
    const quoteStart = line.indexOf('"')
    const quoteEnd = line.indexOf('"', quoteStart + 1)
    const range = suggestion?.range as { startColumn: number; endColumn: number }
    expect(range.startColumn).toBe(quoteStart + 1)
    expect(range.endColumn).toBe(quoteEnd + 2)
  })

  it('sets filterText with the leading quote so Monaco does not filter out every suggestion while typing inside an auto-closed quote pair', () => {
    // Regression test for: with the range widened to swallow the opening quote (above), Monaco
    // fuzzy-matches the typed prefix (e.g. `"OrderD`) against each item's filterText/label. A plain
    // label like `OrderDate` has no leading quote to match, so every suggestion got filtered out —
    // the widget showed no suggestions at all, even on an explicit Ctrl+Space re-invoke.
    const pgInfoRef = createQuotedPgInfoRef()
    // Only `OrderD` has been typed so far; the closing quote is auto-closed immediately after it.
    const line = 'select * from test_orders where "OrderD"'
    const word = 'OrderD'

    const provider = getPgsqlCompletionProvider(monaco, pgInfoRef)
    const context = { triggerCharacter: undefined } as unknown as ProvideCompletionItemsParams[2]
    const token = {} as ProvideCompletionItemsParams[3]

    const result = provider.provideCompletionItems(
      createQuotedIdentModel(line, word),
      createQuotedIdentPosition(line, word),
      context,
      token
    ) as languages.CompletionList

    const suggestion = result.suggestions.find((s) => s.label === 'OrderDate')
    expect(suggestion?.filterText).toBe('"OrderDate')
  })

  it('replaces only the pre-existing opening quote instead of doubling it up when the closing quote has been deleted', () => {
    // Regression test for: type `"OrderD` (auto-closed to `"OrderD|"`), delete the closing quote
    // (leaving `"OrderD` with no closing quote at all), then accept the `OrderDate` suggestion —
    // previously produced `""OrderDate` because `isQuoted` required both quotes to be present, so
    // the range left the existing opening quote untouched while the quoted insertText added another.
    const pgInfoRef = createQuotedPgInfoRef()
    const line = 'select * from test_orders where "OrderD'
    const word = 'OrderD'

    const provider = getPgsqlCompletionProvider(monaco, pgInfoRef)
    const context = { triggerCharacter: undefined } as unknown as ProvideCompletionItemsParams[2]
    const token = {} as ProvideCompletionItemsParams[3]

    const result = provider.provideCompletionItems(
      createQuotedIdentModel(line, word),
      createQuotedIdentPosition(line, word),
      context,
      token
    ) as languages.CompletionList

    const suggestion = result.suggestions.find((s) => s.label === 'OrderDate')
    expect(suggestion?.insertText).toBe('"OrderDate"')
    expect(suggestion?.filterText).toBe('"OrderDate')

    // Only the opening quote gets swallowed — there's no closing quote to swallow, so the end
    // boundary stays at the word's own end (the cursor position).
    const quoteStart = line.indexOf('"')
    const range = suggestion?.range as { startColumn: number; endColumn: number }
    expect(range.startColumn).toBe(quoteStart + 1)
    expect(range.endColumn).toBe(line.length + 1)
  })

  it('still quotes an all-lowercase column when the closing quote is missing', () => {
    // formatInsertText only force-quotes on `isQuoted`, since a mixed-case column already forces
    // quoting via its own `hasUpperCase` check regardless of `isQuoted`. An all-lowercase column
    // has no such fallback, so this exercises `isQuoted` being detected from the opening quote alone.
    const pgInfoRef = createQuotedPgInfoRef()
    pgInfoRef.current.tableColumns[0].columns.push({ attname: 'orderdate', data_type: 'date' })
    const line = 'select * from test_orders where "orderd'
    const word = 'orderd'

    const provider = getPgsqlCompletionProvider(monaco, pgInfoRef)
    const context = { triggerCharacter: undefined } as unknown as ProvideCompletionItemsParams[2]
    const token = {} as ProvideCompletionItemsParams[3]

    const result = provider.provideCompletionItems(
      createQuotedIdentModel(line, word),
      createQuotedIdentPosition(line, word),
      context,
      token
    ) as languages.CompletionList

    const suggestion = result.suggestions.find((s) => s.label === 'orderdate')
    expect(suggestion?.insertText).toBe('"orderdate"')
  })
})