@model-ts/dynamodb 4.0.0 → 4.2.0

package/src/in-memory/treap.ts

@@ -0,0 +1,239 @@
+interface TreapNode<V> {
+  key: string
+  priority: number
+  value: V
+  left: TreapNode<V> | null
+  right: TreapNode<V> | null
+}
+
+export interface TreapBounds {
+  lower?: { key: string; inclusive: boolean }
+  upper?: { key: string; inclusive: boolean }
+}
+
+export class DeterministicTreap<V> {
+  private root: TreapNode<V> | null = null
+  private _size = 0
+
+  get size() {
+    return this._size
+  }
+
+  insert(key: string, value: V, priority: number): void {
+    const [nextRoot, inserted] = this.insertNode(this.root, key, value, priority)
+    this.root = nextRoot
+    if (inserted) this._size += 1
+  }
+
+  remove(key: string): boolean {
+    const [nextRoot, removed] = this.removeNode(this.root, key)
+    this.root = nextRoot
+    if (removed) this._size -= 1
+    return removed
+  }
+
+  has(key: string): boolean {
+    let current = this.root
+    while (current) {
+      if (key === current.key) return true
+      current = key < current.key ? current.left : current.right
+    }
+    return false
+  }
+
+  *iterate(
+    direction: "asc" | "desc",
+    bounds: TreapBounds = {}
+  ): IterableIterator<{ key: string; value: V }> {
+    if (direction === "asc") {
+      yield* this.iterateAsc(bounds)
+      return
+    }
+
+    yield* this.iterateDesc(bounds)
+  }
+
+  clear() {
+    this.root = null
+    this._size = 0
+  }
+
+  private rotateRight(node: TreapNode<V>): TreapNode<V> {
+    const left = node.left
+    if (!left) return node
+
+    node.left = left.right
+    left.right = node
+    return left
+  }
+
+  private rotateLeft(node: TreapNode<V>): TreapNode<V> {
+    const right = node.right
+    if (!right) return node
+
+    node.right = right.left
+    right.left = node
+    return right
+  }
+
+  private insertNode(
+    node: TreapNode<V> | null,
+    key: string,
+    value: V,
+    priority: number
+  ): [TreapNode<V>, boolean] {
+    if (!node) {
+      return [
+        {
+          key,
+          priority,
+          value,
+          left: null,
+          right: null,
+        },
+        true,
+      ]
+    }
+
+    if (key === node.key) {
+      node.value = value
+      return [node, false]
+    }
+
+    if (key < node.key) {
+      const [next, inserted] = this.insertNode(node.left, key, value, priority)
+      node.left = next
+
+      if (node.left && node.left.priority < node.priority) {
+        node = this.rotateRight(node)
+      }
+
+      return [node, inserted]
+    }
+
+    const [next, inserted] = this.insertNode(node.right, key, value, priority)
+    node.right = next
+
+    if (node.right && node.right.priority < node.priority) {
+      node = this.rotateLeft(node)
+    }
+
+    return [node, inserted]
+  }
+
+  private removeNode(
+    node: TreapNode<V> | null,
+    key: string
+  ): [TreapNode<V> | null, boolean] {
+    if (!node) return [null, false]
+
+    if (key < node.key) {
+      const [next, removed] = this.removeNode(node.left, key)
+      node.left = next
+      return [node, removed]
+    }
+
+    if (key > node.key) {
+      const [next, removed] = this.removeNode(node.right, key)
+      node.right = next
+      return [node, removed]
+    }
+
+    if (!node.left) return [node.right, true]
+    if (!node.right) return [node.left, true]
+
+    if (node.left.priority < node.right.priority) {
+      const rotated = this.rotateRight(node)
+      const [next, removed] = this.removeNode(rotated.right, key)
+      rotated.right = next
+      return [rotated, removed]
+    }
+
+    const rotated = this.rotateLeft(node)
+    const [next, removed] = this.removeNode(rotated.left, key)
+    rotated.left = next
+    return [rotated, removed]
+  }
+
+  private *iterateAsc(bounds: TreapBounds): IterableIterator<{ key: string; value: V }> {
+    const stack: TreapNode<V>[] = []
+
+    const pushLeft = (node: TreapNode<V> | null) => {
+      while (node) {
+        if (this.isBelowLowerBound(node.key, bounds.lower)) {
+          node = node.right
+        } else {
+          stack.push(node)
+          node = node.left
+        }
+      }
+    }
+
+    pushLeft(this.root)
+
+    while (stack.length) {
+      const node = stack.pop()!
+
+      if (this.isAboveUpperBound(node.key, bounds.upper)) {
+        return
+      }
+
+      if (!this.isBelowLowerBound(node.key, bounds.lower)) {
+        yield { key: node.key, value: node.value }
+      }
+
+      pushLeft(node.right)
+    }
+  }
+
+  private *iterateDesc(bounds: TreapBounds): IterableIterator<{ key: string; value: V }> {
+    const stack: TreapNode<V>[] = []
+
+    const pushRight = (node: TreapNode<V> | null) => {
+      while (node) {
+        if (this.isAboveUpperBound(node.key, bounds.upper)) {
+          node = node.left
+        } else {
+          stack.push(node)
+          node = node.right
+        }
+      }
+    }
+
+    pushRight(this.root)
+
+    while (stack.length) {
+      const node = stack.pop()!
+
+      if (this.isBelowLowerBound(node.key, bounds.lower)) {
+        return
+      }
+
+      if (!this.isAboveUpperBound(node.key, bounds.upper)) {
+        yield { key: node.key, value: node.value }
+      }
+
+      pushRight(node.left)
+    }
+  }
+
+  private isBelowLowerBound(
+    key: string,
+    lower?: { key: string; inclusive: boolean }
+  ): boolean {
+    if (!lower) return false
+
+    if (lower.inclusive) return key < lower.key
+    return key <= lower.key
+  }
+
+  private isAboveUpperBound(
+    key: string,
+    upper?: { key: string; inclusive: boolean }
+  ): boolean {
+    if (!upper) return false
+
+    if (upper.inclusive) return key > upper.key
+    return key >= upper.key
+  }
+}

