PotentialStorage/test/kv/indexes.test.ts

import { describe, expect, it } from 'vitest'
import fc from 'fast-check'
import { KVStorage, type KVStorageOptions } from '../../src/kv/node.js'
import { InMemoryKVDriver } from '../../src/kv/memory.js'
import type { KVDriver } from '../../src/kv/driver.js'
import { encodeIndexValue } from '../../src/core/encoding.js'
import { collect, CountingDriver, type Capability } from '../helpers.js'

type User = { email?: string; age?: number; tags?: string[] }
type Post = { createdAt: string; tags: string[] }

const indexes: KVStorageOptions['indexes'] = {
  users: {
    byEmail: (u: User) => u.email,
    byAge: (u: User) => u.age,
    byTag: (u: User) => u.tags,
  },
  'users/*/posts': { byDate: (p: Post) => p.createdAt, byTag: (p: Post) => p.tags },
}

function setup(hide: Capability[] = []) {
  const driver = new CountingDriver(new InMemoryKVDriver(), hide)
  return { driver, db: new KVStorage(driver, { indexes }) }
}

const modes: [string, Capability[]][] = [
  ['atomic batch', []],
  ['non-atomic', ['batch']],
]

describe.each(modes)('indexes (%s)', (_, hide) => {
  it('writes entries under <node>/~i/<name>/<enc>/<child>', async () => {
    const { db, driver } = setup(hide)
    await db.child('users').set('123', { email: 'sada@x.io', age: 30 })
    const keys = (await collect(driver.inner.scan({}))).map(([k]) => k)
    expect(keys).toEqual([`users/123`, `users/~i/byAge/${encodeIndexValue(30)}/123`, `users/~i/byEmail/${encodeIndexValue('sada@x.io')}/123`])
  })

  it('finds by exact value and keeps entries in sync on update and delete', async () => {
    const { db } = setup(hide)
    const users = db.child('users')
    await users.set('1', { email: 'a@x', age: 30 })
    await users.set('2', { email: 'b@x', age: 30 })
    expect(await collect(users.findBy('byEmail', 'a@x'))).toEqual([['1', { email: 'a@x', age: 30 }]])
    expect((await collect(users.findBy('byAge', 30))).map(([k]) => k)).toEqual(['1', '2'])
    await users.set('1', { email: 'c@x', age: 31 })
    expect(await collect(users.findBy('byEmail', 'a@x'))).toEqual([])
    expect((await collect(users.findBy('byEmail', 'c@x'))).map(([k]) => k)).toEqual(['1'])
    await users.delete('2')
    expect(await collect(users.findBy('byAge', 30))).toEqual([])
  })

  it('self-form writes on a child update the parent index', async () => {
    const { db } = setup(hide)
    const u = db.child('users').child('7')
    await u.set({ email: 'self@x' })
    expect((await collect(db.child('users').findBy('byEmail', 'self@x'))).map(([k]) => k)).toEqual(['7'])
    await u.delete()
    expect(await collect(db.child('users').findBy('byEmail', 'self@x'))).toEqual([])
  })

  it('deleteTree removes the child entry', async () => {
    const { db, driver } = setup(hide)
    await db.child('users').set('7', { email: 'e' })
    await db.child('users').child('7').child('posts').set('p', { createdAt: '2024', tags: [] })
    await db.child('users').deleteTree('7')
    expect(await collect(driver.inner.scan({}))).toEqual([])
  })

  it('supports multi-entry values, de-duplicated in range queries', async () => {
    const { db } = setup(hide)
    const users = db.child('users')
    await users.set('1', { tags: ['a', 'b', 'a'] })
    await users.set('2', { tags: ['b'] })
    expect((await collect(users.findBy('byTag', 'b'))).map(([k]) => k)).toEqual(['1', '2'])
    expect((await collect(users.findBy('byTag', { gte: 'a', lt: 'c' }))).map(([k]) => k)).toEqual(['1', '2'])
    await users.set('1', { tags: ['c'] })
    expect((await collect(users.findBy('byTag', 'b'))).map(([k]) => k)).toEqual(['2'])
  })

  it('range queries are ordered by value and honour bounds, limit and reverse', async () => {
    const { db } = setup(hide)
    const users = db.child('users')
    await users.setMany([
      ['a', { age: 10 }],
      ['b', { age: 9 }],
      ['c', { age: -1.5 }],
      ['d', { age: 100 }],
      ['e', { age: 20 }],
    ])
    const ages = async (q: object, o = {}) => (await collect(users.findBy('byAge', q as never, o))).map(([, v]) => (v as User).age)
    expect(await ages({})).toEqual([-1.5, 9, 10, 20, 100])
    expect(await ages({ gte: 9, lt: 20 })).toEqual([9, 10])
    expect(await ages({ gt: 9, lte: 20 })).toEqual([10, 20])
    expect(await ages({ gte: 0 }, { limit: 2 })).toEqual([9, 10])
    expect(await ages({ lt: 100 }, { reverse: true })).toEqual([20, 10, 9, -1.5])
  })

  it('applies pattern indexes with * to nested nodes only', async () => {
    const { db } = setup(hide)
    const posts = db.child('users').child('u1').child('posts')
    await posts.set('p1', { createdAt: '2024-02', tags: ['x'] })
    await posts.set('p2', { createdAt: '2024-01', tags: ['y'] })
    expect((await collect(posts.findBy('byDate', { gte: '2024' }))).map(([k]) => k)).toEqual(['p2', 'p1'])
    await expect(db.child('other').findBy('byDate', 'x').next()).rejects.toThrow(/No index/)
  })

  it('handles child keys and values needing escaping', async () => {
    const { db } = setup(hide)
    const users = db.child('users')
    await users.set('a/b%~', { email: 'we/ird%~@x' })
    expect(await collect(users.findBy('byEmail', 'we/ird%~@x'))).toEqual([['a/b%~', { email: 'we/ird%~@x' }]])
  })
})

