[MySQL] feat(mysql): Support customIndex

drizzle-kit/src/serializer/mysqlSchema.ts

@@ -9,6 +9,7 @@ const index = object({
 	using: enumType(['btree', 'hash']).optional(),
 	algorithm: enumType(['default', 'inplace', 'copy']).optional(),
 	lock: enumType(['default', 'none', 'shared', 'exclusive']).optional(),
+	raw: string().optional(),
 }).strict();
 
 const fk = object({
@@ -219,22 +220,32 @@ export type CheckConstraint = TypeOf<typeof checkConstraint>;
 export type View = TypeOf<typeof view>;
 export type ViewSquashed = TypeOf<typeof viewSquashed>;
 
+/**
+ * encode from `;` to `%3B` to avoid split error
+ */
+export const encodeRaw = (raw?: string) => raw?.replaceAll(';', '%3B');
+
+/**
+ * decode from `%3B` to `;` to recover original value
+ */
+export const decodeRaw = (raw?: string) => raw?.replaceAll('%3B', ';');
+
 export const MySqlSquasher = {
 	squashIdx: (idx: Index) => {
 		index.parse(idx);
-		return `${idx.name};${idx.columns.join(',')};${idx.isUnique};${idx.using ?? ''};${idx.algorithm ?? ''};${
-			idx.lock ?? ''
-		}`;
+		return `${idx.name};${idx.columns.join(',')};${idx.isUnique};${idx.using ?? ''};${idx.algorithm ?? ''};${idx.lock ?? ''
+			};${idx.raw ? encodeRaw(idx.raw) : ''}`;
 	},
 	unsquashIdx: (input: string): Index => {
-		const [name, columnsString, isUnique, using, algorithm, lock] = input.split(';');
+		const [name, columnsString, isUnique, using, algorithm, lock, raw] = input.split(';');
 		const destructed = {
 			name,
 			columns: columnsString.split(','),
 			isUnique: isUnique === 'true',
 			using: using ? using : undefined,
 			algorithm: algorithm ? algorithm : undefined,
 			lock: lock ? lock : undefined,
+			raw: raw ? raw : undefined,
 		};
 		return index.parse(destructed);
 	},
@@ -253,9 +264,8 @@ export const MySqlSquasher = {
 		return { name, columns: columns.split(',') };
 	},
 	squashFK: (fk: ForeignKey) => {
-		return `${fk.name};${fk.tableFrom};${fk.columnsFrom.join(',')};${fk.tableTo};${fk.columnsTo.join(',')};${
-			fk.onUpdate ?? ''
-		};${fk.onDelete ?? ''}`;
+		return `${fk.name};${fk.tableFrom};${fk.columnsFrom.join(',')};${fk.tableTo};${fk.columnsTo.join(',')};${fk.onUpdate ?? ''
+			};${fk.onDelete ?? ''}`;
 	},
 	unsquashFK: (input: string): ForeignKey => {
 		const [

drizzle-kit/src/serializer/mysqlSerializer.ts

@@ -367,6 +367,7 @@ export const generateMySqlSnapshot = (
 				using: value.config.using,
 				algorithm: value.config.algorythm,
 				lock: value.config.lock,
+				raw: value.config.raw,
 			};
 		});
 
@@ -996,4 +997,4 @@ AND
 		},
 		internal: internals,
 	};
-};
+};

drizzle-kit/src/sqlgenerator.ts

@@ -83,7 +83,7 @@ import {
 	JsonStatement,
 } from './jsonStatements';
 import { Dialect } from './schemaValidator';
-import { MySqlSquasher } from './serializer/mysqlSchema';
+import { MySqlSquasher, decodeRaw } from './serializer/mysqlSchema';
 import { PgSquasher, policy } from './serializer/pgSchema';
 import { SingleStoreSquasher } from './serializer/singlestoreSchema';
 import { SQLiteSchemaSquashed, SQLiteSquasher } from './serializer/sqliteSchema';
