Feat: Set operations (union / intersect / except) support

drizzle-orm/src/mysql-core/dialect.ts

@@ -19,7 +19,7 @@ import { Subquery, SubqueryConfig } from '~/subquery.ts';
 import { getTableName, Table } from '~/table.ts';
 import { orderSelectedFields, type UpdateSet } from '~/utils.ts';
 import { View } from '~/view.ts';
-import { DrizzleError, ViewBaseConfig } from '../index.ts';
+import { DrizzleError, type Name, ViewBaseConfig } from '../index.ts';
 import { MySqlColumn } from './columns/common.ts';
 import type { MySqlDeleteConfig } from './query-builders/delete.ts';
 import type { MySqlInsertConfig } from './query-builders/insert.ts';
@@ -204,6 +204,7 @@ export class MySqlDialect {
 			offset,
 			lockingClause,
 			distinct,
+			setOperators,
 		}: MySqlSelectConfig,
 	): SQL {
 		const fieldsList = fieldsFlat ?? orderSelectedFields<MySqlColumn>(fields);
@@ -331,7 +332,75 @@ export class MySqlDialect {
 			}
 		}
 
-		return sql`${withSql}select${distinctSql} ${selection} from ${tableSql}${joinsSql}${whereSql}${groupBySql}${havingSql}${orderBySql}${limitSql}${offsetSql}${lockingClausesSql}`;
+		const finalQuery =
+			sql`${withSql}select${distinctSql} ${selection} from ${tableSql}${joinsSql}${whereSql}${groupBySql}${havingSql}${orderBySql}${limitSql}${offsetSql}${lockingClausesSql}`;
+
+		if (setOperators.length > 0) {
+			return this.buildSetOperations(finalQuery, setOperators);
+		}
+
+		return finalQuery;
+	}
+
+	buildSetOperations(leftSelect: SQL, setOperators: MySqlSelectConfig['setOperators']): SQL {
+		const [setOperator, ...rest] = setOperators;
+
+		if (!setOperator) {
+			throw new Error('Cannot pass undefined values to any set operator');
+		}
+
+		if (rest.length === 0) {
+			return this.buildSetOperationQuery({ leftSelect, setOperator });
+		}
+
+		// Some recursive magic here
+		return this.buildSetOperations(
+			this.buildSetOperationQuery({ leftSelect, setOperator }),
+			rest,
+		);
+	}
+
+	buildSetOperationQuery({
+		leftSelect,
+		setOperator: { type, isAll, rightSelect, limit, orderBy, offset },
+	}: { leftSelect: SQL; setOperator: MySqlSelectConfig['setOperators'][number] }): SQL {
+		const leftChunk = sql`(${leftSelect.getSQL()}) `;
+		const rightChunk = sql`(${rightSelect.getSQL()})`;
+
+		let orderBySql;
+		if (orderBy && orderBy.length > 0) {
+			const orderByValues: (SQL<unknown> | Name)[] = [];
+
+			// The next bit is necessary because the sql operator replaces ${table.column} with `table`.`column`
+			// which is invalid MySql syntax, Table from one of the SELECTs cannot be used in global ORDER clause
+			for (const orderByUnit of orderBy) {
+				if (is(orderByUnit, MySqlColumn)) {
+					orderByValues.push(sql.identifier(orderByUnit.name));
+				} else if (is(orderByUnit, SQL)) {
+					for (let i = 0; i < orderByUnit.queryChunks.length; i++) {
+						const chunk = orderByUnit.queryChunks[i];
+
+						if (is(chunk, MySqlColumn)) {
+							orderByUnit.queryChunks[i] = sql.identifier(chunk.name);
+						}
+					}
+
+					orderByValues.push(sql`${orderByUnit}`);
+				} else {
+					orderByValues.push(sql`${orderByUnit}`);
+				}
+			}
+
+			orderBySql = sql` order by ${sql.join(orderByValues, sql`, `)} `;
+		}
+
+		const limitSql = limit ? sql` limit ${limit}` : undefined;
+
+		const operatorChunk = sql.raw(`${type} ${isAll ? 'all ' : ''}`);
+
+		const offsetSql = offset ? sql` offset ${offset}` : undefined;
+
+		return sql`${leftChunk}${operatorChunk}${rightChunk}${orderBySql}${limitSql}${offsetSql}`;
 	}
 
 	buildInsertQuery({ table, values, ignore, onConflict }: MySqlInsertConfig): SQL {
@@ -631,6 +700,7 @@ export class MySqlDialect {
 					where,
 					limit,
 					offset,
+					setOperators: [],
 				});
 
 				where = undefined;
