[Beyonce]: Allow Compound Partition/Sort Keys

README.md

@@ -47,7 +47,7 @@ Tables:
 
 #### A note on `partitionKey` and `sortKey` syntax
 
-Beyonce expects you to specify your partition and sort keys as tuple-2s, e.g. `[Author, $id]`. The first element is a "key prefix" and the 2nd must be a field on your model. For example we set the primary key of the `Author` model above to: `[Author, $id]`, would result in the key: `Author-$id`, where `$id` is the value of a specific Author's id.
+Beyonce expects you to specify your partition and sort keys as arrays, e.g. `[Author, $id]`. The first element is a "key prefix" and all subsequent items must be field names on your model. For example we set the primary key of the `Author` model above to: `[Author, $id]`, would result in the key: `Author-$id`, where `$id` is the value of a specific Author's id. And if we wanted a compound key we could do `[Author, $id, $name]`. You can specify compound keys for both partition and sort keys.
 
 Using the example above, if we wanted to place `Books` under the same partition key, then we'd need to set the `Book` model's `partitionKey` to `[Author, $authorId]`.
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ginger.io/beyonce",
-  "version": "0.0.37",
+  "version": "0.0.38",
   "description": "Type-safe DynamoDB query builder for TypeScript. Designed with single-table architecture in mind.",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

src/main/codegen/generateModels.ts

@@ -1,16 +1,20 @@
 import { Model } from "./types"
+import { formatKeyComponent } from "./util"
 
 export function generateModels(models: Model[]) {
   return models.map(generateModel)
 }
 
 function generateModel(model: Model): string {
-  const [pkPrefix, pk] = model.keys.partitionKey
-  const [skPrefix, sk] = model.keys.sortKey
+  const [pkPrefix, ...pkComponents] = model.keys.partitionKey.map(
+    formatKeyComponent
+  )
+
+  const [skPrefix, ...skComponents] = model.keys.sortKey.map(formatKeyComponent)
 
   return `export const ${model.name}Model = ${model.tableName}Table
     .model<${model.name}>(ModelType.${model.name})
-    .partitionKey("${pkPrefix}", "${pk.replace("$", "")}")
-    .sortKey("${skPrefix}", "${sk.replace("$", "")}")
+    .partitionKey(${pkPrefix}, ${pkComponents.join(", ")})
+    .sortKey(${skPrefix}, ${skComponents.join(", ")})
   `
 }

src/main/codegen/types.ts

@@ -17,7 +17,7 @@ export type PartitionDefinition = {
 }
 
 export type ModelDefinition = Keys & Fields
