supabase/packages/pg-meta/src/query/Query.utils.ts

import { format, ident, literal } from '../pg-format'
import type { Dictionary, Filter, QueryPagination, QueryTable, Sort } from './types'

export function countQuery(
  table: QueryTable,
  options?: {
    filters?: Filter[]
  }
) {
  let query = `select count(*) from ${queryTable(table)}`
  const { filters } = options ?? {}
  if (filters) {
    query = applyFilters(query, filters)
  }
  return query + ';'
}

export function truncateQuery(
  table: QueryTable,
  options?: {
    // [Joshen] yet to implement cascade from UI, just adding first
    cascade?: boolean
  }
) {
  let query = `truncate ${queryTable(table)}`
  const { cascade } = options ?? {}
  if (cascade) {
    query += ' cascade'
  }
  return query + ';'
}

export function deleteQuery(
  table: QueryTable,
  filters?: Filter[],
  options?: {
    returning?: boolean
    enumArrayColumns?: string[]
  }
) {
  if (!filters || filters.length === 0) {
    throw new Error('no filters for this delete query')
  }
  let query = `delete from ${queryTable(table)}`
  const { returning, enumArrayColumns } = options ?? {}
  if (filters) {
    query = applyFilters(query, filters)
  }
  if (returning) {
    query +=
      enumArrayColumns === undefined || enumArrayColumns.length === 0
        ? ` returning *`
        : ` returning *, ${enumArrayColumns.map((x) => `${ident(x)}::text[]`).join(',')}`
  }
  return query + ';'
}

export function insertQuery(
  table: QueryTable,
  values: Dictionary<any>[],
  options?: {
    returning?: boolean
    enumArrayColumns?: string[]
  }
) {
  if (!values || values.length === 0) {
    throw new Error('no value to insert')
  }
  const { returning, enumArrayColumns } = options ?? {}
  const queryColumns = Object.keys(values[0])
    .map((x) => ident(x))
    .join(',')
  let query = ''
  if (queryColumns.length == 0) {
    query = format(
      'insert into %1$s select from jsonb_populate_recordset(null::%1$s, %2$s)',
      queryTable(table),
      literal(JSON.stringify(values))
    )
  } else {
    query = format(
      'insert into %1$s (%2$s) select %2$s from jsonb_populate_recordset(null::%1$s, %3$s)',
      queryTable(table),
      queryColumns,
      literal(JSON.stringify(values))
    )
  }
  if (returning) {
    query +=
      enumArrayColumns === undefined || enumArrayColumns.length === 0
        ? ` returning *`
        : ` returning *, ${enumArrayColumns.map((x) => `${ident(x)}::text[]`).join(',')}`
  }
  return query + ';'
}

export function selectQuery(
  table: QueryTable,
  columns?: string,
  options?: {
    filters?: Filter[]
    pagination?: QueryPagination
    sorts?: Sort[]
  },
  isFinal = true,
  isCTE = false
) {
  let query = ''
  const queryColumn = columns ?? '*'
  query += `select ${queryColumn} from ${isCTE ? queryCTE(table) : queryTable(table)}`

  const { filters, pagination, sorts } = options ?? {}
  if (filters) {
    query = applyFilters(query, filters)
  }
  if (sorts) {
    query = applySorts(query, sorts)
  }
  if (pagination) {
    const { limit, offset } = pagination ?? {}
    query += ` limit ${literal(limit)} offset ${literal(offset)}`
  }
  return `${query}${isFinal ? ';' : ''}`
}

export function updateQuery(
  table: QueryTable,
  value: Dictionary<any>,
  options?: {
    filters?: Filter[]
    returning?: boolean
    enumArrayColumns?: string[]
  }
) {
  const { filters, returning, enumArrayColumns } = options ?? {}
  if (!filters || filters.length === 0) {
    throw new Error('no filters for this update query')
  }
  const queryColumns = Object.keys(value)
    .map((x) => ident(x))
    .join(',')
  let query = format(
    'update %1$s set (%2$s) = (select %2$s from json_populate_record(null::%1$s, %3$s))',
    queryTable(table),
    queryColumns,
    literal(JSON.stringify(value))
  )
  if (filters) {
    query = applyFilters(query, filters)
  }
  if (returning) {
    query +=
      enumArrayColumns === undefined || enumArrayColumns.length === 0
        ? ` returning *`
        : ` returning *, ${enumArrayColumns.map((x) => `${ident(x)}::text[]`).join(',')}`
  }

  return query + ';'
}

//============================================================
// Filter Utils
//============================================================

