feat: add zod-based schema parsing capability

Previously, generated row types accounted for JSON/JSONB columns by
making them generic parameters to the row/insert types, but the actual
DataSource didn't really account for that at all.

Thus, when a caller had a generic type, they would have to do a bump of
hoop-jumping to actually verify the schema they are _claiming_ applies
to the column.

This change adds the ability to specify Zod schemas to be used to
validate JSON/JSONB columns. Those schemas are used at
INSERT/UPDATE-time, but *NOT* at SELECT time.

Parsing columns coming from the database may be something added in the
future, but for now it is much safer to avoid.

package.json

@@ -73,9 +73,11 @@
     "semantic-release": "^19.0.2",
     "ts-jest": "^27.1.3",
     "ts-node": "^10.4.0",
-    "typescript": "^4.5.5"
+    "typescript": "^4.5.5",
+    "zod": "^3.14.4"
   },
   "peerDependencies": {
-    "typescript": ">= 3.9.0"
+    "typescript": ">= 3.9.0",
+    "zod": "^3.14.4"
   }
 }

src/datasource/DBDataSource.ts

@@ -25,6 +25,8 @@ import {
 } from './queries/types'
 import { KeyValueCache } from 'apollo-server-caching'
 import { AsyncLocalStorage } from 'async_hooks'
+import type { ZodSchema } from 'zod'
+import { isSqlToken } from './queries/utils'
 
 export interface QueryOptions<TRowType, TResultType = TRowType>
   extends BuilderOptions<TRowType> {
@@ -66,7 +68,11 @@ interface ExtendedDatabasePool<TRowType> extends DatabasePool {
 export default class DBDataSource<
   TRowType,
   TContext = unknown,
-  TInsertType extends { [K in keyof TRowType]?: unknown } = TRowType
+  TInsertType extends { [K in keyof TRowType]?: unknown } = TRowType,
+  TColumnTypes extends Record<keyof TRowType, string> = Record<
+    keyof TRowType,
+    string
+  >
 > implements DataSource<TContext>
 {
   protected normalizers: KeyNormalizers = {
@@ -129,7 +135,12 @@ export default class DBDataSource<
      *
      * EVERY DATASOURCE MUST PROVIDE THIS AS STUFF WILL BREAK OTHERWISE. SORRY.
      */
-    protected readonly columnTypes: Record<keyof TRowType, string>
+    protected readonly columnTypes: TColumnTypes,
+    protected readonly columnSchemas?: {
+      [K in keyof TRowType as TColumnTypes[K] extends 'json' | 'jsonb'
+        ? K
+        : never]?: ZodSchema
+    }
   ) {
     this.pool = pool as ExtendedDatabasePool<TRowType>
     this.pool.async ||= new AsyncLocalStorage()
@@ -313,6 +324,12 @@ export default class DBDataSource<
       }
     }
 
+    if (isArray(rows)) {
+      rows.map((row) => this.parseColumnSchemas(row))
+    } else {
+      rows = this.parseColumnSchemas(rows)
+    }
+
     const query = this.builder.insert(rows, options)
     return await this.query(query, options)
   }
@@ -357,6 +374,7 @@ export default class DBDataSource<
     data: UpdateSet<TRowType>,
     options?: QueryOptions<TRowType>
   ): Promise<TRowType | readonly TRowType[] | null> {
+    data = this.parseColumnSchemas(data)
     const query = this.builder.update(data, options)
     return await this.query(query, options)
   }
@@ -563,6 +581,26 @@ export default class DBDataSource<
     )
   }
 
