Skip to content

Commit d1ac398

Browse files
waltosswaltoss
committed
Add custom SQL builder functions for relation columns
Allow passing a function to `from`/`to` in relation config that receives the aliased table and returns custom SQL. This enables type casting, JSON extraction, and other SQL transformations in relation joins. Example usage: ```ts r.one.addresses({ from: (table) => sql`${table}.${sql.identifier('id')}::varchar`, to: r.addresses.addressableId, }) ``` Changes: - Add RelationColumnBuilder type for builder functions - Update OneConfig/ManyConfig to accept functions - Update One/Many constructors to handle functions - Add buildColumnSQL helper in relationToSQL - Add integration tests for type casting and JSON extraction
1 parent 05230d9 commit d1ac398

File tree

2 files changed

+199
-37
lines changed

2 files changed

+199
-37
lines changed

‎drizzle-orm/src/relations.ts‎

Lines changed: 71 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,12 @@ export type RelationsRecord = Record<string, AnyRelation>;
273273
export type EmptyRelations = {};
274274
export type AnyRelations = TablesRelationalConfig;
275275

276+
/** Function that builds column SQL given the aliased table */
277+
export type RelationColumnBuilder = (table: SQL) => SQL;
278+
279+
/** A relation column can be either a Column or a builder function */
280+
export type RelationColumnValue = Column<any> | RelationColumnBuilder;
281+
276282
export abstract class Relation<
277283
TTargetTableName extends string = string,
278284
> {
@@ -281,8 +287,8 @@ export abstract class Relation<
281287
declare public readonly relationType: 'many' | 'one';
282288

283289
fieldName!: string;
284-
sourceColumns!: Column<any>[];
285-
targetColumns!: Column<any>[];
290+
sourceColumns!: RelationColumnValue[];
291+
targetColumns!: RelationColumnValue[];
286292
alias: string | undefined;
287293
where: AnyTableFilter | undefined;
288294
sourceTable!: SchemaEntry;
@@ -330,30 +336,38 @@ export class One<
330336
this.alias = config?.alias;
331337
this.where = config?.where;
332338
if (config?.from) {
333-
this.sourceColumns = ((Array.isArray(config.from)
339+
this.sourceColumns = (Array.isArray(config.from)
334340
? config.from
335-
: [config.from]) as RelationsBuilderColumnBase[]).map((it: RelationsBuilderColumnBase) => {
336-
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
337-
this.sourceColumnTableNames.push(it._.tableName);
338-
return it._.column as Column;
341+
: [config.from]).map((it: RelationConfigValue) => {
342+
if (isRelationsBuilderColumnBase(it)) {
343+
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
344+
this.sourceColumnTableNames.push(it._.tableName);
345+
return it._.column as Column;
346+
}
347+
return it; // It's a function, pass through
339348
});
340349
}
341350
if (config?.to) {
342351
this.targetColumns = (Array.isArray(config.to)
343352
? config.to
344-
: [config.to]).map((it: RelationsBuilderColumnBase) => {
345-
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
346-
this.targetColumnTableNames.push(it._.tableName);
347-
return it._.column as Column;
353+
: [config.to]).map((it: RelationConfigValue) => {
354+
if (isRelationsBuilderColumnBase(it)) {
355+
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
356+
this.targetColumnTableNames.push(it._.tableName);
357+
return it._.column as Column;
358+
}
359+
return it; // It's a function, pass through
348360
});
349361
}
350362

351363
if (this.throughTable) {
352364
this.through = {
353-
source: (Array.isArray(config?.from) ? config.from : config?.from ? [config.from] : []).map((
354-
c,
355-
) => c._.through!),
356-
target: (Array.isArray(config?.to) ? config.to : config?.to ? [config.to] : []).map((c) => c._.through!),
365+
source: (Array.isArray(config?.from) ? config.from : config?.from ? [config.from] : [])
366+
.filter(isRelationsBuilderColumnBase)
367+
.map((c) => c._.through!),
368+
target: (Array.isArray(config?.to) ? config.to : config?.to ? [config.to] : [])
369+
.filter(isRelationsBuilderColumnBase)
370+
.map((c) => c._.through!),
357371
};
358372
}
359373
this.optional = (config?.optional ?? true) as TOptional;
@@ -378,29 +392,37 @@ export class Many<TTargetTableName extends string> extends Relation<TTargetTable
378392
this.alias = config?.alias;
379393
this.where = config?.where;
380394
if (config?.from) {
381-
this.sourceColumns = ((Array.isArray(config.from)
395+
this.sourceColumns = (Array.isArray(config.from)
382396
? config.from
383-
: [config.from]) as RelationsBuilderColumnBase[]).map((it: RelationsBuilderColumnBase) => {
384-
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
385-
this.sourceColumnTableNames.push(it._.tableName);
386-
return it._.column as Column;
397+
: [config.from]).map((it: RelationConfigValue) => {
398+
if (isRelationsBuilderColumnBase(it)) {
399+
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
400+
this.sourceColumnTableNames.push(it._.tableName);
401+
return it._.column as Column;
402+
}
403+
return it; // It's a function, pass through
387404
});
388405
}
389406
if (config?.to) {
390407
this.targetColumns = (Array.isArray(config.to)
391408
? config.to
392-
: [config.to]).map((it: RelationsBuilderColumnBase) => {
393-
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
394-
this.targetColumnTableNames.push(it._.tableName);
395-
return it._.column as Column;
409+
: [config.to]).map((it: RelationConfigValue) => {
410+
if (isRelationsBuilderColumnBase(it)) {
411+
this.throughTable ??= it._.through ? tables[it._.through._.tableName]! as SchemaEntry : undefined;
412+
this.targetColumnTableNames.push(it._.tableName);
413+
return it._.column as Column;
414+
}
415+
return it; // It's a function, pass through
396416
});
397417
}
398418
if (this.throughTable) {
399419
this.through = {
400-
source: (Array.isArray(config?.from) ? config.from : config?.from ? [config.from] : []).map((
401-
c,
402-
) => c._.through!),
403-
target: (Array.isArray(config?.to) ? config.to : config?.to ? [config.to] : []).map((c) => c._.through!),
420+
source: (Array.isArray(config?.from) ? config.from : config?.from ? [config.from] : [])
421+
.filter(isRelationsBuilderColumnBase)
422+
.map((c) => c._.through!),
423+
target: (Array.isArray(config?.to) ? config.to : config?.to ? [config.to] : [])
424+
.filter(isRelationsBuilderColumnBase)
425+
.map((c) => c._.through!),
404426
};
405427
}
406428
}
@@ -1017,9 +1039,15 @@ export type AnyTableFilter = TableFilter<
10171039
FieldSelection
10181040
>;
10191041

1042+
export type RelationConfigValue = RelationsBuilderColumnBase | RelationColumnBuilder;
1043+
1044+
function isRelationsBuilderColumnBase(value: RelationConfigValue): value is RelationsBuilderColumnBase {
1045+
return typeof value !== 'function' && '_' in value;
1046+
}
1047+
10201048
export interface OneConfig<TTargetTable extends SchemaEntry, TOptional extends boolean> {
1021-
from?: RelationsBuilderColumnBase | [RelationsBuilderColumnBase, ...RelationsBuilderColumnBase[]];
1022-
to?: RelationsBuilderColumnBase | [RelationsBuilderColumnBase, ...RelationsBuilderColumnBase[]];
1049+
from?: RelationConfigValue | [RelationConfigValue, ...RelationConfigValue[]];
1050+
to?: RelationConfigValue | [RelationConfigValue, ...RelationConfigValue[]];
10231051
where?: TableFilter<TTargetTable>;
10241052
optional?: TOptional;
10251053
alias?: string;
@@ -1031,8 +1059,8 @@ export type AnyOneConfig = OneConfig<
10311059
>;
10321060

10331061
export interface ManyConfig<TTargetTable extends SchemaEntry> {
1034-
from?: RelationsBuilderColumnBase | [RelationsBuilderColumnBase, ...RelationsBuilderColumnBase[]];
1035-
to?: RelationsBuilderColumnBase | [RelationsBuilderColumnBase, ...RelationsBuilderColumnBase[]];
1062+
from?: RelationConfigValue | [RelationConfigValue, ...RelationConfigValue[]];
1063+
to?: RelationConfigValue | [RelationConfigValue, ...RelationConfigValue[]];
10361064
where?: TableFilter<TTargetTable>;
10371065
alias?: string;
10381066
}
@@ -1537,6 +1565,13 @@ export interface BuiltRelationFilters {
15371565
joinCondition?: SQL;
15381566
}
15391567

1568+
function buildColumnSQL(col: RelationColumnValue, table: SQL, casing: CasingCache): SQL {
1569+
if (typeof col === 'function') {
1570+
return col(table);
1571+
}
1572+
return sql`${table}.${sql.identifier(casing.getColumnCasing(col))}`;
1573+
}
1574+
15401575
export function relationToSQL(
15411576
casing: CasingCache,
15421577
relation: Relation,
@@ -1549,7 +1584,7 @@ export function relationToSQL(
15491584
const t = relation.through!.source[i]!;
15501585

15511586
return eq(
1552-
sql`${sourceTable}.${sql.identifier(casing.getColumnCasing(s))}`,
1587+
buildColumnSQL(s, sql`${sourceTable}`, casing),
15531588
sql`${throughTable!}.${sql.identifier(is(t._.column, Column) ? casing.getColumnCasing(t._.column) : t._.key)}`,
15541589
);
15551590
});
@@ -1559,7 +1594,7 @@ export function relationToSQL(
15591594

15601595
return eq(
15611596
sql`${throughTable!}.${sql.identifier(is(t._.column, Column) ? casing.getColumnCasing(t._.column) : t._.key)}`,
1562-
sql`${targetTable}.${sql.identifier(casing.getColumnCasing(s))}`,
1597+
buildColumnSQL(s, sql`${targetTable}`, casing),
15631598
);
15641599
});
15651600

@@ -1578,8 +1613,8 @@ export function relationToSQL(
15781613
const t = relation.targetColumns[i]!;
15791614

15801615
return eq(
1581-
sql`${sourceTable}.${sql.identifier(casing.getColumnCasing(s))}`,
1582-
sql`${targetTable}.${sql.identifier(casing.getColumnCasing(t))}`,
1616+
buildColumnSQL(s, sql`${sourceTable}`, casing),
1617+
buildColumnSQL(t, sql`${targetTable}`, casing),
15831618
);
15841619
});
15851620

‎integration-tests/tests/pg/common-rqb.ts‎

Lines changed: 128 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// eslint-disable-next-line @typescript-eslint/consistent-type-imports
22
import { and, eq, inArray, isNotNull, not, or, sql } from 'drizzle-orm';
33
import type { PgColumnBuilder } from 'drizzle-orm/pg-core';
4-
import { bigint, integer, pgEnum, pgTable, serial, text, timestamp } from 'drizzle-orm/pg-core';
4+
import { bigint, integer, jsonb, pgEnum, pgTable, serial, text, timestamp } from 'drizzle-orm/pg-core';
55
import { describe, expect, expectTypeOf } from 'vitest';
66
import type { Test } from './instrumentation';
77

@@ -1108,5 +1108,132 @@ export function tests(test: Test) {
11081108
expect(res).toStrictEqual([{ id: 1, status: 1 }, { id: 2, status: 0 }]);
11091109
},
11101110
);
1111+
1112+
test.concurrent(
1113+
'RQB v2 relation with custom SQL builder function (type casting)',
1114+
async ({ push, createDB }) => {
1115+
const users = pgTable('rqb_users_sqlbuilder_1', {
1116+
id: bigint('id', { mode: 'bigint' }).primaryKey(),
1117+
name: text('name').notNull(),
1118+
});
1119+
1120+
const posts = pgTable('rqb_posts_sqlbuilder_1', {
1121+
id: serial('id').primaryKey(),
1122+
authorId: text('author_id').notNull(),
1123+
content: text('content'),
1124+
});
1125+
1126+
await push({ users, posts });
1127+
const db = createDB({ users, posts }, (r) => ({
1128+
users: {
1129+
posts: r.many.posts({
1130+
from: (table) => sql`${table}.${sql.identifier('id')}::varchar`,
1131+
to: r.posts.authorId,
1132+
}),
1133+
},
1134+
posts: {
1135+
author: r.one.users({
1136+
from: r.posts.authorId,
1137+
to: (table) => sql`${table}.${sql.identifier('id')}::varchar`,
1138+
}),
1139+
},
1140+
}));
1141+
1142+
await db.insert(users).values([
1143+
{ id: 1n, name: 'Alice' },
1144+
{ id: 2n, name: 'Bob' },
1145+
]);
1146+
1147+
await db.insert(posts).values([
1148+
{ id: 1, authorId: '1', content: 'Post by Alice' },
1149+
{ id: 2, authorId: '1', content: 'Another post by Alice' },
1150+
{ id: 3, authorId: '2', content: 'Post by Bob' },
1151+
]);
1152+
1153+
const userWithPosts = await db.query.users.findFirst({
1154+
where: { id: 1n },
1155+
with: {
1156+
posts: {
1157+
orderBy: { id: 'asc' },
1158+
},
1159+
},
1160+
});
1161+
1162+
expect(userWithPosts).toStrictEqual({
1163+
id: 1n,
1164+
name: 'Alice',
1165+
posts: [
1166+
{ id: 1, authorId: '1', content: 'Post by Alice' },
1167+
{ id: 2, authorId: '1', content: 'Another post by Alice' },
1168+
],
1169+
});
1170+
1171+
const postWithAuthor = await db.query.posts.findFirst({
1172+
where: { id: 3 },
1173+
with: {
1174+
author: true,
1175+
},
1176+
});
1177+
1178+
expect(postWithAuthor).toStrictEqual({
1179+
id: 3,
1180+
authorId: '2',
1181+
content: 'Post by Bob',
1182+
author: { id: 2n, name: 'Bob' },
1183+
});
1184+
},
1185+
);
1186+
1187+
test.concurrent(
1188+
'RQB v2 relation with custom SQL builder for JSON extraction',
1189+
async ({ push, createDB }) => {
1190+
const products = pgTable('rqb_products_jsonb_1', {
1191+
id: serial('id').primaryKey(),
1192+
name: text('name').notNull(),
1193+
metadata: jsonb('metadata').$type<{ categoryId: string }>(),
1194+
});
1195+
1196+
const categories = pgTable('rqb_categories_jsonb_1', {
1197+
id: text('id').primaryKey(),
1198+
name: text('name').notNull(),
1199+
});
1200+
1201+
await push({ products, categories });
1202+
const db = createDB({ products, categories }, (r) => ({
1203+
products: {
1204+
category: r.one.categories({
1205+
from: (table) => sql`${table}.${sql.identifier('metadata')}->>'categoryId'`,
1206+
to: r.categories.id,
1207+
optional: true,
1208+
}),
1209+
},
1210+
}));
1211+
1212+
await db.insert(categories).values([
1213+
{ id: 'cat1', name: 'Electronics' },
1214+
{ id: 'cat2', name: 'Books' },
1215+
]);
1216+
1217+
await db.insert(products).values([
1218+
{ id: 1, name: 'Phone', metadata: { categoryId: 'cat1' } },
1219+
{ id: 2, name: 'Laptop', metadata: { categoryId: 'cat1' } },
1220+
{ id: 3, name: 'Novel', metadata: { categoryId: 'cat2' } },
1221+
]);
1222+
1223+
const productWithCategory = await db.query.products.findFirst({
1224+
where: { id: 1 },
1225+
with: {
1226+
category: true,
1227+
},
1228+
});
1229+
1230+
expect(productWithCategory).toStrictEqual({
1231+
id: 1,
1232+
name: 'Phone',
1233+
metadata: { categoryId: 'cat1' },
1234+
category: { id: 'cat1', name: 'Electronics' },
1235+
});
1236+
},
1237+
);
11111238
});
11121239
}

0 commit comments

Comments
 (0)