Skip to content

Commit 0c51662

Browse files
0xcadams0xcadams
authored
Merge pull request #128 from IderAghbal/main
feat: rework table and column inclusion to be an allow-list
2 parents 2bd150e + 7a995a7 commit 0c51662

4 files changed

Lines changed: 188 additions & 69 deletions

File tree

‎src/relations.ts‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ type DrizzleToZeroSchema<
278278
ZeroTableBuilderSchema<
279279
K & string,
280280
TDrizzleSchema[K],
281-
TColumnConfig[K],
281+
NonNullable<TColumnConfig[K]>,
282282
TCasing
283283
>
284284
>
@@ -496,11 +496,17 @@ const drizzleZeroConfig = <
496496

497497
const tableConfig = config?.tables?.[tableName as keyof TColumnConfig];
498498

499-
// skip tables that don't have a config
500-
if (tableConfig === false) {
499+
if (
500+
config?.tables !== undefined &&
501+
(tableConfig === false || tableConfig === undefined)
502+
) {
501503
debugLog(
502504
config?.debug,
503-
`Skipping table ${String(tableName)} - no config provided`,
505+
`Skipping table ${String(tableName)} - ${
506+
tableConfig === false
507+
? "explicitly excluded"
508+
: "not mentioned in config"
509+
}`,
504510
);
505511
continue;
506512
}

‎src/tables.ts‎

Lines changed: 69 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -56,20 +56,22 @@ type TypeOverride<TCustomType> = {
5656
* @template TTable The Drizzle table type
5757
*/
5858
export type ColumnsConfig<TTable extends Table> =
59-
| false
60-
| Flatten<{
61-
/**
62-
* The columns to include in the Zero schema.
63-
* Set to true to use default mapping, or provide a TypeOverride for custom mapping.
64-
*/
65-
readonly [KColumn in ColumnNames<TTable>]:
66-
| boolean
67-
| ColumnBuilder<
68-
TypeOverride<
69-
ZeroTypeToTypescriptType[DrizzleDataTypeToZeroType[Columns<TTable>[KColumn]["dataType"]]]
70-
>
71-
>;
72-
}>;
59+
| boolean
60+
| Partial<
61+
Flatten<{
62+
/**
63+
* The columns to include in the Zero schema.
64+
* Set to true to use default mapping, or provide a TypeOverride for custom mapping.
65+
*/
66+
readonly [KColumn in ColumnNames<TTable>]:
67+
| boolean
68+
| ColumnBuilder<
69+
TypeOverride<
70+
ZeroTypeToTypescriptType[DrizzleDataTypeToZeroType[Columns<TTable>[KColumn]["dataType"]]]
71+
>
72+
>;
73+
}>
74+
>;
7375

7476
/**
7577
* Maps a Drizzle column type to its corresponding Zero type.
@@ -262,20 +264,39 @@ const createZeroTableBuilder = <
262264
const tableColumns = getTableColumns(table);
263265
const tableConfig = getTableConfigForDatabase(table);
264266

265-
const primaryKeysFromColumns: string[] = [];
266-
267-
const columnsMapped = typedEntries(tableColumns).reduce(
268-
(acc, [key, column]) => {
269-
const columnConfig = columns?.[key as keyof TColumnConfig];
267+
const columnNameToStableKey = new Map<string, string>(
268+
typedEntries(tableColumns).map(([key, column]) => [
269+
column.name,
270+
String(key),
271+
]),
272+
);
270273

271-
if (columnConfig === false) {
272-
debugLog(
273-
debug,
274-
`Skipping column ${String(key)} because columnConfig is false`,
275-
);
274+
const primaryKeys = new Set<string>();
275+
for (const [key, column] of typedEntries(tableColumns)) {
276+
if (column.primary) {
277+
primaryKeys.add(String(key));
278+
}
279+
}
276280

277-
return acc;
281+
for (const pk of tableConfig.primaryKeys) {
282+
for (const pkColumn of pk.columns) {
283+
const key = columnNameToStableKey.get(pkColumn.name);
284+
if (key) {
285+
primaryKeys.add(String(key));
278286
}
287+
}
288+
}
289+
290+
const isColumnBuilder = (value: unknown): value is ColumnBuilder<any> =>
291+
typeof value === "object" && value !== null && "schema" in value;
292+
293+
const columnsMapped = typedEntries(tableColumns).reduce(
294+
(acc, [key, column]) => {
295+
const columnConfig =
296+
typeof columns === "object" && columns !== null
297+
? columns[key as keyof TColumnConfig]
298+
: undefined;
299+
const isColumnConfigOverride = isColumnBuilder(columnConfig);
279300

280301
// From https://github.com/drizzle-team/drizzle-orm/blob/e5c63db0df0eaff5cae8321d97a77e5b47c5800d/drizzle-kit/src/serializer/utils.ts#L5
281302
const resolvedColumnName =
@@ -285,21 +306,30 @@ const createZeroTableBuilder = <
285306
? toCamelCase(column.name)
286307
: toSnakeCase(column.name);
287308

288-
if (
289-
typeof columnConfig !== "boolean" &&
290-
typeof columnConfig !== "object" &&
291-
typeof columnConfig !== "undefined"
292-
) {
293-
throw new Error(
294-
`drizzle-zero: Invalid column config for column ${resolvedColumnName} - expected boolean or ColumnBuilder but was ${typeof columnConfig}`,
295-
);
309+
if (typeof columns === "object" && columns !== null) {
310+
if (
311+
columnConfig !== undefined &&
312+
typeof columnConfig !== "boolean" &&
313+
!isColumnConfigOverride
314+
) {
315+
throw new Error(
316+
`drizzle-zero: Invalid column config for column ${resolvedColumnName} - expected boolean or ColumnBuilder but was ${typeof columnConfig}`,
317+
);
318+
}
319+
320+
if (
321+
columnConfig !== true &&
322+
!isColumnConfigOverride &&
323+
!primaryKeys.has(String(key))
324+
) {
325+
debugLog(
326+
debug,
327+
`Skipping non-primary column ${resolvedColumnName} because it was not explicitly included in the config.`,
328+
);
329+
return acc;
330+
}
296331
}
297332

298-
const isColumnBuilder = (value: unknown): value is ColumnBuilder<any> =>
299-
typeof value === "object" && value !== null && "schema" in value;
300-
301-
const isColumnConfigOverride = isColumnBuilder(columnConfig);
302-
303333
const type =
304334
drizzleColumnTypeToZeroType[
305335
column.columnType as keyof DrizzleColumnTypeToZeroType
@@ -326,10 +356,6 @@ const createZeroTableBuilder = <
326356
? columnConfig.schema.optional
327357
: false;
328358

329-
if (column.primary) {
330-
primaryKeysFromColumns.push(String(key));
331-
}
332-
333359
if (columnConfig && typeof columnConfig !== "boolean") {
334360
return {
335361
...acc,
@@ -362,19 +388,7 @@ const createZeroTableBuilder = <
362388
{} as Record<string, any>,
363389
);
364390

365-
const primaryKeys = [
366-
...primaryKeysFromColumns,
367-
...tableConfig.primaryKeys.flatMap((k) =>
368-
k.columns.map((c) =>
369-
getDrizzleColumnKeyFromColumnName({
370-
columnName: c.name,
371-
table: c.table,
372-
}),
373-
),
374-
),
375-
];
376-
377-
if (!primaryKeys.length) {
391+
if (primaryKeys.size === 0) {
378392
throw new Error(
379393
`drizzle-zero: No primary keys found in table - ${actualTableName}. Did you forget to define a primary key?`,
380394
);

‎src/types.ts‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,18 @@ export type ColumnIndexKeys<TTable extends Table> = {
1414
* @template TDrizzleSchema - The complete Drizzle schema
1515
*/
1616
export type TableColumnsConfig<TDrizzleSchema extends Record<string, unknown>> =
17-
Flatten<{
18-
/**
19-
* The columns to include in the Zero schema.
20-
*/
21-
readonly [K in keyof TDrizzleSchema as TDrizzleSchema[K] extends Table<any>
22-
? K
23-
: never]: TDrizzleSchema[K] extends Table<any>
24-
? ColumnsConfig<TDrizzleSchema[K]>
25-
: never;
26-
}>;
17+
Partial<
18+
Flatten<{
19+
/**
20+
* The columns to include in the Zero schema.
21+
*/
22+
readonly [K in keyof TDrizzleSchema as TDrizzleSchema[K] extends Table<any>
23+
? K
24+
: never]: TDrizzleSchema[K] extends Table<any>
25+
? ColumnsConfig<TDrizzleSchema[K]>
26+
: never;
27+
}>
28+
>;
2729

2830
/**
2931
* A default config type which includes all tables in the Drizzle schema.

‎tests/config.test.ts‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import { describe, expect, test } from "vitest";
2+
import { drizzleZeroConfig } from "../src/relations";
3+
import { pgTable, serial, text, primaryKey } from "drizzle-orm/pg-core";
4+
5+
describe("drizzleZeroConfig with explicit table and column configuration", () => {
6+
const users = pgTable("users", {
7+
id: serial("id").primaryKey(),
8+
name: text("name").notNull(),
9+
email: text("email").notNull(),
10+
phone: text("phone"),
11+
});
12+
13+
const posts = pgTable("posts", {
14+
id: serial("id").primaryKey(),
15+
title: text("title").notNull(),
16+
content: text("content"),
17+
userId: serial("user_id").references(() => users.id),
18+
});
19+
20+
const comments = pgTable("comments", {
21+
id: serial("id").primaryKey(),
22+
content: text("content").notNull(),
23+
postId: serial("post_id").references(() => posts.id),
24+
});
25+
26+
const usersToPosts = pgTable(
27+
"users_to_posts",
28+
{
29+
userId: serial("user_id").references(() => users.id),
30+
postId: serial("post_id").references(() => posts.id),
31+
role: text("role", { enum: ["owner", "editor"] }),
32+
},
33+
(t: any) => ({
34+
pk: primaryKey({ columns: [t.userId, t.postId] }),
35+
}),
36+
);
37+
38+
const drizzleSchema = { users, posts, comments, usersToPosts };
39+
40+
test("should include all tables and columns when no config is provided", () => {
41+
const schema = drizzleZeroConfig(drizzleSchema);
42+
43+
// All tables should be present
44+
expect(Object.keys(schema.tables).length).toBe(4);
45+
expect(new Set(Object.keys(schema.tables))).toStrictEqual(
46+
new Set(["users", "posts", "comments", "usersToPosts"]),
47+
);
48+
49+
// All columns should be present
50+
expect(
51+
new Set(Object.keys((schema.tables as any).users.columns)),
52+
).toStrictEqual(new Set(["id", "name", "email", "phone"]));
53+
expect(
54+
new Set(Object.keys((schema.tables as any).posts.columns)),
55+
).toStrictEqual(new Set(["id", "title", "content", "userId"]));
56+
expect(
57+
new Set(Object.keys((schema.tables as any).comments.columns)),
58+
).toStrictEqual(new Set(["id", "content", "postId"]));
59+
expect(
60+
new Set(Object.keys((schema.tables as any).usersToPosts.columns)),
61+
).toStrictEqual(new Set(["userId", "postId", "role"]));
62+
});
63+
64+
test("should handle explicit table and column configurations", () => {
65+
const schema = drizzleZeroConfig(drizzleSchema, {
66+
tables: {
67+
users: true, // include all columns
68+
usersToPosts: {
69+
userId: false, // will be included anyway because it is part of the primary key
70+
// role is not mentioned, should be excluded
71+
},
72+
posts: false,
73+
// comments table is not mentioned, should be excluded
74+
},
75+
});
76+
77+
// `users` and `usersToPosts` should be in the schema
78+
expect(Object.keys(schema.tables).length).toBe(2);
79+
expect(new Set(Object.keys(schema.tables))).toStrictEqual(
80+
new Set(["users", "usersToPosts"]),
81+
);
82+
83+
// `posts` and `comments` should be excluded
84+
expect((schema.tables as any).posts).toBe(undefined);
85+
expect((schema.tables as any).comments).toBe(undefined);
86+
87+
// `users` table should have `id` (pk) and `name`
88+
expect(
89+
new Set(Object.keys((schema.tables as any).users.columns)),
90+
).toStrictEqual(new Set(["id", "name", "email", "phone"]));
91+
92+
// `usersToPosts` table should have all its columns
93+
expect(
94+
new Set(Object.keys((schema.tables as any).usersToPosts.columns)),
95+
).toStrictEqual(new Set(["userId", "postId"]));
96+
});
97+
});

0 commit comments

Comments
 (0)