@@ -653,6 +723,7 @@ export class MySqlDialect {
 				limit,
 				offset,
 				orderBy,
+				setOperators: [],
 			});
 		} else {
 			result = this.buildSelectQuery({
@@ -667,6 +738,7 @@ export class MySqlDialect {
 				limit,
 				offset,
 				orderBy,
+				setOperators: [],
 			});
 		}
 
@@ -921,6 +993,7 @@ export class MySqlDialect {
 					where,
 					limit,
 					offset,
+					setOperators: [],
 				});
 
 				where = undefined;
@@ -942,6 +1015,7 @@ export class MySqlDialect {
 				limit,
 				offset,
 				orderBy,
+				setOperators: [],
 			});
 		} else {
 			result = this.buildSelectQuery({
@@ -955,6 +1029,7 @@ export class MySqlDialect {
 				limit,
 				offset,
 				orderBy,
+				setOperators: [],
 			});
 		}
 

drizzle-orm/src/mysql-core/query-builders/select.ts

@@ -14,27 +14,34 @@ import type {
 	JoinType,
 	SelectMode,
 	SelectResult,
+	SetOperator,
 } from '~/query-builders/select.types.ts';
 import { QueryPromise } from '~/query-promise.ts';
 import { type Query, SQL } from '~/sql/index.ts';
 import { SelectionProxyHandler, Subquery, SubqueryConfig } from '~/subquery.ts';
 import { Table } from '~/table.ts';
-import { applyMixins, getTableColumns, getTableLikeName, type ValueOrArray } from '~/utils.ts';
+import { applyMixins, getTableColumns, getTableLikeName, haveSameKeys, type ValueOrArray } from '~/utils.ts';
 import { orderSelectedFields } from '~/utils.ts';
 import { ViewBaseConfig } from '~/view-common.ts';
 import { type ColumnsSelection, View } from '~/view.ts';
 import type {
+	AnyMySqlSelect,
 	CreateMySqlSelectFromBuilderMode,
+	GetMySqlSetOperators,
 	LockConfig,
 	LockStrength,
+	MySqlCreateSetOperatorFn,
 	MySqlJoinFn,
 	MySqlSelectConfig,
 	MySqlSelectDynamic,
 	MySqlSelectHKT,
 	MySqlSelectHKTBase,
 	MySqlSelectPrepare,
 	MySqlSelectWithout,
+	MySqlSetOperatorExcludedMethods,
+	MySqlSetOperatorWithResult,
 	SelectedFields,
+	SetOperatorRightSelect,
 } from './select.types.ts';
 
 export class MySqlSelectBuilder<
@@ -121,7 +128,7 @@ export abstract class MySqlSelectQueryBuilderBase<
 		: {},
 	TDynamic extends boolean = false,
 	TExcludedMethods extends string = never,
-	TResult = SelectResult<TSelection, TSelectMode, TNullabilityMap>[],
+	TResult extends any[] = SelectResult<TSelection, TSelectMode, TNullabilityMap>[],
 	TSelectedFields extends ColumnsSelection = BuildSubquerySelection<TSelection, TNullabilityMap>,
 > extends TypedQueryBuilder<TSelectedFields, TResult> {
 	static readonly [entityKind]: string = 'MySqlSelectQueryBuilder';
@@ -164,6 +171,7 @@ export abstract class MySqlSelectQueryBuilderBase<
 			table,
 			fields: { ...fields },
 			distinct,
+			setOperators: [],
 		};
 		this.isPartialSelect = isPartialSelect;
 		this.session = session;
@@ -260,6 +268,61 @@ export abstract class MySqlSelectQueryBuilderBase<
 
 	fullJoin = this.createJoin('full');
 
+	private createSetOperator(
+		type: SetOperator,
+		isAll: boolean,
+	): <TValue extends MySqlSetOperatorWithResult<TResult>>(
+		rightSelection:
+			| ((setOperators: GetMySqlSetOperators) => SetOperatorRightSelect<TValue, TResult>)
+			| SetOperatorRightSelect<TValue, TResult>,
+	) => MySqlSelectWithout<
+		this,
+		TDynamic,
+		MySqlSetOperatorExcludedMethods,
+		true
+	> {
+		return (rightSelection) => {
+			const rightSelect = (typeof rightSelection === 'function'
+				? rightSelection(getMySqlSetOperators())
+				: rightSelection) as TypedQueryBuilder<
+					any,
+					TResult
+				>;
+
+			if (!haveSameKeys(this.getSelectedFields(), rightSelect.getSelectedFields())) {
+				throw new Error(
+					'Set operator error (union / intersect / except): selected fields are not the same or are in a different order',
+				);
+			}
+
+			this.config.setOperators.push({ type, isAll, rightSelect });
+			return this as any;
+		};
+	}
+
+	union = this.createSetOperator('union', false);
+
+	unionAll = this.createSetOperator('union', true);
+
+	intersect = this.createSetOperator('intersect', false);
+
+	intersectAll = this.createSetOperator('intersect', true);
+
+	except = this.createSetOperator('except', false);
+
+	exceptAll = this.createSetOperator('except', true);
+
+	/** @internal */
+	addSetOperators(setOperators: MySqlSelectConfig['setOperators']): MySqlSelectWithout<
+		this,
+		TDynamic,
+		MySqlSetOperatorExcludedMethods,
+		true
+	> {
+		this.config.setOperators.push(...setOperators);
+		return this as any;
+	}
+
 	where(
 		where: ((aliases: this['_']['selection']) => SQL | undefined) | SQL | undefined,
 	): MySqlSelectWithout<this, TDynamic, 'where'> {
@@ -329,20 +392,41 @@ export abstract class MySqlSelectQueryBuilderBase<
 					new SelectionProxyHandler({ sqlAliasedBehavior: 'alias', sqlBehavior: 'sql' }),
 				) as TSelection,
 			);
