Fixed extras inferrence with SQL.Aliased, returned columns argume…

…nt relations' `where.$raw` callback
drizzle-team · Dec 20, 2024 · 661b6f2 · 661b6f2
1 parent 0ab568f
commit 661b6f2
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/drizzle-orm/src/relations.ts b/drizzle-orm/src/relations.ts
@@ -550,9 +550,11 @@ export type BuildQueryResult<
 						K in NonUndefinedKeysOnly<
 							ReturnTypeOrValue<TFullSelection['extras']>
 						>
-					]: Assume<
-						ReturnTypeOrValue<TFullSelection['extras']>[K],
-						SQL
+					]: ReturnType<
+						Assume<
+							ReturnTypeOrValue<TFullSelection['extras']>[K],
+							SQLWrapper
+						>['getSQL']
 					>['_']['type'];
 				}
 				: {})
@@ -722,7 +724,12 @@ export type RelationsFilter<TColumns extends Record<string, Column>> =
 	& {
 		$or?: RelationsFilter<TColumns>[];
 		$not?: RelationsFilter<TColumns>[];
-		$raw?: (operators: Operators) => SQL;
+		$raw?: (
+			columns: {
+				[K in keyof TColumns]: TColumns[K];
+			},
+			operators: Operators,
+		) => SQL;
 	};
 
 export interface OneConfig<
@@ -832,11 +839,11 @@ export class RelationsHelperStatic<TTables extends Record<string, Table>> {
 
 		for (const [tableName, table] of Object.entries(tables)) {
 			one[tableName] = (config) => {
-				return new One(table, config);
+				return new One(table, config as DBQueryConfig<'one'>);
 			};
 
 			many[tableName] = (config) => {
-				return new Many(table, config);
+				return new Many(table, config as DBQueryConfig<'many'>);
 			};
 		}
 
@@ -1020,7 +1027,11 @@ export function relationFilterToSQL(
 
 		switch (target) {
 			case '$raw': {
-				if (value) parts.push((value as (operators: Operators) => SQL)(operators));
+				if (value) {
+					parts.push(
+						(value as (table: Record<string, Column>, operators: Operators) => SQL)(table[Columns], operators),
+					);
+				}
 
 				continue;
 			}

diff --git a/drizzle-orm/src/sql/sql.ts b/drizzle-orm/src/sql/sql.ts
@@ -549,7 +549,7 @@ export namespace sql {
 }
 
 export namespace SQL {
-	export class Aliased<T = unknown> implements SQLWrapper {
+	export class Aliased<T = unknown> implements SQLWrapper<T> {
 		static readonly [entityKind]: string = 'SQL.Aliased';
 
 		declare _: {
@@ -561,17 +561,17 @@ export namespace SQL {
 		isSelectionField = false;
 
 		constructor(
-			readonly sql: SQL,
+			readonly sql: SQL<T>,
 			readonly fieldAlias: string,
 		) {}
 
-		getSQL(): SQL {
-			return this.sql;
+		getSQL(): SQL<T> {
+			return this.sql as SQL<T>;
 		}
 
 		/** @internal */
 		clone() {
-			return new Aliased(this.sql, this.fieldAlias);
+			return new Aliased<T>(this.sql, this.fieldAlias);
 		}
 	}
 }