-export type Keys = { partitionKey: [string, string]; sortKey: [string, string] }
+export type Keys = { partitionKey: string[]; sortKey: string[] }
 export type Fields = { [fieldName: string]: string }
 
 export type GSIDefinition = {

src/main/codegen/util.ts

@@ -11,3 +11,7 @@ export function groupBy<T extends { [key: string]: any }, U extends keyof T>(
 
   return Object.entries(groupedItems)
 }
+
+export function formatKeyComponent(component: string): string {
+  return `"${component.replace("$", "")}"`
+}

src/main/dynamo/Model.ts

@@ -10,9 +10,9 @@ export class Model<
   constructor(
     private table: Table,
     private partitionKeyPrefix: string,
-    private partitionKeyField: U,
+    private partitionKeyFields: U[],
     private sortKeyPrefix: string,
-    private sortKeyField: V,
+    private sortKeyFields: V[],
     readonly modelTag: string
   ) {}
 
@@ -22,24 +22,27 @@ export class Model<
     const {
       partitionKeyPrefix,
       sortKeyPrefix,
-      partitionKeyField,
-      sortKeyField,
+      partitionKeyFields,
+      sortKeyFields,
     } = this
+    const pkComponents = partitionKeyFields.map((_) => params[_])
+    const skComponents = sortKeyFields.map((_) => params[_])
     return new PartitionAndSortKey(
       this.table.partitionKeyName,
-      this.buildKey(partitionKeyPrefix, params[partitionKeyField]),
+      this.buildKey(partitionKeyPrefix, ...pkComponents),
 
       this.table.sortKeyName,
-      this.buildKey(sortKeyPrefix, params[sortKeyField]),
+      this.buildKey(sortKeyPrefix, ...skComponents),
       this.modelTag
     )
   }
 
   partitionKey(params: { [X in U]: string }): PartitionKeyAndSortKeyPrefix<T> {
-    const { partitionKeyPrefix, partitionKeyField } = this
+    const { partitionKeyPrefix, partitionKeyFields } = this
+    const pkComponents = partitionKeyFields.map((_) => params[_])
     return new PartitionKeyAndSortKeyPrefix(
       this.table.partitionKeyName,
-      this.buildKey(partitionKeyPrefix, params[partitionKeyField]),
+      this.buildKey(partitionKeyPrefix, ...pkComponents),
       this.table.sortKeyName,
       this.sortKeyPrefix,
       this.modelTag
@@ -50,18 +53,18 @@ export class Model<
     const {
       partitionKeyPrefix,
       sortKeyPrefix,
-      partitionKeyField,
-      sortKeyField,
+      partitionKeyFields,
+      sortKeyFields,
       modelTag,
     } = this
 
     const fieldsWithTag = { ...fields, model: modelTag } as T
 
-    const pk = this.buildKey(
-      partitionKeyPrefix,
-      fieldsWithTag[partitionKeyField]
-    )
-    const sk = this.buildKey(sortKeyPrefix, fieldsWithTag[sortKeyField])
+    const pkComponents = partitionKeyFields.map((_) => fieldsWithTag[_])
+    const skComponents = sortKeyFields.map((_) => fieldsWithTag[_])
+
+    const pk = this.buildKey(partitionKeyPrefix, ...pkComponents)
+    const sk = this.buildKey(sortKeyPrefix, ...skComponents)
 
     return {
       ...fieldsWithTag,
@@ -70,21 +73,21 @@ export class Model<
     }
   }
 
-  private buildKey(prefix: string, key: string): string {
-    return `${prefix}-${key}`
+  private buildKey(...components: string[]): string {
+    return components.join("-")
   }
 }
 
 export class PartitionKeyBuilder<T extends TaggedModel> {
   constructor(private table: Table, private modelTag: string) {}
   partitionKey<U extends keyof T>(
     prefix: string,
-    partitionKeyField: U
+    ...partitionKeyFields: U[]
   ): SortKeyBuilder<T, U> {
     return new SortKeyBuilder(
       this.table,
       prefix,
-      partitionKeyField,
+      partitionKeyFields,
       this.modelTag
     )
   }
@@ -94,17 +97,20 @@ export class SortKeyBuilder<T extends TaggedModel, U extends keyof T> {
   constructor(
     private table: Table,
     private partitionKeyPrefix: string,
-    private partitionKeyField: U,
+    private partitionKeyFields: U[],
     private modelTag: string
   ) {}
 
-  sortKey<V extends keyof T>(prefix: string, sortKeyField: V): Model<T, U, V> {
+  sortKey<V extends keyof T>(
+    prefix: string,
+    ...sortKeyFields: V[]
+  ): Model<T, U, V> {
     return new Model(
       this.table,
       this.partitionKeyPrefix,
-      this.partitionKeyField,
+      this.partitionKeyFields,
       prefix,
-      sortKeyField,
+      sortKeyFields,
       this.modelTag
     )
   }

src/test/dynamo/Beyonce.test.ts

@@ -1,5 +1,7 @@
 import { FixedDataKeyProvider, JayZ } from "@ginger.io/jay-z"
+import { DynamoDB } from "aws-sdk"
 import crypto from "crypto"
+import { Beyonce } from "../../main/dynamo/Beyonce"
 import {
   aMusicianWithTwoSongs,
   byModelAndIdGSI,
@@ -9,18 +11,24 @@ import {
   MusicianPartition,
   Song,
   SongModel,
+  table,
 } from "./models"
 import { setup } from "./util"
-import { Beyonce } from "../../main/dynamo/Beyonce"
-import { table } from "./models"
-import { DynamoDB } from "aws-sdk"
 
 describe("Beyonce", () => {
   // Without encryption
   it("should put and retrieve an item using pk + sk", async () => {
     await testPutAndRetrieveItem()
   })
 
+  it("should put and retrieve a model with a compound partition key", async () => {
+    await testPutAndRetrieveCompoundPartitionKey()
+  })
+
+  it("should put and retrieve a model with a compound sort key", async () => {
+    await testPutAndRetrieveCompoundSortKey()
+  })
+
   it("should update a top-level item attribute", async () => {
     const db = await setup()
     const [musician, _, __] = aMusicianWithTwoSongs()
@@ -213,6 +221,16 @@ describe("Beyonce", () => {
     await testPutAndRetrieveItem(jayZ)
   })
 
+  it("should put and retrieve a model with a compound partition key with jayZ", async () => {
+    const jayZ = await createJayZ()
+    await testPutAndRetrieveCompoundPartitionKey(jayZ)
+  })
+
+  it("should put and retrieve a model with a compound sort key", async () => {
+    const jayZ = await createJayZ()
+    await testPutAndRetrieveCompoundSortKey(jayZ)
+  })
+
   it("should put and delete an item using pk + sk with jayZ", async () => {
     await testPutAndDeleteItem()
   })
@@ -321,6 +339,61 @@ async function testEmptyQuery(jayZ?: JayZ) {
   expect(result).toEqual({ musician: [], song: [] })
 }
 
+async function testPutAndRetrieveCompoundPartitionKey(jayZ?: JayZ) {
+  interface Person {
+    model: "person"
+    first: string
+    last: string
+    sortKey: string
+  }
+
+  const PersonModel = table
+    .model<Person>("person")
+    .partitionKey("Person", "first", "last")
+    .sortKey("Person", "sortKey")
+
+  const db = await setup(jayZ)
+  const model = PersonModel.create({
+    first: "Bob",
+    last: "Smith",
+    sortKey: "sortKey-123",
+  })
+  await db.put(model)
+
+  const result = await db.get(
+    PersonModel.key({ first: "Bob", last: "Smith", sortKey: "sortKey-123" })
+  )
+
+  expect(result).toEqual(model)
+}
+
+async function testPutAndRetrieveCompoundSortKey(jayZ?: JayZ) {
+  interface LineItem {
+    model: "lineItem"
+    orderId: string
+    id: string
+    timestamp: string
+  }
+
+  const LineItemModel = table
+    .model<LineItem>("lineItem")
+    .partitionKey("OrderId", "orderId")
+    .sortKey("LineItem", "id", "timestamp")
+
+  const db = await setup(jayZ)
+  const model = LineItemModel.create({
+    id: "l1",
+    orderId: "o1",
+    timestamp: "456",
+  })
+  await db.put(model)
+
+  const result = await db.get(
+    LineItemModel.key({ id: "l1", orderId: "o1", timestamp: "456" })
+  )
+  expect(result).toEqual(model)
+}
+
 async function testQueryWithPaginatedResults(jayZ?: JayZ) {
   const db = await setup(jayZ)