Add support for bigint attribute type (#435)

* First pass at bigint support. * Add support for bigint. * Fix lint. * INclude bigint in $add support. * Remove debugging alias.
jeremydaly · Feb 17, 2023 · 8c89d80 · 8c89d80
1 parent 26b6043
commit 8c89d80
Show file tree

Hide file tree

Showing 18 changed files with 268 additions and 64 deletions.
diff --git a/src/__tests__/bootstrap.test.ts b/src/__tests__/bootstrap.test.ts
@@ -32,6 +32,16 @@ export const DocumentClient = new AWS.DynamoDB.DocumentClient({
   // convertEmptyValues: true
 })
 
+export const DocumentClientWrappedNumbers = new AWS.DynamoDB.DocumentClient({
+  endpoint: 'http://localhost:4567',
+  region: 'us-east-1',
+  credentials: new AWS.Credentials({
+    accessKeyId: 'test',
+    secretAccessKey: 'test'
+  }),
+  wrapNumbers: true
+})
+
 // Delay helper
 export const delay = (ms: number) => new Promise(res => setTimeout(res, ms))
 

diff --git a/src/__tests__/entity-creation.unit.test.ts b/src/__tests__/entity-creation.unit.test.ts
@@ -1,6 +1,6 @@
 import Table from '../classes/Table'
 import Entity from '../classes/Entity'
-import { DocumentClient } from './bootstrap.test'
+import { DocumentClient, DocumentClientWrappedNumbers } from './bootstrap.test'
 
 const tableAddEntity = jest.spyOn(Table.prototype, 'addEntity').mockReturnValue()
 
@@ -222,7 +222,7 @@ describe('Entity creation', () => {
         }
       } as const)
     expect(result).toThrow(
-      `Invalid or missing type for 'pk'. Valid types are 'string', 'boolean', 'number', 'list', 'map', 'binary', and 'set'.`
+      `Invalid or missing type for 'pk'. Valid types are 'string', 'boolean', 'number', 'bigint', 'list', 'map', 'binary', and 'set'.`
     )
   })
 
@@ -236,7 +236,7 @@ describe('Entity creation', () => {
         }
       } as const)
     expect(result).toThrow(
-      `Invalid or missing type for 'pk'. Valid types are 'string', 'boolean', 'number', 'list', 'map', 'binary', and 'set'.`
+      `Invalid or missing type for 'pk'. Valid types are 'string', 'boolean', 'number', 'bigint', 'list', 'map', 'binary', and 'set'.`
     )
   })
 
@@ -287,7 +287,7 @@ describe('Entity creation', () => {
           test: { type: 'set', setType: 'test' }
         }
       } as const)
-    expect(result).toThrow(`Invalid 'setType', must be 'string', 'number', or 'binary'`)
+    expect(result).toThrow(`Invalid 'setType', must be 'string', 'number', 'bigint', or 'binary'`)
   })
 
   it(`fails when setting an invalid attribute property type`, () => {
@@ -425,6 +425,38 @@ describe('Entity creation', () => {
     })
   })
 
+  it('requires wrapNumbers if DocumentClient provided', () => {
+    const entityDefn = {
+      name: 'TestEntity',
+      attributes: {
+        pk: { partitionKey: true },
+        test: { type: 'bigint' }
+      }
+    } as const
+
+    const TestTable1 = new Table({
+      name: 'test-table',
+      partitionKey: 'pk',
+      sortKey: 'sk',
+      DocumentClient
+    })
+    expect(() => new Entity({ ...entityDefn, table: TestTable1 })).toThrow(
+      'Please set `wrapNumbers: true` in your DocumentClient to avoid losing precision with bigint fields'
+    )
+
+    const TestTable2 = new Table({
+      name: 'test-table',
+      partitionKey: 'pk',
+      sortKey: 'sk',
+      DocumentClient: DocumentClientWrappedNumbers
+    })
+    const TestEntity = new Entity({
+      ...entityDefn,
+      table: TestTable2
+    })
+    expect(TestEntity.schema.attributes).toHaveProperty('test')
+  })
+
   // it('creates entity w/ table', async () => {
 
   //   // Create basic table

diff --git a/src/__tests__/entity.parse.unit.test.ts b/src/__tests__/entity.parse.unit.test.ts
@@ -1,13 +1,15 @@
-import { DocumentClient } from './bootstrap.test'
+import { DocumentClient, DocumentClientWrappedNumbers } from './bootstrap.test'
 
 import Table from '../classes/Table'
 import Entity from '../classes/Entity'
