@model-ts/dynamodb 4.1.0 → 4.2.1

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",
@@ -281,6 +282,61 @@ export interface Sandbox {
 }
 
 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()
 
   const tracked = createTrackedDocClient(docClient, tableName)