Drizzle `withPagination`

paginate.ts

import invariant from 'tiny-invariant'
import type { PgColumn, PgSelect } from 'drizzle-orm/pg-core'
import type { SQL } from 'drizzle-orm'
import { desc, gt, lt, eq, and, or } from 'drizzle-orm'

// using `asc` or `desc` outside this function obscures the
// original column name, so they must be applied internally
type Order = 'desc' | 'asc'
type OrderColumn = PgColumn | [PgColumn, Order]
type Options = {
  pageSize: number
  cursor?: string
  orderColumns: [OrderColumn, ...OrderColumn[]]
  conds?: SQL
}
type Column = {
  drizzle: PgColumn
  order: Order
  fieldName: string
}

export async function withPagination<T extends PgSelect>(
  qb: T,
  opts: Options
): Promise<{ entries: Awaited<T>; next?: string; prev?: string }> {
  let reverse = false
  let { cursor } = opts
  if (cursor?.[0] === '^') {
    reverse = true
    cursor = cursor.slice(1)
  }

  const cols = opts.orderColumns.map(col => {
    const oc: OrderColumn = Array.isArray(col) ? col : [col, 'asc']
    const [drizzle, _order] = oc
    let order = _order
    if (reverse) {
      order = order === 'asc' ? 'desc' : 'asc'
    }
    // TODO: ask for help with mapping column name back to field name
    const fieldName = Object.entries(qb._.selectedFields).find(
      e => e[1] === drizzle
    )?.[0]
    invariant(fieldName, `Pagination unable to map column '${drizzle.name}'`)
    return { drizzle, order, fieldName }
  })

  const viewConds = opts.conds ? [opts.conds] : []
  if (cursor) {
    viewConds.push(buildCursorConds(cols, JSON.parse(atob(cursor))))
  }

  const { pageSize } = opts
  const order = cols.map(col =>
    col.order === 'asc' ? col.drizzle : desc(col.drizzle)
  )

  const rows = await qb
    .where(and(...viewConds))
    .limit(pageSize + 1)
    .orderBy(...order)

  function buildCursor(row: Awaited<T>[number], modifier: string) {
    const values = cols.map(spec => row[spec.fieldName])
    return modifier + btoa(JSON.stringify(values))
  }

  const entries = rows.slice(0, pageSize) as Awaited<T>
  let startCursor: string | undefined
  if (cursor) {
    startCursor = buildCursor(rows[0], reverse ? '' : '^')
  }
  let endCursor: string | undefined
  if (rows.length > pageSize) {
    endCursor = buildCursor(rows[pageSize - 1], reverse ? '^' : '')
  }

  return reverse
    ? {
        entries: entries.reverse() as Awaited<T>,
        prev: endCursor,
        next: startCursor,
      }
    : {
        entries,
        prev: startCursor,
        next: endCursor,
      }
}

// build a complex condition to support multiple sorted columns
// e.g. (col1 > value1) OR (col1 = value1 AND col2 > value2)
function buildCursorConds(cols: Column[], cursor: any[]) {
  const pageConds = cols.map(col => {
    // for each order column, build a condition using the subset of columns
    // from the left-most to the current (inclusive)
    const subset = cols.slice(0, cols.indexOf(col) + 1)
    const subconds = subset.map((subcol, i) => {
      const value = cursor[i]
      if (subcol !== col) {
        // we haven't reached the last column in the subset - use equality
        return eq(subcol.drizzle, value)
      }
      // the final column - use a ranged operator
      const op = subcol.order === 'asc' ? gt : lt
      return op(subcol.drizzle, value)
    })
    return and(...subconds)
  })
  return or(...pageConds) as SQL
}