Skip to content

Commit 62df494

Browse files
authored
fix(stack-drizzle): preserve per-column types in extractEncryptionSchema (#845)
* fix(stack-drizzle): preserve per-column types in extractEncryptionSchema `extractEncryptionSchema` returned the widened `AnyV3Table`, discarding every column's concrete EQL v3 domain. `InferPlaintext` over an extracted schema therefore collapsed to an index signature, so along the documented flow — `extractEncryptionSchema(table)` -> `Encryption({ schemas })` -> `bulkEncryptModels` — a `string` written to an `IntegerOrd` column type-checked, and plain helper columns (`id`, a plain `text()`) were typed as encrypted. Type-safety/DX only: the runtime already recovered each builder correctly via `getEqlV3Column`. Three linked pieces, all in `@cipherstash/stack-drizzle` — the v3 core needed no change: - `makeEqlV3Column` carries its builder `C` on a phantom brand. The brand rides on the column's `_` config, not the builder object: `pgTable()` rebuilds each builder through drizzle's `BuildColumn`, which keeps only `Omit<TBuilder['_'], keyof MakeColumnConfig<…>>` — drizzle's own extension point (`NotNull<T> = T & { _: { notNull: true } }`). A brand on the object would be dropped and extraction would have nothing to read. - `extractEncryptionSchema` is generic over the table and returns `EncryptedTable<Cols> & Cols`, the shape `encryptedTable()` itself returns. `EncryptedColumnsOf<T>` recovers each brand and keeps only the branded columns — the type-level mirror of the existing runtime loop, which is unchanged. - `V3BuilderOf` short-circuits `any` to `AnyEncryptedV3Column`, so an untyped table keeps the widened-but-usable typing it had before rather than branding every column `any`. `__tests__/schema-extraction.test-d.ts` locks both paths: extracted and hand-authored must infer the same plaintext map, reject a field typed against the wrong domain, and pin `encrypt` to the addressed column. Six of its eight assertions fail without this fix. Two integration files relied on the old looseness and are corrected, not worked around: the dynamically-built matrix table names its model type explicitly (its rows are a pure index signature, so inference dropped the `| null` that `nullableTextEq` needs), and the bench comment claiming extraction returns `AnyV3Table` is now false and says so. Closes #589 * fix(stack-drizzle): retain widened extraction types * test(stack-drizzle): cover modifier survival and the json domain Two gaps in the #589 type-level suite, both of which would have shipped green. Drizzle column modifiers. Every existing fixture declares bare columns, but real schemas almost always use `.notNull()`. The brand survives `pgTable()` only because drizzle's `BuildColumn` keeps `Omit<TBuilder['_'], keyof MakeColumnConfig<…> | 'brand' | 'dialect'>`, and modifiers then intersect further into `_` (`NotNull<T> = T & { _: { notNull: true } }`). That is an upstream compatibility contract, not one we control: a drizzle-orm change to either step could drop the brand for modified columns while bare ones kept working. The runtime would not notice — `getEqlV3Column` recovers builders from the runtime carrier and the SQL domain, never from the type — so extraction would silently type a `.notNull()` column as a plain pass-through field while still encrypting it. Verified the guard isolates that class: narrowing `V3BuilderOf` to match only `notNull: false` fails exactly the three new modifier assertions and leaves all fourteen pre-existing ones green. The json domain. `types.Json` is the only domain whose plaintext is a union (`JsonDocument`) rather than a scalar, which makes it the easiest to lose to a widened index signature without another assertion noticing. Covered on both the extracted and hand-authored paths, including a `@ts-expect-error` on a scalar written to a json column. Both blocks fail against the pre-fix source, so they cover the fix rather than passing incidentally. Test-only; the existing changeset already covers the behaviour. * fix(stack-drizzle): widen mixed extracted schemas
1 parent 20cb8c3 commit 62df494

7 files changed

Lines changed: 564 additions & 28 deletions

File tree

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
---
2+
'@cipherstash/stack-drizzle': minor
3+
'stash': patch
4+
---
5+
6+
Type `extractEncryptionSchema` precisely: a Drizzle-extracted schema now preserves each column's concrete EQL v3 domain instead of widening to `AnyV3Table` (#589).
7+
8+
`extractEncryptionSchema` is generic over the Drizzle table (`<T extends PgTable>(table: T)`) and returns `EncryptedTable<Cols> & Cols`, the same shape a hand-written `encryptedTable({...})` returns, when concrete column brands are available. Each column's builder is carried through `pgTable()` on a phantom brand and recovered by a mapped type, which also filters out the table's non-encrypted columns. Tables widened to `PgTable`, and tables containing ordinary `customType` columns recovered from their EQL SQL domain, retain the safe `AnyV3Table` fallback instead of incorrectly becoming an empty or partial schema type.
9+
10+
What this fixes, along the documented flow `extractEncryptionSchema(table)` → `Encryption({ schemas })` → `bulkEncryptModels`:
11+
12+
- `InferPlaintext<typeof schema>` is a precise per-column plaintext map (`{ email: string; age: number }`) rather than an index signature.
13+
- `encryptModel` / `bulkEncryptModels` check each schema field against its own domain's plaintext — a `string` written to an `IntegerOrd` column is now a compile error instead of an encrypt-time failure — and pass plain helper columns (`id`, a plain `text()`) through with their own types rather than typing them as encrypted.
14+
- `schema.email` addresses the column at its concrete type, so `encrypt` / `encryptQuery` pin the value to that column's plaintext.
15+
16+
**Runtime behaviour is unchanged** — the runtime already recovered each column's builder correctly, so this is a type-level fix only. It is `minor` rather than `patch` because code that previously compiled against the widened types can now fail to compile: a model field typed against the wrong domain, or a schema-derived type that relied on the old index signature. Rows whose shape is only known at runtime (a dynamically built table) should name their model type explicitly — `client.bulkEncryptModels<typeof schema, MyRow>(rows, schema)` — rather than being cast back to `AnyV3Table`.
17+
18+
`skills/stash-drizzle` documents the preserved typing and warns against casting an extracted schema to `AnyV3Table` to make an insert compile. A matching update to the separately maintained CipherStash documentation site is required so its Drizzle schema-extraction guidance explains the precise branded typing and the widened fallback for incomplete runtime-recovered column maps.

‎packages/bench/src/drizzle/setup.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,11 @@ export const encryptionBenchTable = extractEncryptionSchema(benchTable)
4141
* so `__unit__/seed-keys.test.ts` is the real drift guard, not the type.
4242
* Drizzle's `InferInsertModel` describes the ENCRYPTED column (the custom type's
4343
* `data` is the EQL envelope, not the plaintext) and degrades these three to
44-
* optional `any`. Deriving from `encryptionBenchTable` is no better:
45-
* `extractEncryptionSchema` returns the widened `AnyV3Table`, whose column map
46-
* is an index signature — a type that admits `encTxt: 'x'` and `encText: 12345`
47-
* alike, both of which the literal below rejects.
44+
* optional `any`. `InferPlaintext<typeof encryptionBenchTable>` is precise per
45+
* column since #589, and it does pin `encText`/`encInt` — but `encJsonb` is a
46+
* `types.Json` column, whose plaintext is the whole `JsonDocument` union. The
47+
* literal below says `{ idx: number; group: number }`, which is what the bench
48+
* queries actually assume, so deriving would still lose the shape that matters.
4849
*
4950
* Update it by hand when `benchTable` changes; the unit test will tell you.
5051
*/
Lines changed: 350 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,350 @@
1+
/**
2+
* Type-level contract for {@link extractEncryptionSchema} (#589).
3+
*
4+
* A Drizzle-extracted schema must be typed as precisely as the hand-authored
5+
* `encryptedTable({…})` it is equivalent to. Before #589 it was not: extraction
6+
* returned the widened `AnyV3Table`, so `InferPlaintext` over an extracted
7+
* schema collapsed to an index signature and every column on the insert path
8+
* accepted (and produced) the union of every domain's plaintext instead of its
9+
* own. There was no runtime bug — the runtime already recovers each column's
10+
* concrete builder — which is exactly why this needs a type test to hold it.
11+
*
12+
* Every assertion below is written to be equally meaningful on BOTH paths, so
13+
* the hand-authored surface cannot regress while the extracted one is fixed.
14+
*/
15+
16+
import {
17+
type AnyV3Table,
18+
type EncryptedIntegerOrdColumn,
19+
type EncryptedJsonColumn,
20+
type EncryptedTextEqColumn,
21+
encryptedTable,
22+
type InferPlaintext,
23+
type JsonDocument,
24+
types as v3Types,
25+
} from '@cipherstash/stack/eql/v3'
26+
import type { Encrypted } from '@cipherstash/stack/types'
27+
import { Encryption, type EncryptionClient } from '@cipherstash/stack/v3'
28+
import {
29+
customType,
30+
type PgTable,
31+
pgTable,
32+
serial,
33+
text,
34+
} from 'drizzle-orm/pg-core'
35+
import { describe, expectTypeOf, it } from 'vitest'
36+
import { extractEncryptionSchema } from '../src/schema-extraction'
37+
import { types } from '../src/types'
38+
39+
// `id` and `nickname` are load-bearing, not scenery: extraction must keep only
40+
// the encrypted (branded) columns, so the mapped type has to filter them out.
41+
// They also cover the other half of the insert contract — a plain helper column
42+
// must pass through the model types untouched rather than being treated as an
43+
// encrypted one.
44+
const users = pgTable('users', {
45+
id: serial('id').primaryKey(),
46+
nickname: text('nickname'),
47+
email: types.TextEq('email'),
48+
age: types.IntegerOrd('age'),
49+
})
50+
51+
/** The hand-authored equivalent of `users`. Same columns, same domains. */
52+
const authored = encryptedTable('users', {
53+
email: v3Types.TextEq('email'),
54+
age: v3Types.IntegerOrd('age'),
55+
})
56+
57+
const extracted = extractEncryptionSchema(users)
58+
59+
// Consumers can legitimately erase a table's concrete column shape at module
60+
// boundaries. Runtime extraction still sees the columns in this case, but the
61+
// private type brand is no longer recoverable.
62+
const widenedUsers: PgTable = users
63+
const widenedExtracted = extractEncryptionSchema(widenedUsers)
64+
65+
// Runtime extraction also supports ordinary Drizzle custom columns by their
66+
// EQL SQL domain, even though those columns never carried our private brand.
67+
const eqlTextEq = customType<{ data: Encrypted }>({
68+
dataType: () => 'eql_v3_text_eq',
69+
})
70+
const sqlTypeFallbackUsers = pgTable('sql_type_fallback_users', {
71+
email: eqlTextEq('email'),
72+
})
73+
const fallbackExtracted = extractEncryptionSchema(sqlTypeFallbackUsers)
74+
75+
// A single visible brand must not make the schema look complete when another
76+
// runtime-encrypted column is only recoverable from its SQL domain.
77+
const mixedUsers = pgTable('mixed_users', {
78+
email: eqlTextEq('email'),
79+
age: types.IntegerOrd('age'),
80+
})
81+
const mixedExtracted = extractEncryptionSchema(mixedUsers)
82+
83+
// Real schemas rarely declare a bare column — `.notNull()` in particular is
84+
// near-universal. Modifiers are the one place the brand could be dropped for
85+
// SOME columns while bare ones keep working, which is why they get their own
86+
// fixture: every other table here is bare, so a modifier-only regression would
87+
// otherwise ship green. See the `modifiers` describe below.
88+
const modified = pgTable('modified', {
89+
bare: types.TextEq('bare'),
90+
notNull: types.TextEq('not_null').notNull(),
91+
unique: types.IntegerOrd('uniq').unique(),
92+
chained: types.IntegerOrd('chained').notNull().unique(),
93+
})
94+
const modifiedExtracted = extractEncryptionSchema(modified)
95+
96+
// `types.Json`'s plaintext is the whole `JsonDocument` union rather than a
97+
// scalar, so it exercises the one domain where a widened index signature and
98+
// the precise type are easiest to confuse.
99+
const documents = pgTable('documents', {
100+
id: serial('id').primaryKey(),
101+
doc: types.Json('doc'),
102+
label: types.TextEq('label'),
103+
})
104+
const documentsExtracted = extractEncryptionSchema(documents)
105+
const documentsAuthored = encryptedTable('documents', {
106+
doc: v3Types.Json('doc'),
107+
label: v3Types.TextEq('label'),
108+
})
109+
110+
/** The plaintext map both schemas must infer. */
111+
type UserPlaintext = { email: string; age: number }
112+
113+
/** The model handed to `bulkEncryptModels`, plain helper columns included. */
114+
const plainRows = [
115+
{ id: 1, nickname: 'ada', email: 'ada@example.com', age: 36 },
116+
]
117+
118+
/**
119+
* What `bulkEncryptModels` must resolve: the two encrypted columns become
120+
* envelopes, the two plain ones pass through with their own types.
121+
*/
122+
type EncryptedUserRow = {
123+
id: number
124+
nickname: string
125+
email: Encrypted
126+
age: Encrypted
127+
}
128+
129+
describe('extractEncryptionSchema - plaintext inference (#589)', () => {
130+
it('infers the same precise plaintext map as the hand-authored table', () => {
131+
expectTypeOf<
132+
InferPlaintext<typeof authored>
133+
>().toEqualTypeOf<UserPlaintext>()
134+
expectTypeOf<
135+
InferPlaintext<typeof extracted>
136+
>().toEqualTypeOf<UserPlaintext>()
137+
})
138+
139+
it('drops the table’s plain, non-encrypted columns', () => {
140+
// The `toEqualTypeOf` above already forbids extra keys; this states the
141+
// filtering intent directly so a failure names the actual cause.
142+
expectTypeOf<keyof InferPlaintext<typeof extracted>>().toEqualTypeOf<
143+
'email' | 'age'
144+
>()
145+
})
146+
147+
it('keeps each column addressable at its concrete domain type', () => {
148+
// `encryptedTable()` returns `EncryptedTable<C> & C`, so the schema doubles
149+
// as a column accessor. Extraction returns the same intersection — that is
150+
// what lets `client.encrypt(v, { table, column: schema.email })` pin `v` to
151+
// the column's own plaintext rather than the union of every domain's.
152+
expectTypeOf(extracted.email).toEqualTypeOf<EncryptedTextEqColumn>()
153+
expectTypeOf(extracted.age).toEqualTypeOf<EncryptedIntegerOrdColumn>()
154+
expectTypeOf(authored.email).toEqualTypeOf<EncryptedTextEqColumn>()
155+
expectTypeOf(authored.age).toEqualTypeOf<EncryptedIntegerOrdColumn>()
156+
})
157+
})
158+
159+
/**
160+
* Regression guard for the brand's survival through Drizzle's column modifiers.
161+
*
162+
* `pgTable()` does not hand builders through — it rebuilds each one via
163+
* drizzle's `BuildColumn`, which keeps only
164+
* `Omit<TBuilder['_'], keyof MakeColumnConfig<…> | 'brand' | 'dialect'>`. The
165+
* brand rides in `_` precisely so it survives that (see `EqlV3Column`), and
166+
* modifiers then intersect further into `_` (`NotNull<T> = T & { _: { notNull:
167+
* true } }`). Both steps have to preserve the symbol key.
168+
*
169+
* This is a drizzle-orm compatibility contract, not one we control: an upstream
170+
* change to how `BuildColumn` or `NotNull` reshape `_` could drop the brand for
171+
* modified columns while bare ones keep working. The runtime would not notice —
172+
* `getEqlV3Column` recovers builders from the runtime carrier and the SQL domain,
173+
* never from the type — so extraction would silently type a `.notNull()` column
174+
* as a plain pass-through field while still encrypting it. Nothing else in CI
175+
* catches that.
176+
*/
177+
describe('extractEncryptionSchema - drizzle column modifiers', () => {
178+
it('keeps modified columns in the extracted schema', () => {
179+
expectTypeOf<
180+
keyof InferPlaintext<typeof modifiedExtracted>
181+
>().toEqualTypeOf<'bare' | 'notNull' | 'unique' | 'chained'>()
182+
})
183+
184+
it('keeps each modified column at its own concrete plaintext type', () => {
185+
expectTypeOf<InferPlaintext<typeof modifiedExtracted>>().toEqualTypeOf<{
186+
bare: string
187+
notNull: string
188+
unique: number
189+
chained: number
190+
}>()
191+
})
192+
193+
it('keeps modified columns addressable at their concrete domain type', () => {
194+
expectTypeOf(
195+
modifiedExtracted.notNull,
196+
).toEqualTypeOf<EncryptedTextEqColumn>()
197+
expectTypeOf(
198+
modifiedExtracted.chained,
199+
).toEqualTypeOf<EncryptedIntegerOrdColumn>()
200+
})
201+
})
202+
203+
/**
204+
* `types.Json` is the only domain whose plaintext is a union rather than a
205+
* scalar, so it is the easiest one to lose to a widened index signature without
206+
* any other assertion noticing.
207+
*/
208+
describe('extractEncryptionSchema - json domain', () => {
209+
it('infers JsonDocument on both paths', () => {
210+
type DocumentPlaintext = { doc: JsonDocument; label: string }
211+
212+
expectTypeOf<
213+
InferPlaintext<typeof documentsExtracted>
214+
>().toEqualTypeOf<DocumentPlaintext>()
215+
expectTypeOf<
216+
InferPlaintext<typeof documentsAuthored>
217+
>().toEqualTypeOf<DocumentPlaintext>()
218+
})
219+
220+
it('keeps the json column addressable at its concrete domain type', () => {
221+
expectTypeOf(documentsExtracted.doc).toEqualTypeOf<EncryptedJsonColumn>()
222+
expectTypeOf(
223+
documentsExtracted.label,
224+
).toEqualTypeOf<EncryptedTextEqColumn>()
225+
})
226+
227+
it('rejects a scalar written to a json column', async () => {
228+
const client = await Encryption({ schemas: [documentsExtracted] })
229+
230+
expectTypeOf(client.encrypt).toBeCallableWith(
231+
{ note: 'ok' },
232+
{ table: documentsExtracted, column: documentsExtracted.doc },
233+
)
234+
// @ts-expect-error - JsonDocument is object | array | null, never a scalar
235+
client.encrypt(36, {
236+
table: documentsExtracted,
237+
column: documentsExtracted.doc,
238+
})
239+
})
240+
})
241+
242+
describe('extractEncryptionSchema - unbranded table fallback', () => {
243+
it('preserves the widened schema type for a table annotated as PgTable', () => {
244+
expectTypeOf(widenedExtracted).toEqualTypeOf<AnyV3Table>()
245+
expectTypeOf<keyof InferPlaintext<typeof widenedExtracted>>().toEqualTypeOf<
246+
string | number
247+
>()
248+
})
249+
250+
it('preserves the widened schema type for columns recovered by SQL type', () => {
251+
expectTypeOf(fallbackExtracted).toEqualTypeOf<AnyV3Table>()
252+
expectTypeOf<
253+
keyof InferPlaintext<typeof fallbackExtracted>
254+
>().toEqualTypeOf<string | number>()
255+
})
256+
257+
it('still types fields as encrypted through widened model operations', async () => {
258+
const client = await Encryption({ schemas: [widenedExtracted] })
259+
const result = await client.bulkEncryptModels(
260+
[{ email: 'ada@example.com' }],
261+
widenedExtracted,
262+
)
263+
if (result.failure) throw result.failure
264+
265+
expectTypeOf(result.data).toEqualTypeOf<Array<{ email: Encrypted }>>()
266+
})
267+
268+
it('widens a mixed branded and SQL-derived schema as a complete unit', async () => {
269+
expectTypeOf(mixedExtracted).toEqualTypeOf<AnyV3Table>()
270+
271+
const client = await Encryption({ schemas: [mixedExtracted] })
272+
const result = await client.bulkEncryptModels(
273+
[{ email: 'ada@example.com', age: 36 }],
274+
mixedExtracted,
275+
)
276+
if (result.failure) throw result.failure
277+
278+
expectTypeOf(result.data).toEqualTypeOf<
279+
Array<{ email: Encrypted; age: Encrypted }>
280+
>()
281+
})
282+
})
283+
284+
/**
285+
* The documented Drizzle flow, end to end:
286+
* `extractEncryptionSchema(table)` -> `Encryption({ schemas })` -> `bulkEncryptModels`.
287+
*
288+
* This is the path the issue reports as widened, so it is asserted on the real
289+
* factory rather than on a hand-named client type — the whole point is that the
290+
* concrete tuple survives from extraction through to the model types.
291+
*
292+
* These `it` bodies are type-checked, never executed: vitest's typecheck mode
293+
* does not run `.test-d.ts` files, and they are outside the runtime suite's
294+
* `include`. So `Encryption()` here needs no credentials.
295+
*/
296+
describe('extractEncryptionSchema - documented flow (#589)', () => {
297+
it('carries the extracted schema into the client tuple', async () => {
298+
const client = await Encryption({ schemas: [extracted] })
299+
expectTypeOf(client).toEqualTypeOf<
300+
EncryptionClient<readonly [typeof extracted]>
301+
>()
302+
})
303+
304+
it('types bulkEncryptModels precisely against an extracted schema', async () => {
305+
const client = await Encryption({ schemas: [extracted] })
306+
const result = await client.bulkEncryptModels(plainRows, extracted)
307+
if (result.failure) throw result.failure
308+
309+
expectTypeOf(result.data).toEqualTypeOf<EncryptedUserRow[]>()
310+
})
311+
312+
it('types bulkEncryptModels precisely against the hand-authored schema', async () => {
313+
const client = await Encryption({ schemas: [authored] })
314+
const result = await client.bulkEncryptModels(plainRows, authored)
315+
if (result.failure) throw result.failure
316+
317+
expectTypeOf(result.data).toEqualTypeOf<EncryptedUserRow[]>()
318+
})
319+
320+
it('rejects a field typed against the wrong domain on either path', async () => {
321+
const onExtracted = await Encryption({ schemas: [extracted] })
322+
const onAuthored = await Encryption({ schemas: [authored] })
323+
const wrong = [{ email: 'ada@example.com', age: 'thirty-six' }]
324+
325+
// This is the assertion with teeth. While extraction returned `AnyV3Table`,
326+
// every column's plaintext was the union of every domain's, so a `string`
327+
// in an `IntegerOrd` column type-checked and the mismatch only surfaced at
328+
// encrypt time. If that widening ever returns, the suppression below goes
329+
// unused and this test fails.
330+
// @ts-expect-error - `age` is IntegerOrd: its plaintext is number, not string
331+
onExtracted.bulkEncryptModels(wrong, extracted)
332+
// @ts-expect-error - the hand-authored path has always rejected this; it must keep doing so
333+
onAuthored.bulkEncryptModels(wrong, authored)
334+
})
335+
336+
it('pins encrypt() to the addressed column’s own plaintext type', async () => {
337+
const client = await Encryption({ schemas: [extracted] })
338+
339+
expectTypeOf(client.encrypt).toBeCallableWith('ada@example.com', {
340+
table: extracted,
341+
column: extracted.email,
342+
})
343+
expectTypeOf(client.encrypt).toBeCallableWith(36, {
344+
table: extracted,
345+
column: extracted.age,
346+
})
347+
// @ts-expect-error - the email column encrypts a string, not a number
348+
client.encrypt(36, { table: extracted, column: extracted.email })
349+
})
350+
})

0 commit comments

Comments
 (0)