@@ -3480,7 +3480,7 @@ class CreateMySqlIndexConvertor extends Convertor {
 
 	convert(statement: JsonCreateIndexStatement): string {
 		// should be changed
-		const { name, columns, isUnique } = MySqlSquasher.unsquashIdx(
+		const { name, columns, isUnique, raw } = MySqlSquasher.unsquashIdx(
 			statement.data,
 		);
 		const indexPart = isUnique ? 'UNIQUE INDEX' : 'INDEX';
@@ -3495,7 +3495,7 @@ class CreateMySqlIndexConvertor extends Convertor {
 			})
 			.join(',');
 
-		return `CREATE ${indexPart} \`${name}\` ON \`${statement.tableName}\` (${uniqueString});`;
+		return raw ? decodeRaw(raw) as string : `CREATE ${indexPart} \`${name}\` ON \`${statement.tableName}\` (${uniqueString});`;
 	}
 }
 
@@ -4028,4 +4028,4 @@ insert into users(id, name, typed) values (4, 'name4', 'four');
 insert into users(id, name, typed) values (5, 'name5', 'five');
 
 drop type __venum;
-`;
+`;

drizzle-kit/tests/mysql-generated.test.ts

@@ -1,5 +1,5 @@
 import { SQL, sql } from 'drizzle-orm';
-import { int, mysqlTable, text } from 'drizzle-orm/mysql-core';
+import { bigint, customType, customIndex, int, mysqlTable, text, varchar } from 'drizzle-orm/mysql-core';
 import { expect, test } from 'vitest';
 import { diffTestSchemasMysql } from './schemaDiffer';
 
@@ -1288,3 +1288,93 @@ test('generated as string: change generated constraint', async () => {
 		"ALTER TABLE `users` ADD `gen_name` text GENERATED ALWAYS AS (`users`.`name` || 'hello') VIRTUAL;",
 	]);
 });
+
+export const vector = customType<{
+	data: string;
+	config: { length: number };
+	configRequired: true;
+}>({
+	dataType(config) {
+		return `VECTOR(${config.length})`;
+	},
+});
+
+test.only('generated as string: custom index using sql', async () => {
+	const from = {};
+	const to = {
+		users: mysqlTable('users', {
+			id: bigint({ mode: "bigint" }).autoincrement().primaryKey(),
+			name: varchar({ length: 255 }).notNull(),
+			embedding: vector("embedding", { length: 3 }),
+		}, () => {
+			return {
+				idx_embedding_oceanbase: customIndex({
+					name: 'idx_embedding_oceanbase',
+					raw: 'CREATE VECTOR INDEX idx_embedding_oceanbase ON users(embedding) WITH (distance=L2, type=hnsw);',
+				}),
+				idx_embedding_oceanbase_without_semicolon: customIndex({
+					name: 'idx_embedding_oceanbase_without_semicolon',
+					raw: 'CREATE VECTOR INDEX idx_embedding_oceanbase_without_semicolon ON users(embedding) WITH (distance=L2, type=hnsw)',
+				}),
+				idx_embedding_mysql: customIndex({
+					name: 'idx_embedding_mysql',
+					raw: 'CREATE VECTOR INDEX idx_embedding_mysql ON users(embedding) SECONDARY_ENGINE_ATTRIBUTE=\'{"type":"spann", "distance":"cosine"}\';',
+				}),
+			};
+		}),
+	};
+
+	const { statements, sqlStatements } = await diffTestSchemasMysql(
+		from,
+		to,
+		[],
+	);
+
+	expect(statements).toStrictEqual([
+		{
+			type: 'create_table',
+			tableName: 'users',
+			schema: undefined,
+			columns: [{ "name": "id", "type": "bigint", "primaryKey": false, "notNull": true, "autoincrement": true }, { "name": "name", "type": "varchar(255)", "primaryKey": false, "notNull": true, "autoincrement": false }, { "name": "embedding", "type": "VECTOR(3)", "primaryKey": false, "notNull": false, "autoincrement": false }],
+			compositePKs: ['users_id;id'],
+			compositePkName: 'users_id',
+			uniqueConstraints: [],
+			internals: { tables: {}, indexes: {} },
+			checkConstraints: []
+		},
+		{
+			type: 'create_index',
+			tableName: 'users',
+			data: 'idx_embedding_oceanbase;;false;;;;CREATE VECTOR INDEX idx_embedding_oceanbase ON users(embedding) WITH (distance=L2, type=hnsw)%3B',
+			schema: undefined,
+			internal: { tables: {}, indexes: {} }
+		},
+		{
+			type: 'create_index',
+			tableName: 'users',
+			data: 'idx_embedding_oceanbase_without_semicolon;;false;;;;CREATE VECTOR INDEX idx_embedding_oceanbase_without_semicolon ON users(embedding) WITH (distance=L2, type=hnsw)%3B',
+			schema: undefined,
+			internal: { tables: {}, indexes: {} }
+		},
+		{
+			type: 'create_index',
+			tableName: 'users',
+			data: `idx_embedding_mysql;;false;;;;CREATE VECTOR INDEX idx_embedding_mysql ON users(embedding) SECONDARY_ENGINE_ATTRIBUTE='{"type":"spann", "distance":"cosine"}'%3B`,
+			schema: undefined,
+			internal: { tables: {}, indexes: {} }
+		}
+	]);
+
+	expect(sqlStatements).toStrictEqual([
+		`CREATE TABLE \`users\` (
+	\`id\` bigint AUTO_INCREMENT NOT NULL,
+	\`name\` varchar(255) NOT NULL,
+	\`embedding\` VECTOR(3),
+	CONSTRAINT \`users_id\` PRIMARY KEY(\`id\`)
+);
+`,
+		'CREATE VECTOR INDEX idx_embedding_oceanbase ON users(embedding) WITH (distance=L2, type=hnsw);',
+		'CREATE VECTOR INDEX idx_embedding_oceanbase_without_semicolon ON users(embedding) WITH (distance=L2, type=hnsw);',
+		'CREATE VECTOR INDEX idx_embedding_mysql ON users(embedding) SECONDARY_ENGINE_ATTRIBUTE=\'{"type":"spann", "distance":"cosine"}\';',
+	]);
+});

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