+import { toDynamoBigInt } from '../lib/utils'
+import DynamoDB from 'aws-sdk/clients/dynamodb'
 
 const TestTable = new Table({
   name: 'test-table',
   partitionKey: 'pk',
   sortKey: 'sk',
-  DocumentClient
+  DocumentClient: DocumentClientWrappedNumbers
 })
 
 const TestEntity = new Entity({
@@ -19,6 +21,8 @@ const TestEntity = new Entity({
     test_string_coerce: { type: 'string' },
     test_number: { type: 'number', alias: 'count', coerce: false },
     test_number_coerce: { type: 'number', default: 0 },
+    test_bigint: { type: 'bigint', coerce: false },
+    test_bigint_coerce: { type: 'bigint' },
     test_boolean: { type: 'boolean', coerce: false },
     test_boolean_coerce: { type: 'boolean' },
     test_list: { type: 'list' },
@@ -30,8 +34,10 @@ const TestEntity = new Entity({
     test_string_set_type: { type: 'set', setType: 'string' },
     test_number_set_type: { type: 'set', setType: 'number' },
     test_binary_set_type: { type: 'set', setType: 'binary' },
+    test_bigint_set_type: { type: 'set', setType: 'bigint' },
     test_string_set_type_coerce: { type: 'set', setType: 'string', coerce: true },
     test_number_set_type_coerce: { type: 'set', setType: 'number', coerce: true },
+    test_bigint_set_type_coerce: { type: 'set', setType: 'bigint', coerce: true },
     test_binary: { type: 'binary' },
     simple_string: 'string',
     format_simple_string: {
@@ -197,4 +203,58 @@ describe('parse', () => {
       test_composite2: 'email'
     })
   })
+
+  it('parses wrapped numbers', () => {
+    const wrap = (value: number) =>
+      DynamoDB.Converter.output({ N: value.toString() }, { wrapNumbers: true })
+
+    const item = TestEntity.parse({
+      pk: 'test@test.com',
+      sk: 'bigint',
+      test_number: wrap(1234.567),
+      test_number_coerce: wrap(-0.0023)
+    })
+    expect(item).toEqual({
+      email: 'test@test.com',
+      test_type: 'bigint',
+      count: 1234.567,
+      test_number_coerce: -0.0023
+    })
+  })
+
+  it('parses bigints', () => {
+    const item = TestEntity.parse({
+      pk: 'test@test.com',
+      sk: 'bigint',
+      test_bigint: toDynamoBigInt(BigInt('90071992547409911234')),
+      test_bigint_coerce: '12345'
+    })
+    expect(item).toEqual({
+      email: 'test@test.com',
+      test_type: 'bigint',
+      test_bigint: BigInt('90071992547409911234'),
+      test_bigint_coerce: BigInt('12345')
+    })
+  })
+
+  it('parses bigint sets', () => {
+    const item = TestEntity.parse({
+      pk: 'test@test.com',
+      sk: 'bigint',
+      test_bigint_set_type: DocumentClient.createSet([
+        toDynamoBigInt(BigInt('90071992547409911234')),
+        toDynamoBigInt(BigInt('-90071992547409911234')),
+        1234
+      ]),
+    })
+    expect(item).toEqual({
+      email: 'test@test.com',
+      test_type: 'bigint',
+      test_bigint_set_type: [
+        BigInt('90071992547409911234'),
+        BigInt('-90071992547409911234'),
+        BigInt(1234),
+      ]
+    })
+  })
 }) // end parse
diff --git a/src/__tests__/parseTableAttributes.unit.test.ts b/src/__tests__/parseTableAttributes.unit.test.ts
@@ -19,7 +19,7 @@ describe('parseTableAttributes', () => {
     expect(() => {
       parseTableAttributes(Object.assign({}, attrs, { test: {} }), 'pk', 'sk')
     }).toThrow(
-      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'list', 'map', 'binary', and 'set'.`
+      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'bigint', 'list', 'map', 'binary', and 'set'.`
     )
   })
 
@@ -43,7 +43,7 @@ describe('parseTableAttributes', () => {
     expect(() => {
       parseTableAttributes(Object.assign({}, attrs, { test: 'not-a-type' }), 'pk', 'sk')
     }).toThrow(
-      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'list', 'map', 'binary', and 'set'.`
+      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'bigint', 'list', 'map', 'binary', and 'set'.`
     )
   })
 
@@ -67,7 +67,7 @@ describe('parseTableAttributes', () => {
     expect(() => {
       parseTableAttributes(Object.assign({}, attrs, { test: { type: 'not-a-type' } }), 'pk', 'sk')
     }).toThrow(
-      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'list', 'map', 'binary', and 'set'.`
+      `Invalid or missing type for 'test'. Valid types are 'string', 'boolean', 'number', 'bigint', 'list', 'map', 'binary', and 'set'.`
     )
   })
 
@@ -88,7 +88,7 @@ describe('parseTableAttributes', () => {
         'pk',
         'sk'
       )
-    }).toThrow(`Invalid 'setType', must be 'string', 'number', or 'binary'`)
+    }).toThrow(`Invalid 'setType', must be 'string', 'number', 'bigint' or 'binary'`)
   })
 
   it('fails when attribute has an invalid config option', async () => {

diff --git a/src/__tests__/type-infering.unit.test.ts b/src/__tests__/type-infering.unit.test.ts
@@ -52,12 +52,12 @@ type ExpectedQueryOpts<
     parseAsEntity: string
     select: DocumentClientType.Select
     filters: ConditionsOrFilters<FilteredAttributes>
-    eq: string | number
-    lt: string | number
-    lte: string | number
-    gt: string | number
-    gte: string | number
-    between: [string, string] | [number, number]
+    eq: string | number | bigint
+    lt: string | number | bigint
+    lte: string | number | bigint
+    gt: string | number | bigint
+    gte: string | number | bigint
+    between: [string, string] | [number, number] | [bigint, bigint]
     beginsWith: string
     startKey: {}
   }

diff --git a/src/__tests__/validateTypes.unit.test.ts b/src/__tests__/validateTypes.unit.test.ts
@@ -1,3 +1,4 @@
+import { toDynamoBigInt } from '../lib/utils'
 import validateTypes from '../lib/validateTypes'
 
 import { DocumentClient } from './bootstrap.test'
@@ -163,4 +164,34 @@ describe('validateTypes', () => {
       validateTypes()({ type: 'set', setType: 'string', coerce: true }, 'attr', 'test')
     }).toThrow(`DocumentClient required for this operation`)
   })
+
+  it('coerces strings to bigints', async () => {
+    const result = validateTypes(DocumentClient)({ type: 'bigint', coerce: true }, 'attr', '123000000000000000000001')
+    expect(result).toEqual(toDynamoBigInt(BigInt('123000000000000000000001')))
+  })
+
+  it('coerces numbers to bigints', async () => {
+    const result = validateTypes(DocumentClient)({ type: 'bigint', coerce: true }, 'attr', 12e10)
+    expect(result).toEqual(toDynamoBigInt(BigInt('120000000000')))
+  })
+
+  it('fails if non-bigint values if coerce is not true', async () => {
+    expect(() => {
+      validateTypes(DocumentClient)({ type: 'bigint' }, 'attr', '123000000000000000000001')
+    }).toThrow(`'attr' must be of type bigint`)
+    expect(() => {
+      validateTypes(DocumentClient)({ type: 'bigint' }, 'attr', 12e10)
+    }).toThrow(`'attr' must be of type bigint`)
+  })
+
+  it('converts arrays of bigints to sets', async () => {
+    const result = validateTypes(DocumentClient)({ type: 'set', setType: 'bigint' }, 'attr', [
+      BigInt(-1234),
+      BigInt('123000000000000000000001')
+    ])
+    expect(result).toEqual(DocumentClient.createSet([
+      toDynamoBigInt(BigInt(-1234)),
+      toDynamoBigInt(BigInt('123000000000000000000001'))
+    ], { validate: false }))
+  })
 })
diff --git a/src/classes/Entity/Entity.ts b/src/classes/Entity/Entity.ts
@@ -1021,7 +1021,7 @@ class Entity<Name extends string = string,
       ) {
         // If a number or a set and adding
         if (
-          ['number', 'set'].includes(mapping.type) &&
+          ['bigint', 'number', 'set'].includes(mapping.type) &&
           data[field]?.$add !== undefined &&
           data[field]?.$add !== null
         ) {

diff --git a/src/classes/Entity/types.ts b/src/classes/Entity/types.ts
@@ -32,7 +32,7 @@ export interface EntityConstructor<
 }
 
 export type KeyAttributeDefinition = {
-  type: 'string' | 'number' | 'binary'
+  type: 'string' | 'number' | 'bigint' | 'binary'
   // 🔨 TOIMPROVE: Probably typable
   default: any
   hidden: boolean
@@ -213,6 +213,7 @@ export type FromDynamoData<T extends DynamoDBTypes> = {
   string: string
   boolean: boolean
   number: number
+  bigint: bigint
   list: any[]
   map: any
   binary: any
@@ -320,13 +321,13 @@ export type ConditionOrFilter<Attributes extends A.Key = A.Key> = (
     negate: boolean
     entity: string
     // 🔨 TOIMPROVE: Probably typable
-    eq: string | number | boolean | null
-    ne: string | number | boolean | null
+    eq: string | number | bigint | boolean | null
+    ne: string | number | bigint | boolean | null
     lt: string | number
-    lte: string | number
-    gt: string | number
-    gte: string | number
-    between: [string, string] | [number, number]
+    lte: string | bigint | number
+    gt: string | bigint | number
+    gte: string | bigint | number
+    between: [string, string] | [number, number] | [bigint, bigint]
     beginsWith: string
     in: any[]
   }>
@@ -507,6 +508,11 @@ export type AttributeUpdateInput<AttributeType> =
       | { $delete?: number[]; $add?: number[]; $prepend?: AttributeType; $append?: number[] }
       | string[]
     >
+  | If<
+      B.Or<A.Equals<AttributeType, bigint[]>, A.Equals<AttributeType, bigint[] | undefined>>,
+      | { $delete?: bigint[]; $add?: bigint[]; $prepend?: AttributeType; $append?: bigint[] }
+      | string[]
+    >
   | If<
       B.Or<A.Equals<AttributeType, string[]>, A.Equals<AttributeType, string[] | undefined>>,
       | { $delete?: string[]; $add?: string[]; $prepend?: AttributeType; $append?: string[] }
@@ -518,6 +524,7 @@ export type AttributeUpdateInput<AttributeType> =
       | boolean[]
     >
   | If<B.Or<A.Equals<AttributeType, FromDynamoData<'number'>>, A.Equals<AttributeType, FromDynamoData<'number'> | undefined>>, { $add?: number }>
+  | If<B.Or<A.Equals<AttributeType, FromDynamoData<'bigint'>>, A.Equals<AttributeType, FromDynamoData<'bigint'> | undefined>>, { $add?: bigint }>
 
 export type DeleteOptionsReturnValues = 'NONE' | 'ALL_OLD'
 

diff --git a/src/classes/Table/Table.ts b/src/classes/Table/Table.ts
@@ -89,7 +89,7 @@ class Table<Name extends string, PartitionKey extends A.Key, SortKey extends A.K
 
   // Validate and sets the document client (extend with options.convertEmptyValues because it's not typed)
   set DocumentClient(
-    docClient: (DocumentClient & { options?: { convertEmptyValues: boolean } }) | undefined,
+    docClient: (DocumentClient & { options?: { convertEmptyValues: boolean; wrapNumbers: boolean } }) | undefined,
   ) {
     // If a valid document client
     // @ts-ignore

diff --git a/src/classes/Table/types.ts b/src/classes/Table/types.ts
@@ -25,8 +25,8 @@ export interface TableConstructor<
   removeNullAttributes?: boolean
 }
 
-export type DynamoDBTypes = 'string' | 'boolean' | 'number' | 'list' | 'map' | 'binary' | 'set'
-export type DynamoDBKeyTypes = 'string' | 'number' | 'binary'
+export type DynamoDBTypes = 'string' | 'boolean' | 'number' | 'bigint' | 'list' | 'map' | 'binary' | 'set'
+export type DynamoDBKeyTypes = 'string' | 'number' | 'bigint' | 'binary'
 
 export interface executeParse {
   execute?: boolean
@@ -65,12 +65,12 @@ export type $QueryOptions<
     reverse: boolean
     select: DocumentClient.Select
     // 🔨 TOIMPROVE: Probably typable (should be the same as sort key)
-    eq: string | number
-    lt: string | number
-    lte: string | number
-    gt: string | number
-    gte: string | number
-    between: [string, string] | [number, number]
+    eq: string | number | bigint
+    lt: string | number | bigint
+    lte: string | number | bigint
+    gt: string | number | bigint
+    gte: string | number | bigint
+    between: [string, string] | [number, number] | [bigint, bigint]
     beginsWith: string
     startKey: {}
   }