function applyFilters(query: string, filters: Filter[]) {
  if (filters.length === 0) return query
  query += ` where ${filters
    .map((filter) => {
      // Handle composite values
      if (Array.isArray(filter.column)) {
        switch (filter.operator) {
          case 'in':
            return inTupleFilterSql(filter)
          case '=':
          case '<>':
          case '>':
          case '<':
          case '>=':
          case '<=':
            return defaultTupleFilterSql(filter)
          default:
            throw new Error(`Cannot use ${filter.operator} operator in a tuple filter`)
        }
      }

      switch (filter.operator) {
        case 'in':
          return inFilterSql(filter)
        case 'is':
          return isFilterSql(filter)
        case '~~':
        case '~~*':
        case '!~~':
        case '!~~*':
          return castColumnToText(filter)
        default:
          return `${ident(filter.column)} ${filter.operator} ${filterLiteral(filter.value)}`
      }
    })
    .join(' and ')}`
  return query
}

function inFilterSql(filter: Filter) {
  let values: Array<unknown>
  if (Array.isArray(filter.value)) {
    values = filter.value.map((x) => filterLiteral(x))
  } else {
    const filterValueTxt = String(filter.value)
    values = filterValueTxt.split(',').map((x) => filterLiteral(x))
  }
  return `${ident(filter.column)} ${filter.operator} (${values.join(',')})`
}

function defaultTupleFilterSql(filter: Filter) {
  if (!Array.isArray(filter.column)) {
    throw new Error('Use standard applyFilters for single column')
  }
  if (!Array.isArray(filter.value)) {
    throw new Error('Tuple filter value must be an array')
  }
  if (filter.value.length !== filter.column.length) {
    throw new Error('Tuple filter value must have the same length as the column array')
  }

  const columns = `(${filter.column.map((c) => ident(c)).join(', ')})`
  const values = `(${filter.value.map((v) => filterLiteral(v)).join(', ')})`
  return `${columns} ${filter.operator} ${values}`
}

function inTupleFilterSql(filter: Filter) {
  if (!Array.isArray(filter.column)) {
    throw new Error('Use inFilterSql for single columns')
  }
  if (!Array.isArray(filter.value)) {
    throw new Error(`Values for a tuple 'in' filter must be an array`)
  }

  const columns = `(${filter.column.map((c) => ident(c)).join(', ')})`

  const values = filter.value.map((v) => {
    if (Array.isArray(v)) {
      if (v.length !== filter.column.length) {
        throw new Error(`Tuple value length must match column length`)
      }
      return `(${v.map((x) => filterLiteral(x)).join(', ')})`
    } else {
      const filterValueTxt = String(v)
      const currValues = filterValueTxt.split(',')
      if (currValues.length !== filter.column.length) {
        throw new Error(`Tuple value length must match column length`)
      }
      return `(${currValues.map((x) => filterLiteral(x)).join(', ')})`
    }
  })

  return `${columns} ${filter.operator} (${values.join(', ')})`
}

function isFilterSql(filter: Filter) {
  const filterValueTxt = String(filter.value)
  switch (filterValueTxt) {
    case 'null':
    case 'false':
    case 'true':
    case 'not null':
      return `${ident(filter.column)} ${filter.operator} ${filterValueTxt}`
    default:
      return `${ident(filter.column)} ${filter.operator} ${filterLiteral(filter.value)}`
  }
}

function castColumnToText(filter: Filter) {
  return `${ident(filter.column)}::text ${filter.operator} ${filterLiteral(filter.value)}`
}

function filterLiteral(value: any) {
  if (typeof value === 'string') {
    if (value?.startsWith('ARRAY[') && value?.endsWith(']')) {
      return value
    } else {
      return literal(value)
    }
  }
  return value
}

//============================================================
// Sort Utils
//============================================================

function applySorts(query: string, sorts: Sort[]) {
  const validSorts = sorts.filter((sort) => sort.column)
  if (validSorts.length === 0) return query
  query += ` order by ${validSorts
    .map((x) => {
      const order = x.ascending ? 'asc' : 'desc'
      const nullOrder = x.nullsFirst ? 'nulls first' : 'nulls last'
      return `${ident(x.table)}.${ident(x.column)} ${order} ${nullOrder}`
    })
    .join(', ')}`
  return query
}

//============================================================
// Misc
//============================================================

function queryTable(table: QueryTable) {
  return `${ident(table.schema)}.${ident(table.name)}`
}

function queryCTE(table: QueryTable) {
  return `${ident(table.name)}`
}