From 62a31001fc296d123bb24726d5c66d2d2f62d8e6 Mon Sep 17 00:00:00 2001 From: dengfuping Date: Tue, 17 Dec 2024 21:15:33 +0800 Subject: [PATCH] feat(mysql): Support customIndex --- drizzle-kit/src/serializer/mysqlSchema.ts | 24 +++-- drizzle-kit/src/serializer/mysqlSerializer.ts | 3 +- drizzle-kit/src/sqlgenerator.ts | 8 +- drizzle-kit/tests/mysql-generated.test.ts | 92 ++++++++++++++++++- drizzle-orm/src/mysql-core/indexes.ts | 26 +++++- 5 files changed, 135 insertions(+), 18 deletions(-) diff --git a/drizzle-kit/src/serializer/mysqlSchema.ts b/drizzle-kit/src/serializer/mysqlSchema.ts index 3a6fb9179..ad47e49ec 100644 --- a/drizzle-kit/src/serializer/mysqlSchema.ts +++ b/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,15 +220,24 @@ export type CheckConstraint = TypeOf; export type View = TypeOf; export type ViewSquashed = TypeOf; +/** + * 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(','), @@ -235,6 +245,7 @@ export const MySqlSquasher = { 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 [ diff --git a/drizzle-kit/src/serializer/mysqlSerializer.ts b/drizzle-kit/src/serializer/mysqlSerializer.ts index aaa1acb82..2865e1ff8 100644 --- a/drizzle-kit/src/serializer/mysqlSerializer.ts +++ b/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, }; -}; +}; \ No newline at end of file diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 26adaf531..f08f1a135 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/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; -`; +`; \ No newline at end of file diff --git a/drizzle-kit/tests/mysql-generated.test.ts b/drizzle-kit/tests/mysql-generated.test.ts index 3531582d0..0c2e1e9f7 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, 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"}\';', + ]); +}); diff --git a/drizzle-orm/src/mysql-core/indexes.ts b/drizzle-orm/src/mysql-core/indexes.ts index 5b73b1d30..0b16386c2 100644 --- a/drizzle-orm/src/mysql-core/indexes.ts +++ b/drizzle-orm/src/mysql-core/indexes.ts @@ -27,6 +27,11 @@ 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; @@ -34,7 +39,7 @@ 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,7 +51,7 @@ 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'; @@ -54,10 +59,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, }; } @@ -99,6 +107,10 @@ export type GetColumnsTableName = TColumns extends >[] ? TTableName : never; +export interface OptionalIndexConfig extends Omit { + columns?: IndexColumn[]; +} + export function index(name: string): IndexBuilderOn { return new IndexBuilderOn(name, false); } @@ -106,3 +118,7 @@ export function index(name: string): IndexBuilderOn { 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};` }); +}