[MySQL] feat(mysql): Support vectorIndex and secondaryEngineAttribute

drizzle-kit/src/serializer/mysqlSchema.ts

@@ -6,9 +6,11 @@ const index = object({
 	name: string(),
 	columns: string().array(),
 	isUnique: boolean(),
+	vector: boolean().optional(),
 	using: enumType(['btree', 'hash']).optional(),
 	algorithm: enumType(['default', 'inplace', 'copy']).optional(),
 	lock: enumType(['default', 'none', 'shared', 'exclusive']).optional(),
+	secondaryEngineAttribute: string().optional(),
 }).strict();
 
 const fk = object({
@@ -222,19 +224,19 @@ export type ViewSquashed = TypeOf<typeof viewSquashed>;
 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.vector};${idx.using ?? ''};${idx.algorithm ?? ''};${idx.lock ?? ''};${idx.secondaryEngineAttribute ?? ''}`;
 	},
 	unsquashIdx: (input: string): Index => {
-		const [name, columnsString, isUnique, using, algorithm, lock] = input.split(';');
+		const [name, columnsString, isUnique, vector, using, algorithm, lock, secondaryEngineAttribute] = input.split(';');
 		const destructed = {
 			name,
 			columns: columnsString.split(','),
 			isUnique: isUnique === 'true',
+			vector: vector === 'true',
 			using: using ? using : undefined,
 			algorithm: algorithm ? algorithm : undefined,
 			lock: lock ? lock : undefined,
+			secondaryEngineAttribute: secondaryEngineAttribute ? secondaryEngineAttribute : undefined,
 		};
 		return index.parse(destructed);
 	},

drizzle-kit/src/serializer/mysqlSerializer.ts

@@ -364,9 +364,11 @@ export const generateMySqlSnapshot = (
 				name,
 				columns: indexColumns,
 				isUnique: value.config.unique ?? false,
+				vector: value.config.vector ?? false,
 				using: value.config.using,
 				algorithm: value.config.algorythm,
 				lock: value.config.lock,
+				secondaryEngineAttribute: value.config.secondaryEngineAttribute,
 			};
 		});
 

drizzle-kit/src/sqlgenerator.ts

@@ -3480,10 +3480,10 @@ class CreateMySqlIndexConvertor extends Convertor {
 
 	convert(statement: JsonCreateIndexStatement): string {
 		// should be changed
-		const { name, columns, isUnique } = MySqlSquasher.unsquashIdx(
+		const { name, columns, isUnique, vector, secondaryEngineAttribute } = MySqlSquasher.unsquashIdx(
 			statement.data,
 		);
-		const indexPart = isUnique ? 'UNIQUE INDEX' : 'INDEX';
+		const indexPart = isUnique ? 'UNIQUE INDEX' : vector ? 'VECTOR INDEX' : 'INDEX';
 
 		const uniqueString = columns
 			.map((it) => {
@@ -3494,8 +3494,9 @@ class CreateMySqlIndexConvertor extends Convertor {
 					: `\`${it}\``;
 			})
 			.join(',');
+		const secondaryEngineAttributeString = secondaryEngineAttribute ? ` SECONDARY_ENGINE_ATTRIBUTE='${secondaryEngineAttribute}'` : '';
 
-		return `CREATE ${indexPart} \`${name}\` ON \`${statement.tableName}\` (${uniqueString});`;
+		return `CREATE ${indexPart} \`${name}\` ON \`${statement.tableName}\` (${uniqueString})${secondaryEngineAttributeString};`;
 	}
 }
 

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, index, int, mysqlTable, text, varchar, vectorIndex } from 'drizzle-orm/mysql-core';
 import { expect, test } from 'vitest';
 import { diffTestSchemasMysql } from './schemaDiffer';
 
@@ -1288,3 +1288,91 @@ 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: vector index and secondaryEngineAttribute', async () => {
+	const from = {};
+	const to = {
+		users: mysqlTable('users', {
+			id: bigint({ mode: "bigint" }).autoincrement().primaryKey(),
+			name: varchar({ length: 255 }).notNull(),
+			embedding: vector("embedding", { length: 3 }),
+		}, (table) => {
+			return {
+				idx_embedding: vectorIndex({
+					name: "idx_embedding",
+					secondaryEngineAttribute: '{"type":"spann", "distance":"cosine"}',
+				}).on(table.embedding),
+			};
+		}),
+	};
+
+	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;embedding;false;true;;;;{"type":"spann", "distance":"cosine"}',
+			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` ON `users` (`embedding`) SECONDARY_ENGINE_ATTRIBUTE=\'{"type":"spann", "distance":"cosine"}\';',
+	]);
+});

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

@@ -3,6 +3,8 @@ import type { SQL } from '~/sql/sql.ts';
 import type { AnyMySqlColumn, MySqlColumn } from './columns/index.ts';
 import type { MySqlTable } from './table.ts';
 
+export type MysqlIndexMethod = 'btree' | 'hash' | 'hnsw' | (string & {});
+
 interface IndexConfig {
 	name: string;
 
@@ -13,6 +15,11 @@ interface IndexConfig {
 	 */
 	unique?: boolean;
 
+	/**
+	 * If true, the index will be created as `create vector index` instead of `create index`.
+	 */
+	vector?: boolean;
+
 	/**
 	 * If set, the index will be created as `create index ... using { 'btree' | 'hash' }`.
 	 */
@@ -27,17 +34,22 @@ interface IndexConfig {
 	 * If set, adds locks to the index creation.
 	 */
 	lock?: 'default' | 'none' | 'shared' | 'exclusive';
+
+	secondaryEngineAttribute?: 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 | Omit<IndexConfig, 'columns'>, private unique: boolean) { }
 
 	on(...columns: [IndexColumn, ...IndexColumn[]]): IndexBuilder {
-		return new IndexBuilder(this.name, columns, this.unique);
+		return typeof this.name === 'object' ? new IndexBuilder({
+			...this.name,
+			columns,
+		}) : new IndexBuilder(this.name, columns, this.unique);
 	}
 }
 
@@ -54,10 +66,13 @@ export class IndexBuilder implements AnyIndexBuilder {
 	/** @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,
 		};
 	}
@@ -106,3 +121,7 @@ export function index(name: string): IndexBuilderOn {
 export function uniqueIndex(name: string): IndexBuilderOn {
 	return new IndexBuilderOn(name, true);
 }
+
+export function vectorIndex(config: Omit<IndexConfig, 'columns' | 'vector'>): IndexBuilderOn {
+	return new IndexBuilderOn({ ...config, vector: true }, false);
+}