Partial Key Queries

src/main/codegen/generateGSIs.ts

@@ -1,4 +1,5 @@
 import { Fields, Table } from "./types"
+import {buildKey, replaceKeyDollar} from "./util"
 
 export function generateGSIs(tables: Table[]): string {
   return tables.map(generateGSIsForTable).join("\n\n")

src/main/codegen/generateModels.ts

@@ -1,4 +1,5 @@
 import { Model } from "./types"
+import { buildKey } from "./util"
 
 export function generateModels(models: Model[]) {
   return models.map(generateModel)
@@ -10,7 +11,8 @@ function generateModel(model: Model): string {
 
   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}", ${buildKey(pk)})
+    .sortKey("${skPrefix}", ${buildKey(sk)})
   `
 }
+

src/main/codegen/generateTables.ts

@@ -1,4 +1,5 @@
 import { Table } from "./types"
+import { buildKeyArray } from "./util"
 
 export function generateTables(tables: Table[]): string {
   const tableCode = tables.map((table) => {
@@ -24,13 +25,18 @@ function generateEncryptionBlacklist(table: Table): string {
     .forEach((model) => {
       const [, pk] = model.keys.partitionKey
       const [, sk] = model.keys.sortKey
-      encryptionBlacklistSet.add(pk.replace("$", ""))
-      encryptionBlacklistSet.add(sk.replace("$", ""))
+
+      const pks = buildKeyArray(model.keys.partitionKey[1])
+      const sks = buildKeyArray(model.keys.sortKey[1])
+      pks.forEach((key: string) => encryptionBlacklistSet.add(key))
+      sks.forEach((key: string) => encryptionBlacklistSet.add(key))
     })
 
   table.gsis.forEach(({ name, partitionKey, sortKey }) => {
-    encryptionBlacklistSet.add(partitionKey.replace("$", ""))
-    encryptionBlacklistSet.add(sortKey.replace("$", ""))
+    const pk = buildKeyArray(partitionKey)
+    const sk = buildKeyArray(sortKey)
+    pk.forEach((key: string) => encryptionBlacklistSet.add(key))
+    sk.forEach((key: string) => encryptionBlacklistSet.add(key))
   })
 
   return JSON.stringify(Array.from(encryptionBlacklistSet))

src/main/codegen/util.ts

@@ -11,3 +11,27 @@ export function groupBy<T extends { [key: string]: any }, U extends keyof T>(
 
   return Object.entries(groupedItems)
 }
+
+export function replaceKeyDollar(key: string | string[]) {
+  if (Array.isArray(key)) {
+    return key.map((k) => k.replace("$", ""))
+  }
+  return key.replace("$", "")
+}
+
+export function buildKey(key: string | string[]) {
+  const cleanKey = replaceKeyDollar(key);
+  if (Array.isArray(cleanKey)) {
+    const keys = cleanKey.map((k) => `"${k}"`).join(', ')
+    return `[${keys}]`
+  }
+  return `"${cleanKey.replace("$", "")}"`
+}
+
+export function buildKeyArray(key: string | string[]) {
+  const keyOrKeys = replaceKeyDollar(key)
+  if (Array.isArray(keyOrKeys)) {
+    return keyOrKeys
+  }
+  return [keyOrKeys]
+}

src/main/dynamo/Beyonce.ts

@@ -120,7 +120,7 @@ export class Beyonce {
       })
       .promise()
 
-    if (unprocessedKeys !== undefined) {
+    if (unprocessedKeys?.entries !== undefined) {
       console.error("Some keys didn't process", unprocessedKeys)
     }
 

src/main/dynamo/GSI.ts

@@ -20,7 +20,7 @@ export class GSI<T extends Model<any, any, any>> {
     this.modelTags = models.map((_) => _.modelTag)
   }
 
-  key(partitionKey: string): PartitionKey<ExtractFields<T>> {
+  key(partitionKey: string | string[]): PartitionKey<ExtractFields<T>> {
     return new PartitionKey(this.partitionKeyName, partitionKey, this.modelTags)
   }
 }

src/main/dynamo/Model.ts

@@ -1,20 +1,21 @@
 import { PartitionAndSortKey, PartitionKeyAndSortKeyPrefix } from "./keys"
 import { Table } from "./Table"
-import { TaggedModel } from "./types"
+import { TaggedModel, AtLeastOne } from "./types"
+import { key } from "aws-sdk/clients/signer"
 
 export class Model<
   T extends TaggedModel,
   U extends keyof T,
   V extends keyof T
-> {
+  > {
   constructor(
     private table: Table,
     private partitionKeyPrefix: string,
-    private partitionKeyField: U,
+    private partitionKeyField: U | U[],
     private sortKeyPrefix: string,
-    private sortKeyField: V,
+    private sortKeyField: V | V[],
     readonly modelTag: string
-  ) {}
+  ) { }
 
   key(
     params: { [X in U]: string } & { [Y in V]: string }
@@ -27,10 +28,36 @@ export class Model<
     } = this
     return new PartitionAndSortKey(
       this.table.partitionKeyName,
-      this.buildKey(partitionKeyPrefix, params[partitionKeyField]),
+      this.buildKey(
+        partitionKeyPrefix,
+        this._buildPartitionKey(params, partitionKeyField)),
 
       this.table.sortKeyName,
-      this.buildKey(sortKeyPrefix, params[sortKeyField]),
+      this.buildKey(
+        sortKeyPrefix,
+        this._buildSortKey(params, sortKeyField)),
+      this.modelTag
+    )
+  }
+
+  partialKey(
+    params: AtLeastOne<{ [X in U]: string }> & AtLeastOne<{ [Y in V]: string }>
+  ): PartitionKeyAndSortKeyPrefix<T> {
+    const {
+      partitionKeyPrefix,
+      sortKeyPrefix,
+      partitionKeyField,
+      sortKeyField,
+    } = this
+    return new PartitionKeyAndSortKeyPrefix(
+      this.table.partitionKeyName,
+      this.buildKey(
+        partitionKeyPrefix,
+        this._buildPartitionKey(params, partitionKeyField)),
+      this.table.sortKeyName,
+      this.buildKey(
+        sortKeyPrefix,
+        this._buildSortKey(params, sortKeyField)),
       this.modelTag
     )
   }
@@ -39,7 +66,9 @@ export class Model<
     const { partitionKeyPrefix, partitionKeyField } = this
     return new PartitionKeyAndSortKeyPrefix(
       this.table.partitionKeyName,
-      this.buildKey(partitionKeyPrefix, params[partitionKeyField]),
+      this.buildKey(
+        partitionKeyPrefix,
+        this._buildPartitionKey(params, partitionKeyField)),
       this.table.sortKeyName,
       this.sortKeyPrefix,
       this.modelTag
@@ -59,9 +88,10 @@ export class Model<
 
     const pk = this.buildKey(
       partitionKeyPrefix,
-      fieldsWithTag[partitionKeyField]
-    )
-    const sk = this.buildKey(sortKeyPrefix, fieldsWithTag[sortKeyField])
+      this._buildPartitionKey(fieldsWithTag, partitionKeyField))
+    const sk = this.buildKey(
+      sortKeyPrefix,
+      this._buildSortKey(fieldsWithTag, sortKeyField))
 
     return {
       ...fieldsWithTag,
@@ -70,13 +100,49 @@ export class Model<
     }
   }
 
-  private buildKey(prefix: string, key: string): string {
+  /* TODO: Had issue mapping unioned fields of T & V */
+  private _buildPartitionKey(model: AtLeastOne<{ [X in U]: string }>, fields: U | U[]): string | string[] {
+    if (Array.isArray(fields)) {
+      const pkey: string[] = [];
+      for (const field of fields) {
+        const value = model[field]
+        if (value) {
+          pkey.push(value)
+        } else {
+          break
+        }
+      };
+      return pkey
+    }
+    return model[fields as U]
+  }
+
+  private _buildSortKey(model: { [Y in V]: string }, fields: V | V[]): string | string[] {
+    if (Array.isArray(fields)) {
+      const skey: string[] = [];
+      for (const field of fields) {
+        const value = model[field]
+        if (value) {
+          skey.push(value)
+        } else {
+          break
+        }
+      };
+      return skey
+    }
+    return model[fields as V]
+  }
+
+  private buildKey(prefix: string, key: string | string[]): string {
+    if (Array.isArray(key)) {
+      return `${prefix}-${key.join('-')}`
+    }
     return `${prefix}-${key}`
   }
 }
 
 export class PartitionKeyBuilder<T extends TaggedModel> {
-  constructor(private table: Table, private modelTag: string) {}
+  constructor(private table: Table, private modelTag: string) { }
   partitionKey<U extends keyof T>(
     prefix: string,
     partitionKeyField: U
@@ -94,11 +160,11 @@ export class SortKeyBuilder<T extends TaggedModel, U extends keyof T> {
   constructor(
     private table: Table,
     private partitionKeyPrefix: string,
-    private partitionKeyField: U,
+    private partitionKeyField: U | U[],
     private modelTag: string
-  ) {}
+  ) { }
 
-  sortKey<V extends keyof T>(prefix: string, sortKeyField: V): Model<T, U, V> {
+  sortKey<V extends keyof T>(prefix: string, sortKeyField: V | V[]): Model<T, U, V> {
     return new Model(
       this.table,
       this.partitionKeyPrefix,

src/main/dynamo/QueryBuilder.ts

@@ -49,7 +49,7 @@ type RawQueryResults<T extends TaggedModel> = {
 /** Builds and executes parameters for a DynamoDB Query operation */
 export class QueryBuilder<T extends TaggedModel> extends QueryExpressionBuilder<
   T
-> {
+  > {
   private scanIndexForward: boolean = true
   private modelTags: string[] = getModelTags(this.config)
 

src/main/dynamo/keys.ts

@@ -3,31 +3,31 @@ import { TaggedModel } from "./types"
 export class PartitionKey<T extends TaggedModel> {
   constructor(
     readonly partitionKeyName: string,
-    readonly partitionKey: string,
+    readonly partitionKey: string | string[],
     readonly modelTags: T["model"][]
-  ) {}
+  ) { }
 }
 
 export class PartitionKeyAndSortKeyPrefix<T extends TaggedModel> {
   constructor(
     readonly partitionKeyName: string,
-    readonly partitionKey: string,
+    readonly partitionKey: string | string[],
     readonly sortKeyName: string,
     readonly sortKeyPrefix: string,
     readonly modelTag: T["model"]
-  ) {}
+  ) { }
 }
 
 export class PartitionAndSortKey<T extends TaggedModel> {
   constructor(
     readonly partitionKeyName: string,
-    readonly partitionKey: string,
+    readonly partitionKey: string | string[],
     readonly sortKeyName: string,
-    readonly sortKey: string,
+    readonly sortKey: string | string[],
     readonly modelTag: T["model"]
-  ) {}
+  ) { }
 
   // Required to avoid TypeScript's structural typing from collapsing
   // PartitionAndSortKey<T> and PartitionAndSortKey<U> into the same object.
-  private dummy(item: T) {}
+  private dummy(item: T) { }
 }

src/main/dynamo/types.ts

@@ -11,3 +11,5 @@ export type ExtractKeyType<T> = T extends PartitionAndSortKey<infer U>
 export type GroupedModels<T extends TaggedModel> = {
   [K in T["model"]]: T extends { model: K } ? T[] : never
 }
+
+export type AtLeastOne<T, U = { [K in keyof T]: Pick<T, K> }> = Partial<T> & U[keyof U]

src/test/codegen/generateCode.test.ts

@@ -322,3 +322,49 @@ Tables:
 
   expect(lines).toContainEqual(`name: BestNameEvah`)
 })
+
+it("should generate a complex key model", () => {
+  const result = generateCode(`
+Tables:
+  ComplexLibrary:
+    Partitions:
+      ComplexAuthors:
+        ComplexAuthor:
+          partitionKey: [Author, $id]
+          sortKey: [Author, [$id, $name]]
+          id: string
+          name: string
+`)
+
+  expect(result).toEqual(`import { Table } from "@ginger.io/beyonce"
+
+export const ComplexLibraryTable = new Table({
+  name: "ComplexLibrary",
+  partitionKeyName: "pk",
+  sortKeyName: "sk",
+  encryptionBlacklist: ["id", "name"]
+})
+
+export enum ModelType {
+  ComplexAuthor = "ComplexAuthor"
+}
+
+export interface ComplexAuthor {
+  model: ModelType.ComplexAuthor
+  id: string
+  name: string
+}
+
+export const ComplexAuthorModel = ComplexLibraryTable.model<ComplexAuthor>(
+  ModelType.ComplexAuthor
+)
+  .partitionKey("Author", "id")
+  .sortKey("Author", ["id", "name"])
+
+export type Model = ComplexAuthor
+
+export const ComplexAuthorsPartition = ComplexLibraryTable.partition([
+  ComplexAuthorModel
+])
+`)
+})