package/src/in-memory/utils.ts

@@ -0,0 +1,45 @@
+import crypto from "crypto"
+
+export type InMemoryItem = { [key: string]: any }
+
+export const ITEM_KEY_SEPARATOR = "\u0000"
+
+export const encodeCompositeKey = (...parts: string[]): string =>
+  parts
+    .map((part) => `${part.length}:${part}`)
+    .join(ITEM_KEY_SEPARATOR)
+
+export const encodeItemKey = (pk: string, sk: string): string =>
+  encodeCompositeKey(pk, sk)
+
+export const encodeIndexEntryKey = (rangeKey: string, itemKey: string): string =>
+  `${rangeKey}${ITEM_KEY_SEPARATOR}${itemKey}`
+
+export const cloneItem = <T>(item: T): T => JSON.parse(JSON.stringify(item))
+
+export const stablePriority = (
+  indexName: string,
+  hashKey: string,
+  rangeKey: string,
+  itemKey: string
+): number => {
+  const hash = crypto
+    .createHash("sha256")
+    .update(`${indexName}::${hashKey}::${rangeKey}::${itemKey}`)
+    .digest()
+
+  return hash.readUInt32BE(0)
+}
+
+export const sortItemsByPKSK = (items: InMemoryItem[]): InMemoryItem[] =>
+  [...items].sort((a, b) => {
+    const pkA = String(a.PK ?? "")
+    const pkB = String(b.PK ?? "")
+    if (pkA !== pkB) return pkA < pkB ? -1 : 1
+
+    const skA = String(a.SK ?? "")
+    const skB = String(b.SK ?? "")
+    if (skA !== skB) return skA < skB ? -1 : 1
+
+    return 0
+  })

package/src/sandbox.ts

