From 1b55241964a8fe38a143b3e673a008621f1ee090 Mon Sep 17 00:00:00 2001 From: dengfuping Date: Mon, 16 Dec 2024 18:28:57 +0800 Subject: [PATCH] feat(mysql): Support vectorIndex and secondaryEngineAttribute --- drizzle-kit/src/serializer/mysqlSchema.ts | 10 ++- drizzle-kit/src/serializer/mysqlSerializer.ts | 2 + drizzle-kit/src/sqlgenerator.ts | 7 +- drizzle-kit/tests/mysql-generated.test.ts | 90 ++++++++++++++++++- drizzle-orm/src/mysql-core/indexes.ts | 29 ++++-- 5 files changed, 125 insertions(+), 13 deletions(-) diff --git a/drizzle-kit/src/serializer/mysqlSchema.ts b/drizzle-kit/src/serializer/mysqlSchema.ts index 3a6fb9179..5013c1c28 100644 --- a/drizzle-kit/src/serializer/mysqlSchema.ts +++ b/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; 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); }, diff --git a/drizzle-kit/src/serializer/mysqlSerializer.ts b/drizzle-kit/src/serializer/mysqlSerializer.ts index aaa1acb82..0f22a2938 100644 --- a/drizzle-kit/src/serializer/mysqlSerializer.ts +++ b/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, }; }); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 26adaf531..5ec0a71a6 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/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};`; } } diff --git a/drizzle-kit/tests/mysql-generated.test.ts b/drizzle-kit/tests/mysql-generated.test.ts index 3531582d0..9a83173f0 100644 --- a/drizzle-kit/tests/mysql-generated.test.ts +++ b/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"}\';', + ]); +}); diff --git a/drizzle-orm/src/mysql-core/indexes.ts b/drizzle-orm/src/mysql-core/indexes.ts index 5b73b1d30..19a831c80 100644 --- a/drizzle-orm/src/mysql-core/indexes.ts +++ b/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,6 +34,8 @@ interface IndexConfig { * If set, adds locks to the index creation. */ lock?: 'default' | 'none' | 'shared' | 'exclusive'; + + secondaryEngineAttribute?: string; } export type IndexColumn = MySqlColumn | SQL; @@ -34,10 +43,13 @@ 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, 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): IndexBuilderOn { + return new IndexBuilderOn({ ...config, vector: true }, false); +}