fix(drizzle): cast jsonb selector params to eql_v2_encrypted

Author: tobyhede

packages/drizzle/__tests__/operators-jsonb.test.ts

@@ -0,0 +1,83 @@
+import type { ProtectClient } from '@cipherstash/protect/client'
+import { PgDialect, pgTable } from 'drizzle-orm/pg-core'
+import { describe, expect, it, vi } from 'vitest'
+import { createProtectOperators, encryptedType } from '../src/pg'
+
+const docsTable = pgTable('json_docs', {
+  metadata: encryptedType<Record<string, unknown>>('metadata', {
+    dataType: 'json',
+    searchableJson: true,
+  }),
+})
+
+function createMockProtectClient() {
+  const encryptedSelector = '{"v":"encrypted-selector"}'
+  const encryptQuery = vi.fn(async (terms: unknown[]) => ({
+    data: terms.map(() => encryptedSelector),
+  }))
+
+  return {
+    client: { encryptQuery } as unknown as ProtectClient,
+    encryptQuery,
+    encryptedSelector,
+  }
+}
+
+describe('createProtectOperators JSONB selector typing', () => {
+  it('casts jsonbPathQueryFirst selector params to eql_v2_encrypted', async () => {
+    const { client, encryptQuery, encryptedSelector } =
+      createMockProtectClient()
+    const protectOps = createProtectOperators(client)
+    const dialect = new PgDialect()
+
+    const condition = await protectOps.jsonbPathQueryFirst(
+      docsTable.metadata,
+      '$.profile.email',
+    )
+    const query = dialect.sqlToQuery(condition)
+
+    expect(query.sql).toMatch(
+      /eql_v2\.jsonb_path_query_first\([^,]+,\s*\$\d+::eql_v2_encrypted\)/,
+    )
+    expect(query.params).toHaveLength(1)
+    expect(typeof query.params[0]).toBe('string')
+    expect(encryptQuery).toHaveBeenCalledTimes(1)
+    expect(encryptQuery.mock.calls[0]?.[0]).toMatchObject([
+      { queryType: 'steVecSelector' },
+    ])
+  })
+
+  it('casts jsonbPathExists selector params to eql_v2_encrypted', async () => {
+    const { client } = createMockProtectClient()
+    const protectOps = createProtectOperators(client)
+    const dialect = new PgDialect()
+
+    const condition = await protectOps.jsonbPathExists(
+      docsTable.metadata,
+      '$.profile.email',
+    )
+    const query = dialect.sqlToQuery(condition)
+
+    expect(query.sql).toMatch(
+      /eql_v2\.jsonb_path_exists\([^,]+,\s*\$\d+::eql_v2_encrypted\)/,
+    )
+    expect(query.params).toHaveLength(1)
+    expect(typeof query.params[0]).toBe('string')
+  })
+
+  it('casts jsonbGet selector params to eql_v2_encrypted', async () => {
+    const { client } = createMockProtectClient()
+    const protectOps = createProtectOperators(client)
+    const dialect = new PgDialect()
+
+    const condition = await protectOps.jsonbGet(
+      docsTable.metadata,
+      '$.profile.email',
+    )
+    const query = dialect.sqlToQuery(condition)
+
+    expect(query.sql).toMatch(/->\s*\$\d+::eql_v2_encrypted/)
+    expect(query.params).toHaveLength(1)
+    expect(typeof query.params[0]).toBe('string')
+  })
+})

packages/drizzle/src/pg/operators.ts

@@ -882,6 +882,8 @@ function createJsonbOperator(
   protectTableCache: Map<string, ProtectTable<ProtectTableColumn>>,
 ): Promise<SQL> {
   const { config } = columnInfo
+  const encryptedSelector = (value: unknown) =>
+    sql`${bindIfParam(value, left)}::eql_v2_encrypted`
 
   if (!config?.searchableJson) {
     throw new ProtectOperatorError(
@@ -907,11 +909,11 @@ function createJsonbOperator(
     }
     switch (operator) {
       case 'jsonbPathQueryFirst':
-        return sql`eql_v2.jsonb_path_query_first(${left}, ${bindIfParam(encrypted, left)})`
+        return sql`eql_v2.jsonb_path_query_first(${left}, ${encryptedSelector(encrypted)})`
       case 'jsonbGet':
-        return sql`${left} -> ${bindIfParam(encrypted, left)}`
+        return sql`${left} -> ${encryptedSelector(encrypted)}`
       case 'jsonbPathExists':
-        return sql`eql_v2.jsonb_path_exists(${left}, ${bindIfParam(encrypted, left)})`
+        return sql`eql_v2.jsonb_path_exists(${left}, ${encryptedSelector(encrypted)})`
     }
   }