-			this.config.orderBy = Array.isArray(orderBy) ? orderBy : [orderBy];
+
+			const orderByArray = Array.isArray(orderBy) ? orderBy : [orderBy];
+
+			if (this.config.setOperators.length > 0) {
+				this.config.setOperators.at(-1)!.orderBy = orderByArray;
+			} else {
+				this.config.orderBy = orderByArray;
+			}
 		} else {
-			this.config.orderBy = columns as (MySqlColumn | SQL | SQL.Aliased)[];
+			const orderByArray = columns as (MySqlColumn | SQL | SQL.Aliased)[];
+
+			if (this.config.setOperators.length > 0) {
+				this.config.setOperators.at(-1)!.orderBy = orderByArray;
+			} else {
+				this.config.orderBy = orderByArray;
+			}
 		}
 		return this as any;
 	}
 
 	limit(limit: number): MySqlSelectWithout<this, TDynamic, 'limit'> {
-		this.config.limit = limit;
+		if (this.config.setOperators.length > 0) {
+			this.config.setOperators.at(-1)!.limit = limit;
+		} else {
+			this.config.limit = limit;
+		}
 		return this as any;
 	}
 
 	offset(offset: number): MySqlSelectWithout<this, TDynamic, 'offset'> {
-		this.config.offset = offset;
+		if (this.config.setOperators.length > 0) {
+			this.config.setOperators.at(-1)!.offset = offset;
+		} else {
+			this.config.offset = offset;
+		}
 		return this as any;
 	}
 
@@ -392,7 +476,7 @@ export interface MySqlSelectBase<
 		: {},
 	TDynamic extends boolean = false,
 	TExcludedMethods extends string = never,
-	TResult = SelectResult<TSelection, TSelectMode, TNullabilityMap>[],
+	TResult extends any[] = SelectResult<TSelection, TSelectMode, TNullabilityMap>[],
 	TSelectedFields extends ColumnsSelection = BuildSubquerySelection<TSelection, TNullabilityMap>,
 > extends
 	MySqlSelectQueryBuilderBase<
@@ -463,3 +547,44 @@ export class MySqlSelectBase<
 }
 
 applyMixins(MySqlSelectBase, [QueryPromise]);
+
+function createSetOperator(type: SetOperator, isAll: boolean): MySqlCreateSetOperatorFn {
+	return (leftSelect, rightSelect, ...restSelects) => {
+		const setOperators = [rightSelect, ...restSelects].map((select) => ({
+			type,
+			isAll,
+			rightSelect: select as AnyMySqlSelect,
+		}));
+
+		for (const setOperator of setOperators) {
+			if (!haveSameKeys((leftSelect as any).getSelectedFields(), setOperator.rightSelect.getSelectedFields())) {
+				throw new Error(
+					'Set operator error (union / intersect / except): selected fields are not the same or are in a different order',
+				);
+			}
+		}
+
+		return (leftSelect as AnyMySqlSelect).addSetOperators(setOperators) as any;
+	};
+}
+
+const getMySqlSetOperators = () => ({
+	union,
+	unionAll,
+	intersect,
+	intersectAll,
+	except,
+	exceptAll,
+});
+
+export const union = createSetOperator('union', false);
+
+export const unionAll = createSetOperator('union', true);
+
+export const intersect = createSetOperator('intersect', false);
+
+export const intersectAll = createSetOperator('intersect', true);
+
+export const except = createSetOperator('except', false);
+
+export const exceptAll = createSetOperator('except', true);