@@ -27,14 +27,19 @@ interface IndexConfig {
 	 * If set, adds locks to the index creation.
 	 */
 	lock?: 'default' | 'none' | 'shared' | 'exclusive';
+
+	/**
+	 * If set, use raw sql to create index.
+	 */
+	raw?: string;
 }
 
 export type IndexColumn = MySqlColumn | SQL;
 
 export class IndexBuilderOn {
 	static readonly [entityKind]: string = 'MySqlIndexBuilderOn';
 
-	constructor(private name: string, private unique: boolean) {}
+	constructor(private name: string, private unique: boolean) { }
 
 	on(...columns: [IndexColumn, ...IndexColumn[]]): IndexBuilder {
 		return new IndexBuilder(this.name, columns, this.unique);
@@ -46,18 +51,21 @@ export interface AnyIndexBuilder {
 }
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
-export interface IndexBuilder extends AnyIndexBuilder {}
+export interface IndexBuilder extends AnyIndexBuilder { }
 
 export class IndexBuilder implements AnyIndexBuilder {
 	static readonly [entityKind]: string = 'MySqlIndexBuilder';
 
 	/** @internal */
 	config: IndexConfig;
 
-	constructor(name: string, columns: IndexColumn[], unique: boolean) {
-		this.config = {
+	constructor(name: string, columns: IndexColumn[], unique: boolean);
+	constructor(config: IndexConfig);
+
+	constructor(name: string | IndexConfig, columns?: IndexColumn[], unique?: boolean) {
+		typeof name === 'object' ? this.config = name : this.config = {
 			name,
-			columns,
+			columns: columns as IndexColumn[],
 			unique,
 		};
 	}
@@ -99,10 +107,18 @@ export type GetColumnsTableName<TColumns> = TColumns extends
 	>[] ? TTableName
 	: never;
 
+export interface OptionalIndexConfig extends Omit<IndexConfig, 'columns'> {
+	columns?: IndexColumn[];
+}
+
 export function index(name: string): IndexBuilderOn {
 	return new IndexBuilderOn(name, false);
 }
 
 export function uniqueIndex(name: string): IndexBuilderOn {
 	return new IndexBuilderOn(name, true);
 }
+
+export function customIndex(config: OptionalIndexConfig): IndexBuilder {
+	return new IndexBuilder({ unique: false, columns: [], ...config, raw: config.raw?.endsWith(';') ? config.raw : `${config.raw};` });
+}