@@ -4,6 +4,7 @@ import DynamoDB from "aws-sdk/clients/dynamodb"
 import { formatSnapshotDiff } from "./diff"
 import { Client } from "./client"
 import { GSI_NAMES } from "./gsi"
+import { createInMemoryDocumentClient } from "./in-memory"
 
 const ddb = new DynamoDB({
   accessKeyId: "xxx",
@@ -105,18 +106,241 @@ export const getTableContents = async (
   return acc
 }
 
+// -------------------------------------------------------------------------------------
+// Tracking
+// -------------------------------------------------------------------------------------
+
+interface TrackedEntry {
+  pk: string
+  sk: string
+  original: any | null
+}
+
+const WRITE_METHODS = new Set([
+  "put",
+  "update",
+  "delete",
+  "batchWrite",
+  "transactWrite",
+])
+
+function createTrackedDocClient(
+  original: DynamoDB.DocumentClient,
+  tableName: string
+) {
+  let isTracking = false
+  const trackedKeys = new Map<string, TrackedEntry>()
+
+  const captureKey = async (pk: string, sk: string) => {
+    const compositeKey = `${pk}__${sk}`
+    if (trackedKeys.has(compositeKey)) return
+
+    const { Item } = await original
+      .get({ TableName: tableName, Key: { PK: pk, SK: sk } })
+      .promise()
+
+    trackedKeys.set(compositeKey, { pk, sk, original: Item ?? null })
+  }
+
+  const captureKeysForOperation = async (method: string, params: any) => {
+    switch (method) {
+      case "put":
+        if (params.TableName === tableName && params.Item) {
+          await captureKey(params.Item.PK, params.Item.SK)
+        }
+        break
+      case "update":
+      case "delete":
+        if (params.TableName === tableName && params.Key) {
+          await captureKey(params.Key.PK, params.Key.SK)
+        }
+        break
+      case "batchWrite": {
+        const tableItems = params.RequestItems?.[tableName] || []
+        await Promise.all(
+          tableItems.map((item: any) => {
+            if (item.PutRequest) {
+              return captureKey(
+                item.PutRequest.Item.PK,
+                item.PutRequest.Item.SK
+              )
+            }
+            if (item.DeleteRequest) {
+              return captureKey(
+                item.DeleteRequest.Key.PK,
+                item.DeleteRequest.Key.SK
+              )
+            }
+          })
+        )
+        break
+      }
+      case "transactWrite": {
+        const transactItems = params.TransactItems || []
+        await Promise.all(
+          transactItems
+            .map((item: any) => {
+              if (item.Put?.TableName === tableName) {
+                return captureKey(item.Put.Item.PK, item.Put.Item.SK)
+              }
+              if (item.Update?.TableName === tableName) {
+                return captureKey(item.Update.Key.PK, item.Update.Key.SK)
+              }
+              if (item.Delete?.TableName === tableName) {
+                return captureKey(item.Delete.Key.PK, item.Delete.Key.SK)
+              }
+            })
+            .filter(Boolean)
+        )
+        break
+      }
+    }
+  }
+
+  const proxy = new Proxy(original, {
+    get(target, prop) {
+      const value = (target as any)[prop]
+      if (value === undefined) return undefined
+
+      if (typeof value === "function") {
+        if (isTracking && WRITE_METHODS.has(prop as string)) {
+          return (params: any) => {
+            const request = value.call(target, params)
+            const origPromise = request.promise.bind(request)
+            request.promise = async () => {
+              await captureKeysForOperation(prop as string, params)
+              return origPromise()
+            }
+            return request
+          }
+        }
+        return value.bind(target)
+      }
+
+      return value
+    },
+  })
+
+  return {
+    proxy: proxy as DynamoDB.DocumentClient,
+    startTracking: () => {
+      isTracking = true
+      trackedKeys.clear()
+    },
+    rollback: async () => {
+      isTracking = false
+
+      const entries = Array.from(trackedKeys.values())
+      const toDelete = entries.filter((e) => e.original === null)
+      const toRestore = entries.filter((e) => e.original !== null)
+
+      const deleteChunks = chunksOf(25)(toDelete)
+      const restoreChunks = chunksOf(25)(toRestore)
+
+      await Promise.all([
+        ...deleteChunks.map((chunk) =>
+          original
+            .batchWrite({
+              RequestItems: {
+                [tableName]: chunk.map(({ pk, sk }) => ({
+                  DeleteRequest: { Key: { PK: pk, SK: sk } },
+                })),
+              },
+            })
+            .promise()
+        ),
+        ...restoreChunks.map((chunk) =>
+          original
+            .batchWrite({
+              RequestItems: {
+                [tableName]: chunk.map(({ original: item }) => ({
+                  PutRequest: { Item: item },
+                })),
+              },
+            })
+            .promise()
+        ),
+      ])
+
+      trackedKeys.clear()
+    },
+  }
+}
+
+// -------------------------------------------------------------------------------------
+// Sandbox
+// -------------------------------------------------------------------------------------
+
 export interface Sandbox {
   destroy: () => Promise<void>
   snapshot: () => Promise<{ [key: string]: any }>
   seed: (...args: Array<{ [key: string]: any }>) => Promise<void>
   get: (pk: string, sk: string) => Promise<null | any>
   diff: (before: { [key: string]: any }) => Promise<string>
+  startTracking: () => void
+  rollback: () => Promise<void>
 }
 
 export const createSandbox = async (client: Client): Promise<Sandbox> => {
+  if (process.env.EXPERIMENTAL_DYNAMODB_IN_MEMORY === "1") {
+    const tableName = crypto.randomBytes(20).toString("hex")
+    const inMemoryClient =
+      createInMemoryDocumentClient() as any as DynamoDB.DocumentClient & {
+        __inMemorySnapshot: (name: string) => { [key: string]: any }
+        __inMemoryResetTable: (name: string) => void
+      }
+
+    const tracked = createTrackedDocClient(inMemoryClient, tableName)
+
+    client.setDocumentClient(tracked.proxy)
+    client.setTableName(tableName)
+
+    return {
+      destroy: async () => {
+        inMemoryClient.__inMemoryResetTable(tableName)
+      },
+      snapshot: async () => inMemoryClient.__inMemorySnapshot(tableName),
+      seed: async (...args: Array<{ [key: string]: any }>) => {
+        const chunks = chunksOf(25)(args)
+
+        await Promise.all(
+          chunks.map(async (chunk) => {
+            const items = chunk.map((i) =>
+              typeof i?._model?.__dynamoDBEncode === "function"
+                ? i._model.__dynamoDBEncode(i)
+                : typeof i.encode === "function"
+                ? i.encode()
+                : i
+            )
+
+            return client.documentClient
+              .batchWrite({
+                RequestItems: {
+                  [tableName]: items.map((i) => ({ PutRequest: { Item: i } })),
+                },
+              })
+              .promise()
+          })
+        )
+      },
+      get: (pk: string, sk: string) =>
+        client.documentClient
+          .get({ TableName: tableName, Key: { PK: pk, SK: sk } })
+          .promise()
+          .then(({ Item }) => Item ?? null),
+      diff: async (before) => {
+        const snapshot = inMemoryClient.__inMemorySnapshot(tableName)
+        return formatSnapshotDiff(before, snapshot)
+      },
+      startTracking: tracked.startTracking,
+      rollback: tracked.rollback,
+    }
+  }
+
   const tableName = await createTable()
 
-  client.setDocumentClient(docClient)
+  const tracked = createTrackedDocClient(docClient, tableName)
+  client.setDocumentClient(tracked.proxy)
   client.setTableName(tableName)
 
   return {
@@ -155,5 +379,7 @@ export const createSandbox = async (client: Client): Promise<Sandbox> => {
 
       return formatSnapshotDiff(before, snapshot)
     },
+    startTracking: tracked.startTracking,
+    rollback: tracked.rollback,
   }
 }