describe('index costs', () => {
  it('non-indexed nodes do not read before writing', async () => {
    const { db, driver } = setup()
    driver.reset()
    await db.child('other').set('1', { email: 'x' })
    await db.child('other').setMany([['2', {}]])
    expect(driver.calls).toEqual({ setMany: 2 })
  })

  it('indexed writes: one getMany then one atomic batch', async () => {
    const { db, driver } = setup()
    driver.reset()
    await db.child('users').setMany([
      ['1', { email: 'a' }],
      ['2', { email: 'b' }],
    ])
    expect(driver.calls).toEqual({ getMany: 1, batch: 1 })
  })

  it('findBy with a small limit makes one value read', async () => {
    const { db, driver } = setup()
    await db.child('users').setMany(Array.from({ length: 300 }, (_, i) => [String(i), { age: 1 }] as [string, User]))
    driver.reset()
    expect(await collect(db.child('users').findBy('byAge', 1, { limit: 1 }))).toHaveLength(1)
    expect(driver.calls).toEqual({ scan: 1, getMany: 1 })
    driver.reset()
    expect(await collect(db.child('users').findBy('byAge', 1))).toHaveLength(300)
    expect(driver.calls).toEqual({ scan: 1, getMany: 3 })
  })
})

describe('non-atomic crash recovery (read-repair)', () => {
  /** Fails the n-th mutating call. */
  class CrashingDriver extends CountingDriver {
    failAt = Infinity
    #n = 0
    override setMany(entries: [string, unknown][]) {
      if (++this.#n === this.failAt) return Promise.reject(new Error('crash'))
      return super.setMany(entries)
    }
    override deleteMany(keys: string[]) {
      if (++this.#n === this.failAt) return Promise.reject(new Error('crash'))
      return super.deleteMany(keys)
    }
  }

  it('skips and deletes stale entries left by a crash after the index add', async () => {
    const driver = new CrashingDriver(new InMemoryKVDriver(), ['batch'])
    const users = new KVStorage(driver, { indexes }).child('users')
    await users.set('1', { email: 'old' })
    driver.failAt = 2 + 2 // adds succeed, value write crashes
    await expect(users.set('1', { email: 'new' })).rejects.toThrow('crash')
    // The 'new' entry is a dangling hint; the value still says 'old'.
    expect(await collect(users.findBy('byEmail', 'new'))).toEqual([])
    expect(await collect(users.findBy('byEmail', 'old'))).toEqual([['1', { email: 'old' }]])
    const entries = (await collect(driver.inner.scan({ gte: 'users/~i/' }))).map(([k]) => k)
    expect(entries).toEqual([`users/~i/byEmail/${encodeIndexValue('old')}/1`])
  })

  it('skips stale entries left by a crash before the index remove', async () => {
    const driver = new CrashingDriver(new InMemoryKVDriver(), ['batch'])
    const users = new KVStorage(driver, { indexes }).child('users')
    await users.set('1', { email: 'old' })
    driver.failAt = 2 + 3 // adds + value succeed, remove crashes
    await expect(users.set('1', { email: 'new' })).rejects.toThrow('crash')
    expect(await collect(users.findBy('byEmail', 'old'))).toEqual([])
    expect(await collect(users.findBy('byEmail', 'new'))).toEqual([['1', { email: 'new' }]])
  })

  it('can leave stale entries in place with readRepair: false', async () => {
    const driver = new InMemoryKVDriver()
    await driver.setMany([[`users/~i/byEmail/${encodeIndexValue('x')}/1`, true]])
    const users = new KVStorage(driver, { indexes, readRepair: false }).child('users')
    expect(await collect(users.findBy('byEmail', 'x'))).toEqual([])
    expect(await collect(driver.scan({}))).toHaveLength(1)
  })
})

describe('index declarations', () => {
  it('rejects the same index name on overlapping patterns', () => {
    expect(() => new KVStorage(new InMemoryKVDriver(), { indexes: { 'a/*': { i: () => 1 }, '*/b': { i: () => 2 } } })).toThrow(/overlapping/)
    expect(() => new KVStorage(new InMemoryKVDriver(), { indexes: { 'a/*': { i: () => 1 }, '*/b': { j: () => 2 } } })).not.toThrow()
  })

  it('rejects non-indexable index function results', async () => {
    const db = new KVStorage(new InMemoryKVDriver(), { indexes: { '': { bad: () => ({}) as never } } })
    await expect(db.set('x', 1)).rejects.toThrow(/non-indexable/)
  })
})

describe('property: findBy equals brute-force filter', () => {
  const val = fc.record({
    age: fc.option(fc.integer({ min: -3, max: 3 }), { nil: undefined }),
    tags: fc.array(fc.constantFrom('a', 'b', 'c/d'), { maxLength: 3 }),
  })
  const op = fc.oneof(
    fc.record({ t: fc.constant('set' as const), k: fc.constantFrom('1', '2', '3', 'x/y'), v: val }),
    fc.record({ t: fc.constant('del' as const), k: fc.constantFrom('1', '2', '3', 'x/y') }),
  )
  for (const [mode, hide] of modes) {
    it(mode, async () => {
      await fc.assert(
        fc.asyncProperty(fc.array(op, { maxLength: 20 }), fc.integer({ min: -4, max: 4 }), fc.integer({ min: -4, max: 4 }), async (ops, lo, hi) => {
          const driver: KVDriver = new CountingDriver(new InMemoryKVDriver(), hide)
          const users = new KVStorage(driver, { indexes }).child('users')
          const model = new Map<string, User>()
          for (const o of ops) {
            if (o.t === 'set') {
              await users.set(o.k, o.v)
              model.set(o.k, o.v)
            } else {
              await users.delete(o.k)
              model.delete(o.k)
            }
          }
          const got = (await collect(users.findBy('byAge', { gte: lo, lt: hi }))).map(([k]) => k).sort()
          const exp = [...model]
            .filter(([, v]) => v.age !== undefined && v.age >= lo && v.age < hi)
            .map(([k]) => k)
            .sort()
          expect(got).toEqual(exp)
          for (const tag of ['a', 'b', 'c/d']) {
            const g = (await collect(users.findBy('byTag', tag))).map(([k]) => k).sort()
            expect(g).toEqual(
              [...model]
                .filter(([, v]) => v.tags?.includes(tag))
                .map(([k]) => k)
                .sort(),
            )
          }
          // No stale entries are ever left behind by clean writes.
          let entryCount = 0
          for await (const _ of driver.scan({ gte: 'users/~i/', lt: 'users/~i0' })) entryCount++
          const expectedEntries = [...model.values()].reduce((n, v) => n + (v.age !== undefined ? 1 : 0) + new Set(v.tags).size, 0)
          expect(entryCount).toBe(expectedEntries)
        }),
        { numRuns: 200 },
      )
    })
  }
})