Skip to content

Commit 7cd9d79

Browse files
AndriiShermanAndriiSherman
authored
Merge pull request #3693 from drizzle-team/feature/mysql-force-index
feature/mysql-force-index
2 parents 21dab20 + d9d234e commit 7cd9d79

File tree

7 files changed

+964
-12
lines changed

7 files changed

+964
-12
lines changed

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,18 @@ export class MySqlDialect {
244244
return orderBy && orderBy.length > 0 ? sql` order by ${sql.join(orderBy, sql`, `)}` : undefined;
245245
}
246246

247+
private buildIndex({
248+
indexes,
249+
indexFor,
250+
}: {
251+
indexes: string[] | undefined;
252+
indexFor: 'USE' | 'FORCE' | 'IGNORE';
253+
}): SQL | undefined {
254+
return indexes && indexes.length > 0
255+
? sql` ${sql.raw(indexFor)} INDEX (${sql.raw(indexes.join(`, `))})`
256+
: undefined;
257+
}
258+
247259
buildSelectQuery(
248260
{
249261
withList,
@@ -260,6 +272,9 @@ export class MySqlDialect {
260272
lockingClause,
261273
distinct,
262274
setOperators,
275+
useIndex,
276+
forceIndex,
277+
ignoreIndex,
263278
}: MySqlSelectConfig,
264279
): SQL {
265280
const fieldsList = fieldsFlat ?? orderSelectedFields<MySqlColumn>(fields);
@@ -319,10 +334,15 @@ export class MySqlDialect {
319334
const tableSchema = table[MySqlTable.Symbol.Schema];
320335
const origTableName = table[MySqlTable.Symbol.OriginalName];
321336
const alias = tableName === origTableName ? undefined : joinMeta.alias;
337+
const useIndexSql = this.buildIndex({ indexes: joinMeta.useIndex, indexFor: 'USE' });
338+
const forceIndexSql = this.buildIndex({ indexes: joinMeta.forceIndex, indexFor: 'FORCE' });
339+
const ignoreIndexSql = this.buildIndex({ indexes: joinMeta.ignoreIndex, indexFor: 'IGNORE' });
322340
joinsArray.push(
323341
sql`${sql.raw(joinMeta.joinType)} join${lateralSql} ${
324342
tableSchema ? sql`${sql.identifier(tableSchema)}.` : undefined
325-
}${sql.identifier(origTableName)}${alias && sql` ${sql.identifier(alias)}`} on ${joinMeta.on}`,
343+
}${sql.identifier(origTableName)}${useIndexSql}${forceIndexSql}${ignoreIndexSql}${
344+
alias && sql` ${sql.identifier(alias)}`
345+
} on ${joinMeta.on}`,
326346
);
327347
} else if (is(table, View)) {
328348
const viewName = table[ViewBaseConfig].name;
@@ -359,6 +379,12 @@ export class MySqlDialect {
359379

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

382+
const useIndexSql = this.buildIndex({ indexes: useIndex, indexFor: 'USE' });
383+
384+
const forceIndexSql = this.buildIndex({ indexes: forceIndex, indexFor: 'FORCE' });
385+
386+
const ignoreIndexSql = this.buildIndex({ indexes: ignoreIndex, indexFor: 'IGNORE' });
387+
362388
let lockingClausesSql;
363389
if (lockingClause) {
364390
const { config, strength } = lockingClause;
@@ -371,7 +397,7 @@ export class MySqlDialect {
371397
}
372398

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

376402
if (setOperators.length > 0) {
377403
return this.buildSetOperations(finalQuery, setOperators);

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

Lines changed: 101 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import type { MySqlColumn } from '~/mysql-core/columns/index.ts';
33
import type { MySqlDialect } from '~/mysql-core/dialect.ts';
44
import type { MySqlPreparedQueryConfig, MySqlSession, PreparedQueryHKTBase } from '~/mysql-core/session.ts';
55
import type { SubqueryWithSelection } from '~/mysql-core/subquery.ts';
6-
import type { MySqlTable } from '~/mysql-core/table.ts';
6+
import { MySqlTable } from '~/mysql-core/table.ts';
77
import { TypedQueryBuilder } from '~/query-builders/query-builder.ts';
88
import type {
99
BuildSubquerySelection,
@@ -21,9 +21,11 @@ import type { ColumnsSelection, Placeholder, Query } from '~/sql/sql.ts';
2121
import { SQL, View } from '~/sql/sql.ts';
2222
import { Subquery } from '~/subquery.ts';
2323
import { Table } from '~/table.ts';
24-
import { applyMixins, getTableColumns, getTableLikeName, haveSameKeys, type ValueOrArray } from '~/utils.ts';
25-
import { orderSelectedFields } from '~/utils.ts';
24+
import type { ValueOrArray } from '~/utils.ts';
25+
import { applyMixins, getTableColumns, getTableLikeName, haveSameKeys, orderSelectedFields } from '~/utils.ts';
2626
import { ViewBaseConfig } from '~/view-common.ts';
27+
import type { IndexBuilder } from '../indexes.ts';
28+
import { convertIndexToString, toArray } from '../utils.ts';
2729
import { MySqlViewBase } from '../view-base.ts';
2830
import type {
2931
AnyMySqlSelect,
@@ -45,6 +47,14 @@ import type {
4547
SetOperatorRightSelect,
4648
} from './select.types.ts';
4749

50+
export type IndexForHint = IndexBuilder | string;
51+
52+
export type IndexConfig = {
53+
useIndex?: IndexForHint | IndexForHint[];
54+
forceIndex?: IndexForHint | IndexForHint[];
55+
ignoreIndex?: IndexForHint | IndexForHint[];
56+
};
57+
4858
export class MySqlSelectBuilder<
4959
TSelection extends SelectedFields | undefined,
5060
TPreparedQueryHKT extends PreparedQueryHKTBase,
@@ -78,6 +88,8 @@ export class MySqlSelectBuilder<
7888

7989
from<TFrom extends MySqlTable | Subquery | MySqlViewBase | SQL>(
8090
source: TFrom,
91+
onIndex?: TFrom extends MySqlTable ? IndexConfig
92+
: 'Index hint configuration is allowed only for MySqlTable and not for subqueries or views',
8193
): CreateMySqlSelectFromBuilderMode<
8294
TBuilderMode,
8395
GetSelectTableName<TFrom>,
@@ -105,6 +117,21 @@ export class MySqlSelectBuilder<
105117
fields = getTableColumns<MySqlTable>(source);
106118
}
107119

120+
let useIndex: string[] = [];
121+
let forceIndex: string[] = [];
122+
let ignoreIndex: string[] = [];
123+
if (is(source, MySqlTable) && onIndex && typeof onIndex !== 'string') {
124+
if (onIndex.useIndex) {
125+
useIndex = convertIndexToString(toArray(onIndex.useIndex));
126+
}
127+
if (onIndex.forceIndex) {
128+
forceIndex = convertIndexToString(toArray(onIndex.forceIndex));
129+
}
130+
if (onIndex.ignoreIndex) {
131+
ignoreIndex = convertIndexToString(toArray(onIndex.ignoreIndex));
132+
}
133+
}
134+
108135
return new MySqlSelectBase(
109136
{
110137
table: source,
@@ -114,6 +141,9 @@ export class MySqlSelectBuilder<
114141
dialect: this.dialect,
115142
withList: this.withList,
116143
distinct: this.distinct,
144+
useIndex,
145+
forceIndex,
146+
ignoreIndex,
117147
},
118148
) as any;
119149
}
@@ -156,14 +186,17 @@ export abstract class MySqlSelectQueryBuilderBase<
156186
protected dialect: MySqlDialect;
157187

158188
constructor(
159-
{ table, fields, isPartialSelect, session, dialect, withList, distinct }: {
189+
{ table, fields, isPartialSelect, session, dialect, withList, distinct, useIndex, forceIndex, ignoreIndex }: {
160190
table: MySqlSelectConfig['table'];
161191
fields: MySqlSelectConfig['fields'];
162192
isPartialSelect: boolean;
163193
session: MySqlSession | undefined;
164194
dialect: MySqlDialect;
165195
withList: Subquery[];
166196
distinct: boolean | undefined;
197+
useIndex?: string[];
198+
forceIndex?: string[];
199+
ignoreIndex?: string[];
167200
},
168201
) {
169202
super();
@@ -173,6 +206,9 @@ export abstract class MySqlSelectQueryBuilderBase<
173206
fields: { ...fields },
174207
distinct,
175208
setOperators: [],
209+
useIndex,
210+
forceIndex,
211+
ignoreIndex,
176212
};
177213
this.isPartialSelect = isPartialSelect;
178214
this.session = session;
@@ -187,9 +223,13 @@ export abstract class MySqlSelectQueryBuilderBase<
187223
private createJoin<TJoinType extends JoinType>(
188224
joinType: TJoinType,
189225
): MySqlJoinFn<this, TDynamic, TJoinType> {
190-
return (
226+
return <
227+
TJoinedTable extends MySqlTable | Subquery | MySqlViewBase | SQL,
228+
>(
191229
table: MySqlTable | Subquery | MySqlViewBase | SQL,
192230
on: ((aliases: TSelection) => SQL | undefined) | SQL | undefined,
231+
onIndex?: TJoinedTable extends MySqlTable ? IndexConfig
232+
: 'Index hint configuration is allowed only for MySqlTable and not for subqueries or views',
193233
) => {
194234
const baseTableName = this.tableName;
195235
const tableName = getTableLikeName(table);
@@ -228,7 +268,22 @@ export abstract class MySqlSelectQueryBuilderBase<
228268
this.config.joins = [];
229269
}
230270

231-
this.config.joins.push({ on, table, joinType, alias: tableName });
271+
let useIndex: string[] = [];
272+
let forceIndex: string[] = [];
273+
let ignoreIndex: string[] = [];
274+
if (is(table, MySqlTable) && onIndex && typeof onIndex !== 'string') {
275+
if (onIndex.useIndex) {
276+
useIndex = convertIndexToString(toArray(onIndex.useIndex));
277+
}
278+
if (onIndex.forceIndex) {
279+
forceIndex = convertIndexToString(toArray(onIndex.forceIndex));
280+
}
281+
if (onIndex.ignoreIndex) {
282+
ignoreIndex = convertIndexToString(toArray(onIndex.ignoreIndex));
283+
}
284+
}
285+
286+
this.config.joins.push({ on, table, joinType, alias: tableName, useIndex, forceIndex, ignoreIndex });
232287

233288
if (typeof tableName === 'string') {
234289
switch (joinType) {
@@ -286,6 +341,16 @@ export abstract class MySqlSelectQueryBuilderBase<
286341
* })
287342
* .from(users)
288343
* .leftJoin(pets, eq(users.id, pets.ownerId))
344+
*
345+
* // Select userId and petId with use index hint
346+
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
347+
* userId: users.id,
348+
* petId: pets.id,
349+
* })
350+
* .from(users)
351+
* .leftJoin(pets, eq(users.id, pets.ownerId), {
352+
* useIndex: ['pets_owner_id_index']
353+
* })
289354
* ```
290355
*/
291356
leftJoin = this.createJoin('left');
@@ -315,6 +380,16 @@ export abstract class MySqlSelectQueryBuilderBase<
315380
* })
316381
* .from(users)
317382
* .rightJoin(pets, eq(users.id, pets.ownerId))
383+
*
384+
* // Select userId and petId with use index hint
385+
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
386+
* userId: users.id,
387+
* petId: pets.id,
388+
* })
389+
* .from(users)
390+
* .leftJoin(pets, eq(users.id, pets.ownerId), {
391+
* useIndex: ['pets_owner_id_index']
392+
* })
318393
* ```
319394
*/
320395
rightJoin = this.createJoin('right');
@@ -344,6 +419,16 @@ export abstract class MySqlSelectQueryBuilderBase<
344419
* })
345420
* .from(users)
346421
* .innerJoin(pets, eq(users.id, pets.ownerId))
422+
*
423+
* // Select userId and petId with use index hint
424+
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
425+
* userId: users.id,
426+
* petId: pets.id,
427+
* })
428+
* .from(users)
429+
* .leftJoin(pets, eq(users.id, pets.ownerId), {
430+
* useIndex: ['pets_owner_id_index']
431+
* })
347432
* ```
348433
*/
349434
innerJoin = this.createJoin('inner');
@@ -373,6 +458,16 @@ export abstract class MySqlSelectQueryBuilderBase<
373458
* })
374459
* .from(users)
375460
* .fullJoin(pets, eq(users.id, pets.ownerId))
461+
*
462+
* // Select userId and petId with use index hint
463+
* const usersIdsAndPetIds: { userId: number; petId: number | null }[] = await db.select({
464+
* userId: users.id,
465+
* petId: pets.id,
466+
* })
467+
* .from(users)
468+
* .leftJoin(pets, eq(users.id, pets.ownerId), {
469+
* useIndex: ['pets_owner_id_index']
470+
* })
376471
* ```
377472
*/
378473
fullJoin = this.createJoin('full');

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,17 @@ import type { Assume, ValidateShape } from '~/utils.ts';
2525
import type { MySqlPreparedQueryConfig, PreparedQueryHKTBase, PreparedQueryKind } from '../session.ts';
2626
import type { MySqlViewBase } from '../view-base.ts';
2727
import type { MySqlViewWithSelection } from '../view.ts';
28-
import type { MySqlSelectBase, MySqlSelectQueryBuilderBase } from './select.ts';
28+
import type { IndexConfig, MySqlSelectBase, MySqlSelectQueryBuilderBase } from './select.ts';
2929

3030
export interface MySqlSelectJoinConfig {
3131
on: SQL | undefined;
3232
table: MySqlTable | Subquery | MySqlViewBase | SQL;
3333
alias: string | undefined;
3434
joinType: JoinType;
3535
lateral?: boolean;
36+
useIndex?: string[];
37+
forceIndex?: string[];
38+
ignoreIndex?: string[];
3639
}
3740

3841
export type BuildAliasTable<TTable extends MySqlTable | View, TAlias extends string> = TTable extends Table
@@ -74,6 +77,9 @@ export interface MySqlSelectConfig {
7477
limit?: number | Placeholder;
7578
offset?: number | Placeholder;
7679
}[];
80+
useIndex?: string[];
81+
forceIndex?: string[];
82+
ignoreIndex?: string[];
7783
}
7884

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

121129
export type SelectedFieldsFlat = SelectedFieldsFlatBase<MySqlColumn>;

‎drizzle-orm/src/mysql-core/utils.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type { Index } from './indexes.ts';
99
import { IndexBuilder } from './indexes.ts';
1010
import type { PrimaryKey } from './primary-keys.ts';
1111
import { PrimaryKeyBuilder } from './primary-keys.ts';
12+
import type { IndexForHint } from './query-builders/select.ts';
1213
import { MySqlTable } from './table.ts';
1314
import { type UniqueConstraint, UniqueConstraintBuilder } from './unique-constraint.ts';
1415
import { MySqlViewConfig } from './view-common.ts';
@@ -67,3 +68,13 @@ export function getViewConfig<
6768
...view[MySqlViewConfig],
6869
};
6970
}
71+
72+
export function convertIndexToString(indexes: IndexForHint[]) {
73+
return indexes.map((idx) => {
74+
return typeof idx === 'object' ? idx.config.name : idx;
75+
});
76+
}
77+
78+
export function toArray<T>(value: T | T[]): T[] {
79+
return Array.isArray(value) ? value : [value];
80+
}

0 commit comments

Comments
 (0)