Skip to content

Commit

Permalink
Merge pull request #1218 from Angelelz/feat-union
Browse files Browse the repository at this point in the history
  • Loading branch information
@dankochetov
dankochetov authored Nov 2, 2023
2 parents 08e3043 + a83ea7d commit 1a482ce
Show file tree
Hide file tree
Showing 25 changed files with 3,587 additions and 65 deletions.
79 changes: 77 additions & 2 deletions drizzle-orm/src/mysql-core/dialect.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -204,6 +204,7 @@ export class MySqlDialect {
offset,
lockingClause,
distinct,
setOperators,
}: MySqlSelectConfig,
): SQL {
const fieldsList = fieldsFlat ?? orderSelectedFields<MySqlColumn>(fields);
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -631,6 +700,7 @@ export class MySqlDialect {
where,
limit,
offset,
setOperators: [],
});

where = undefined;
Expand All @@ -653,6 +723,7 @@ export class MySqlDialect {
limit,
offset,
orderBy,
setOperators: [],
});
} else {
result = this.buildSelectQuery({
Expand All @@ -667,6 +738,7 @@ export class MySqlDialect {
limit,
offset,
orderBy,
setOperators: [],
});
}

Expand Down Expand Up @@ -921,6 +993,7 @@ export class MySqlDialect {
where,
limit,
offset,
setOperators: [],
});

where = undefined;
Expand All @@ -942,6 +1015,7 @@ export class MySqlDialect {
limit,
offset,
orderBy,
setOperators: [],
});
} else {
result = this.buildSelectQuery({
Expand All @@ -955,6 +1029,7 @@ export class MySqlDialect {
limit,
offset,
orderBy,
setOperators: [],
});
}

Expand Down
139 changes: 132 additions & 7 deletions drizzle-orm/src/mysql-core/query-builders/select.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<
Expand Down Expand Up @@ -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';
Expand Down Expand Up @@ -164,6 +171,7 @@ export abstract class MySqlSelectQueryBuilderBase<
table,
fields: { ...fields },
distinct,
setOperators: [],
};
this.isPartialSelect = isPartialSelect;
this.session = session;
Expand Down Expand Up @@ -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'> {
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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<
Expand Down Expand Up @@ -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);
Loading

0 comments on commit 1a482ce

Please sign in to comment.