Skip to content

Commit

Permalink
Merge pull request #3693 from drizzle-team/feature/mysql-force-index
Browse files Browse the repository at this point in the history 
feature/mysql-force-index
  • Loading branch information
@AndriiSherman
AndriiSherman authored Dec 12, 2024
2 parents 21dab20 + d9d234e commit 7cd9d79
Show file tree
Hide file tree
Showing 7 changed files with 964 additions and 12 deletions.
30 changes: 28 additions & 2 deletions drizzle-orm/src/mysql-core/dialect.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,18 @@ export class MySqlDialect {
return orderBy && orderBy.length > 0 ? sql` order by ${sql.join(orderBy, sql`, `)}` : undefined;
}

private buildIndex({
indexes,
indexFor,
}: {
indexes: string[] | undefined;
indexFor: 'USE' | 'FORCE' | 'IGNORE';
}): SQL | undefined {
return indexes && indexes.length > 0
? sql` ${sql.raw(indexFor)} INDEX (${sql.raw(indexes.join(`, `))})`
: undefined;
}

buildSelectQuery(
{
withList,
Expand All @@ -260,6 +272,9 @@ export class MySqlDialect {
lockingClause,
distinct,
setOperators,
useIndex,
forceIndex,
ignoreIndex,
}: MySqlSelectConfig,
): SQL {
const fieldsList = fieldsFlat ?? orderSelectedFields<MySqlColumn>(fields);
Expand Down Expand Up @@ -319,10 +334,15 @@ export class MySqlDialect {
const tableSchema = table[MySqlTable.Symbol.Schema];
const origTableName = table[MySqlTable.Symbol.OriginalName];
const alias = tableName === origTableName ? undefined : joinMeta.alias;
const useIndexSql = this.buildIndex({ indexes: joinMeta.useIndex, indexFor: 'USE' });
const forceIndexSql = this.buildIndex({ indexes: joinMeta.forceIndex, indexFor: 'FORCE' });
const ignoreIndexSql = this.buildIndex({ indexes: joinMeta.ignoreIndex, indexFor: 'IGNORE' });
joinsArray.push(
sql`${sql.raw(joinMeta.joinType)} join${lateralSql} ${
tableSchema ? sql`${sql.identifier(tableSchema)}.` : undefined
}${sql.identifier(origTableName)}${alias && sql` ${sql.identifier(alias)}`} on ${joinMeta.on}`,
}${sql.identifier(origTableName)}${useIndexSql}${forceIndexSql}${ignoreIndexSql}${
alias && sql` ${sql.identifier(alias)}`
} on ${joinMeta.on}`,
);
} else if (is(table, View)) {
const viewName = table[ViewBaseConfig].name;
Expand Down Expand Up @@ -359,6 +379,12 @@ export class MySqlDialect {

const offsetSql = offset ? sql` offset ${offset}` : undefined;

const useIndexSql = this.buildIndex({ indexes: useIndex, indexFor: 'USE' });

const forceIndexSql = this.buildIndex({ indexes: forceIndex, indexFor: 'FORCE' });

const ignoreIndexSql = this.buildIndex({ indexes: ignoreIndex, indexFor: 'IGNORE' });

let lockingClausesSql;
if (lockingClause) {
const { config, strength } = lockingClause;
Expand All @@ -371,7 +397,7 @@ export class MySqlDialect {
}

const finalQuery =
sql`${withSql}select${distinctSql} ${selection} from ${tableSql}${joinsSql}${whereSql}${groupBySql}${havingSql}${orderBySql}${limitSql}${offsetSql}${lockingClausesSql}`;
sql`${withSql}select${distinctSql} ${selection} from ${tableSql}${useIndexSql}${forceIndexSql}${ignoreIndexSql}${joinsSql}${whereSql}${groupBySql}${havingSql}${orderBySql}${limitSql}${offsetSql}${lockingClausesSql}`;

if (setOperators.length > 0) {
return this.buildSetOperations(finalQuery, setOperators);
Expand Down
107 changes: 101 additions & 6 deletions drizzle-orm/src/mysql-core/query-builders/select.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import type { MySqlColumn } from '~/mysql-core/columns/index.ts';
import type { MySqlDialect } from '~/mysql-core/dialect.ts';
import type { MySqlPreparedQueryConfig, MySqlSession, PreparedQueryHKTBase } from '~/mysql-core/session.ts';
import type { SubqueryWithSelection } from '~/mysql-core/subquery.ts';
import type { MySqlTable } from '~/mysql-core/table.ts';
import { MySqlTable } from '~/mysql-core/table.ts';
import { TypedQueryBuilder } from '~/query-builders/query-builder.ts';
import type {
BuildSubquerySelection,
Expand All @@ -21,9 +21,11 @@ import type { ColumnsSelection, Placeholder, Query } from '~/sql/sql.ts';
import { SQL, View } from '~/sql/sql.ts';
import { Subquery } from '~/subquery.ts';
import { Table } from '~/table.ts';
import { applyMixins, getTableColumns, getTableLikeName, haveSameKeys, type ValueOrArray } from '~/utils.ts';
import { orderSelectedFields } from '~/utils.ts';
import type { ValueOrArray } from '~/utils.ts';
import { applyMixins, getTableColumns, getTableLikeName, haveSameKeys, orderSelectedFields } from '~/utils.ts';
import { ViewBaseConfig } from '~/view-common.ts';
import type { IndexBuilder } from '../indexes.ts';
import { convertIndexToString, toArray } from '../utils.ts';
import { MySqlViewBase } from '../view-base.ts';
import type {
AnyMySqlSelect,
Expand All @@ -45,6 +47,14 @@ import type {
SetOperatorRightSelect,
} from './select.types.ts';

export type IndexForHint = IndexBuilder | string;

export type IndexConfig = {
useIndex?: IndexForHint | IndexForHint[];
forceIndex?: IndexForHint | IndexForHint[];
ignoreIndex?: IndexForHint | IndexForHint[];
};

export class MySqlSelectBuilder<
TSelection extends SelectedFields | undefined,
TPreparedQueryHKT extends PreparedQueryHKTBase,
Expand Down Expand Up @@ -78,6 +88,8 @@ export class MySqlSelectBuilder<

from<TFrom extends MySqlTable | Subquery | MySqlViewBase | SQL>(
source: TFrom,
onIndex?: TFrom extends MySqlTable ? IndexConfig
: 'Index hint configuration is allowed only for MySqlTable and not for subqueries or views',
): CreateMySqlSelectFromBuilderMode<
TBuilderMode,
GetSelectTableName<TFrom>,
Expand Down Expand Up @@ -105,6 +117,21 @@ export class MySqlSelectBuilder<
fields = getTableColumns<MySqlTable>(source);
}

let useIndex: string[] = [];
let forceIndex: string[] = [];
let ignoreIndex: string[] = [];
if (is(source, MySqlTable) && onIndex && typeof onIndex !== 'string') {
if (onIndex.useIndex) {
useIndex = convertIndexToString(toArray(onIndex.useIndex));
}
if (onIndex.forceIndex) {
forceIndex = convertIndexToString(toArray(onIndex.forceIndex));
}
if (onIndex.ignoreIndex) {
ignoreIndex = convertIndexToString(toArray(onIndex.ignoreIndex));
}
}

return new MySqlSelectBase(
{
table: source,
Expand All @@ -114,6 +141,9 @@ export class MySqlSelectBuilder<
dialect: this.dialect,
withList: this.withList,
distinct: this.distinct,
useIndex,
forceIndex,
ignoreIndex,
},
) as any;
}
Expand Down Expand Up @@ -156,14 +186,17 @@ export abstract class MySqlSelectQueryBuilderBase<
protected dialect: MySqlDialect;

constructor(
{ table, fields, isPartialSelect, session, dialect, withList, distinct }: {
{ table, fields, isPartialSelect, session, dialect, withList, distinct, useIndex, forceIndex, ignoreIndex }: {
table: MySqlSelectConfig['table'];
fields: MySqlSelectConfig['fields'];
isPartialSelect: boolean;
session: MySqlSession | undefined;
dialect: MySqlDialect;
withList: Subquery[];
distinct: boolean | undefined;
useIndex?: string[];
forceIndex?: string[];
ignoreIndex?: string[];
},
) {
super();
Expand All @@ -173,6 +206,9 @@ export abstract class MySqlSelectQueryBuilderBase<
fields: { ...fields },
distinct,
setOperators: [],
useIndex,
forceIndex,
ignoreIndex,
};
this.isPartialSelect = isPartialSelect;
this.session = session;
Expand All @@ -187,9 +223,13 @@ export abstract class MySqlSelectQueryBuilderBase<
private createJoin<TJoinType extends JoinType>(
joinType: TJoinType,
): MySqlJoinFn<this, TDynamic, TJoinType> {
return (
return <
TJoinedTable extends MySqlTable | Subquery | MySqlViewBase | SQL,
>(
table: MySqlTable | Subquery | MySqlViewBase | SQL,
on: ((aliases: TSelection) => SQL | undefined) | SQL | undefined,
onIndex?: TJoinedTable extends MySqlTable ? IndexConfig
: 'Index hint configuration is allowed only for MySqlTable and not for subqueries or views',
) => {
const baseTableName = this.tableName;
const tableName = getTableLikeName(table);
Expand Down Expand Up @@ -228,7 +268,22 @@ export abstract class MySqlSelectQueryBuilderBase<
this.config.joins = [];
}

this.config.joins.push({ on, table, joinType, alias: tableName });
let useIndex: string[] = [];
let forceIndex: string[] = [];
let ignoreIndex: string[] = [];
if (is(table, MySqlTable) && onIndex && typeof onIndex !== 'string') {
if (onIndex.useIndex) {
useIndex = convertIndexToString(toArray(onIndex.useIndex));
}
if (onIndex.forceIndex) {
forceIndex = convertIndexToString(toArray(onIndex.forceIndex));
}
if (onIndex.ignoreIndex) {
ignoreIndex = convertIndexToString(toArray(onIndex.ignoreIndex));
}
}

this.config.joins.push({ on, table, joinType, alias: tableName, useIndex, forceIndex, ignoreIndex });

if (typeof tableName === 'string') {
switch (joinType) {
Expand Down Expand Up @@ -286,6 +341,16 @@ export abstract class MySqlSelectQueryBuilderBase<
* })
* .from(users)
* .leftJoin(pets, eq(users.id, pets.ownerId))
*
* // Select userId and petId with use index hint
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
* userId: users.id,
* petId: pets.id,
* })
* .from(users)
* .leftJoin(pets, eq(users.id, pets.ownerId), {
* useIndex: ['pets_owner_id_index']
* })
* ```
*/
leftJoin = this.createJoin('left');
Expand Down Expand Up @@ -315,6 +380,16 @@ export abstract class MySqlSelectQueryBuilderBase<
* })
* .from(users)
* .rightJoin(pets, eq(users.id, pets.ownerId))
*
* // Select userId and petId with use index hint
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
* userId: users.id,
* petId: pets.id,
* })
* .from(users)
* .leftJoin(pets, eq(users.id, pets.ownerId), {
* useIndex: ['pets_owner_id_index']
* })
* ```
*/
rightJoin = this.createJoin('right');
Expand Down Expand Up @@ -344,6 +419,16 @@ export abstract class MySqlSelectQueryBuilderBase<
* })
* .from(users)
* .innerJoin(pets, eq(users.id, pets.ownerId))
*
* // Select userId and petId with use index hint
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
* userId: users.id,
* petId: pets.id,
* })
* .from(users)
* .leftJoin(pets, eq(users.id, pets.ownerId), {
* useIndex: ['pets_owner_id_index']
* })
* ```
*/
innerJoin = this.createJoin('inner');
Expand Down Expand Up @@ -373,6 +458,16 @@ export abstract class MySqlSelectQueryBuilderBase<
* })
* .from(users)
* .fullJoin(pets, eq(users.id, pets.ownerId))
*
* // Select userId and petId with use index hint
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
* userId: users.id,
* petId: pets.id,
* })
* .from(users)
* .leftJoin(pets, eq(users.id, pets.ownerId), {
* useIndex: ['pets_owner_id_index']
* })
* ```
*/
fullJoin = this.createJoin('full');
Expand Down
10 changes: 9 additions & 1 deletion drizzle-orm/src/mysql-core/query-builders/select.types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,17 @@ import type { Assume, ValidateShape } from '~/utils.ts';
import type { MySqlPreparedQueryConfig, PreparedQueryHKTBase, PreparedQueryKind } from '../session.ts';
import type { MySqlViewBase } from '../view-base.ts';
import type { MySqlViewWithSelection } from '../view.ts';
import type { MySqlSelectBase, MySqlSelectQueryBuilderBase } from './select.ts';
import type { IndexConfig, MySqlSelectBase, MySqlSelectQueryBuilderBase } from './select.ts';

export interface MySqlSelectJoinConfig {
on: SQL | undefined;
table: MySqlTable | Subquery | MySqlViewBase | SQL;
alias: string | undefined;
joinType: JoinType;
lateral?: boolean;
useIndex?: string[];
forceIndex?: string[];
ignoreIndex?: string[];
}

export type BuildAliasTable<TTable extends MySqlTable | View, TAlias extends string> = TTable extends Table
Expand Down Expand Up @@ -74,6 +77,9 @@ export interface MySqlSelectConfig {
limit?: number | Placeholder;
offset?: number | Placeholder;
}[];
useIndex?: string[];
forceIndex?: string[];
ignoreIndex?: string[];
}

export type MySqlJoin<
Expand Down Expand Up @@ -116,6 +122,8 @@ export type MySqlJoinFn<
>(
table: TJoinedTable,
on: ((aliases: T['_']['selection']) => SQL | undefined) | SQL | undefined,
onIndex?: TJoinedTable extends MySqlTable ? IndexConfig
: 'Index hint configuration is allowed only for MySqlTable and not for subqueries or views',
) => MySqlJoin<T, TDynamic, TJoinType, TJoinedTable, TJoinedName>;

export type SelectedFieldsFlat = SelectedFieldsFlatBase<MySqlColumn>;
Expand Down
11 changes: 11 additions & 0 deletions drizzle-orm/src/mysql-core/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import type { Index } from './indexes.ts';
import { IndexBuilder } from './indexes.ts';
import type { PrimaryKey } from './primary-keys.ts';
import { PrimaryKeyBuilder } from './primary-keys.ts';
import type { IndexForHint } from './query-builders/select.ts';
import { MySqlTable } from './table.ts';
import { type UniqueConstraint, UniqueConstraintBuilder } from './unique-constraint.ts';
import { MySqlViewConfig } from './view-common.ts';
Expand Down Expand Up @@ -67,3 +68,13 @@ export function getViewConfig<
...view[MySqlViewConfig],
};
}

export function convertIndexToString(indexes: IndexForHint[]) {
return indexes.map((idx) => {
return typeof idx === 'object' ? idx.config.name : idx;
});
}

export function toArray<T>(value: T | T[]): T[] {
return Array.isArray(value) ? value : [value];
}
Loading

0 comments on commit 7cd9d79

Please sign in to comment.