+  protected parseColumnSchemas<TRow extends AllowSql<TRowType | TInsertType>>(
+    row: TRow
+  ): TRow {
+    if (!this.columnSchemas) {
+      return row
+    }
+    const keys = Object.keys(
+      this.columnSchemas
+    ) as (keyof typeof this.columnSchemas)[]
+    for (const column of keys) {
+      const schema = this.columnSchemas[column]
+      if (!schema || isSqlToken(row[column])) {
+        continue
+      }
+
+      row[column] = schema.parse(row[column])
+    }
+    return row
+  }
+
   private transformResult<TInput, TOutput>(input: TInput): TOutput {
     const transform = this.normalizers.columnToKey
 

src/datasource/__tests__/integration.test.ts

@@ -1,21 +1,26 @@
 import assert from 'assert'
 import { createPool, DatabasePool, sql, SqlSqlToken } from 'slonik'
+import { z, ZodSchema } from 'zod'
 
 import { DBDataSource } from '..'
 
-interface DummyRowType {
+interface DummyRowType<TSchema = DefaultJsonbSchema> {
   id: number
   name: string
   code: string
   withDefault?: string | SqlSqlToken
   camelCase: string
   tsTest: Date
   dateTest: Date
-  jsonbTest: { a: number }
+  jsonbTest: TSchema
   nullable?: string | null
 }
 
-const columnTypes: Record<keyof DummyRowType, string> = {
+interface DefaultJsonbSchema {
+  a: number
+}
+
+const columnTypes = {
   id: 'int8',
   name: 'citext',
   code: 'text',
@@ -25,19 +30,21 @@ const columnTypes: Record<keyof DummyRowType, string> = {
   dateTest: 'date',
   jsonbTest: 'jsonb',
   nullable: 'text',
-}
+} as const
 
 let pool: DatabasePool
 
-const createRow = (values: Partial<DummyRowType>): DummyRowType => {
+const createRow = <TSchema = DefaultJsonbSchema>(
+  values: Partial<DummyRowType<TSchema>>
+): DummyRowType<TSchema> => {
   return {
     id: 1,
     code: '',
     name: '',
     camelCase: '',
     tsTest: new Date('2020-12-05T00:00:00.000Z'),
     dateTest: new Date('2021-04-19'),
-    jsonbTest: { a: 1 },
+    jsonbTest: { a: 1 } as unknown as TSchema,
     nullable: null,
     ...values,
   }
@@ -77,9 +84,20 @@ afterAll(async () => {
   await pool.end()
 })
 
-class TestDataSource extends DBDataSource<DummyRowType> {
-  constructor() {
-    super(pool, 'test_table', columnTypes)
+class TestDataSource<TSchema = DefaultJsonbSchema> extends DBDataSource<
+  DummyRowType<TSchema>,
+  unknown,
+  DummyRowType<TSchema>,
+  typeof columnTypes
+> {
+  constructor(columnSchemas?: {
+    [K in keyof DummyRowType<TSchema> as typeof columnTypes[K] extends
+      | 'json'
+      | 'jsonb'
+      ? K
+      : never]?: ZodSchema
+  }) {
+    super(pool, 'test_table', columnTypes, columnSchemas)
   }
 
   public idLoader = this.loaders.create('id')
@@ -99,12 +117,12 @@ class TestDataSource extends DBDataSource<DummyRowType> {
   )
 
   // these functions are protected, so we're not normally able to access them
-  public testGet: TestDataSource['get'] = this.get
-  public testCount: TestDataSource['count'] = this.count
-  public testCountGroup: TestDataSource['countGroup'] = this.countGroup
-  public testInsert: TestDataSource['insert'] = this.insert
-  public testUpdate: TestDataSource['update'] = this.update
-  public testDelete: TestDataSource['delete'] = this.delete
+  public testGet: TestDataSource<TSchema>['get'] = this.get
+  public testCount: TestDataSource<TSchema>['count'] = this.count
+  public testCountGroup: TestDataSource<TSchema>['countGroup'] = this.countGroup
+  public testInsert: TestDataSource<TSchema>['insert'] = this.insert
+  public testUpdate: TestDataSource<TSchema>['update'] = this.update
+  public testDelete: TestDataSource<TSchema>['delete'] = this.delete
 }
 
 let ds: TestDataSource
@@ -124,8 +142,8 @@ describe('DBDataSource', () => {
       expect(result).toHaveLength(0)
     })
 
-    it('throws an exception when expecting one result', () => {
-      expect(
+    it('throws an exception when expecting one result', async () => {
+      await expect(
         ds.testGet({ expected: 'one' })
       ).rejects.toThrowErrorMatchingInlineSnapshot(`"Resource not found."`)
     })
@@ -531,7 +549,7 @@ describe('DBDataSource', () => {
     })
   })
 
-  describe.only('transaction support', () => {
+  describe('transaction support', () => {
     it('can use transactions successfully', async () => {
       const row: DummyRowType = createRow({
         id: 34,
@@ -610,4 +628,163 @@ describe('DBDataSource', () => {
       ).toBeFalsy()
     })
   })
+
+  describe('json column schemas', () => {
+    let ds: TestDataSource
+
+    beforeEach(() => {
+      ds = new TestDataSource({
+        jsonbTest: z.object({ a: z.number() }),
+      })
+    })
+
+    describe('when inserting', () => {
+      describe('one row', () => {
+        describe('with a valid schema', () => {
+          it('throws no errors', async () => {
+            await expect(
+              ds.testInsert(createRow({ id: 34, jsonbTest: { a: 2222 } }))
+            ).resolves.not.toThrowError()
+          })
+
+          it('inserts the row', async () => {
+            const result = await ds.testInsert(
+              createRow({ id: 34, jsonbTest: { a: 2222 } })
+            )
+            expect(result.id).toEqual(34)
+          })
+        })
+
+        describe('with an invalid schema', () => {
+          it('throws an error', async () => {
+            await expect(
+              // @ts-expect-error testing schema parsing
+              ds.testInsert(createRow({ id: 36, jsonbTest: { a: 'asdf' } }))
+            ).rejects.toThrowError('Expected number, received string')
+          })
+
+          it("doesn't insert the row", async () => {
+            try {
+              await ds.testInsert(
+                // @ts-expect-error testing schema parsing
+                createRow({ id: 36, jsonbTest: { a: 'asdf' } })
+              )
+            } catch {
+              // swallow the error
+            }
+            const result = await ds.testGet()
+            expect(result.length).toEqual(0)
+          })
+        })
+      })
+      describe('multiple rows', () => {
+        describe('with all valid schemas', () => {
+          it('throws no errors', async () => {
+            await expect(
+              ds.testInsert([
+                createRow({ id: 37, jsonbTest: { a: 2222 } }),
+                createRow({ id: 38, jsonbTest: { a: 2223 } }),
+              ])
+            ).resolves.not.toThrowError()
+          })
+
+          it('inserts the rows', async () => {
+            const row1 = createRow({ id: 39, jsonbTest: { a: 2222 } })
+            const row2 = createRow({ id: 40, jsonbTest: { a: 2223 } })
+            const result = await ds.testInsert([row1, row2])
+            expect(result.length).toEqual(2)
+            expect(result).toContainEqual(
+              expect.objectContaining({ id: row1.id })
+            )
+            expect(result).toContainEqual(
+              expect.objectContaining({ id: row2.id })
+            )
+          })
+        })
+
+        describe('with at least one invalid schema', () => {
+          it('throws an error', async () => {
+            await expect(
+              ds.testInsert([
+                // @ts-expect-error testing schema parsing
+                createRow({ id: 41, jsonbTest: { a: 'asdf' } }),
+                createRow({ id: 42, jsonbTest: { a: 2223 } }),
+              ])
+            ).rejects.toThrowError('Expected number, received string')
+          })
+
+          it('inserts no rows', async () => {
+            try {
+              await ds.testInsert([
+                // @ts-expect-error testing schema parsing
+                createRow({ id: 43, jsonbTest: { a: 'asdf' } }),
+                createRow({ id: 44, jsonbTest: { a: 2223 } }),
+              ])
+            } catch {
+              // swallow the error
+            }
+            const result = await ds.testGet({ expected: 'any' })
+            expect(result.length).toBe(0)
+          })
+        })
+      })
+    })
+
+    describe('when updating', () => {
+      const originalValue = 2222
+
+      beforeEach(async () => {
+        const row = createRow({ id: 39, jsonbTest: { a: originalValue } })
+        await ds.testInsert(row)
+      })
+      describe('and given a valid schema', () => {
+        it('throws no errors and inserts correctly', async () => {
+          await expect(
+            ds.testUpdate({ jsonbTest: { a: 1234 } })
+          ).resolves.not.toThrowError()
+        })
+
+        it('updates the row', async () => {
+          const newValue = 9999
+          const result = await ds.testUpdate(
+            { jsonbTest: { a: newValue } },
+            {
+              where: { id: 39 },
+              expected: 'one',
+            }
+          )
+          expect(result.jsonbTest.a).toEqual(newValue)
+        })
+      })
+
+      describe('and given an invalid schema', () => {
+        it('throws an error', async () => {
+          await expect(
+            // @ts-expect-error testing schema parsing
+            ds.testUpdate({ jsonbTest: { a: 'asdf' } })
+          ).rejects.toThrowError('Expected number, received string')
+        })
+
+        it("doesn't update the row", async () => {
+          try {
+            await ds.testUpdate(
+              // @ts-expect-error testing schema parsing
+              { jsonbTest: { a: 'asdf' } },
+              {
+                where: { id: 39 },
+                expected: 'one',
+              }
+            )
+          } catch {
+            // swallow the error
+          }
+          const result = await ds.testGet({
+            where: { id: 39 },
+            expected: 'one',
+          })
+          expect(result.jsonbTest.a).toEqual(originalValue)
+        })
+      })
+    })
+  })
 })

src/datasource/queries/types.ts

@@ -41,11 +41,7 @@ export type GenericSet = Record<
 >
 
 export type UpdateSet<TRowType> = {
-  [K in keyof TRowType]?:
-    | (TRowType[K] extends SerializableValueType | undefined
-        ? TRowType[K]
-        : never)
-    | SqlToken
+  [K in keyof TRowType]?: TRowType[K] | SqlToken
 } & GenericSet
 
 export type ColumnListEntry = string | IdentifierSqlToken | SqlSqlToken