feat(minato): support indexes (#102)

packages/core/src/database.ts

@@ -44,28 +44,28 @@ export namespace Join2 {
   export type Predicate<S, U extends Input<S>> = (args: Parameters<S, U>) => Eval.Expr<boolean>
 }
 
-type CreateMap<T, S> = { [K in keyof T]?: Create<T[K], S> }
-
-export type Create<T, S> =
+type CreateUnit<T, S> =
   | T extends Values<AtomicTypes> ? T
-  : T extends (infer I extends Values<S>)[] ? CreateMap<I, S>[] |
+  : T extends (infer I extends Values<S>)[] ? Create<I, S>[] |
     {
       $literal?: DeepPartial<I>
-      $create?: MaybeArray<CreateMap<I, S>>
-      $upsert?: MaybeArray<CreateMap<I, S>>
+      $create?: MaybeArray<Create<I, S>>
+      $upsert?: MaybeArray<Create<I, S>>
       $connect?: Query.Expr<Flatten<I>>
     }
-  : T extends Values<S> ? CreateMap<T, S> |
+  : T extends Values<S> ? Create<T, S> |
     {
       $literal?: DeepPartial<T>
-      $create?: CreateMap<T, S>
-      $upsert?: CreateMap<T, S>
+      $create?: Create<T, S>
+      $upsert?: Create<T, S>
       $connect?: Query.Expr<Flatten<T>>
     }
   : T extends (infer U)[] ? DeepPartial<U>[]
-  : T extends object ? CreateMap<T, S>
+  : T extends object ? Create<T, S>
   : T
 
+export type Create<T, S> = { [K in keyof T]?: CreateUnit<T[K], S> }
+
 function mergeQuery<T>(base: Query.FieldExpr<T>, query: Query.Expr<Flatten<T>> | ((row: Row<T>) => Query.Expr<Flatten<T>>)): Selection.Callback<T, boolean> {
   if (typeof query === 'function') {
     return (row: any) => {
@@ -129,6 +129,7 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     Object.values(fields).forEach(field => field?.transformers?.forEach(x => driver.define(x)))
 
     await driver.prepare(name)
+    await driver.prepareIndexes(name)
   }
 
   extend<K extends Keys<S>>(name: K, fields: Field.Extension<S[K], N>, config: Partial<Model.Config<FlatKeys<S[K]>>> = {}) {
@@ -189,9 +190,12 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
       }
     })
     // use relation field as primary
-    if (Array.isArray(model.primary) && model.primary.every(key => model.fields[key]?.relation)) {
-      model.primary = deduplicate(model.primary.map(key => model.fields[key]!.relation!.fields).flat())
+    if (Array.isArray(model.primary)) {
+      model.primary = deduplicate(model.primary.map(key => model.fields[key]!.relation?.fields || key).flat())
     }
+    model.unique = model.unique.map(keys => typeof keys === 'string' ? model.fields[keys]!.relation?.fields || keys
+      : keys.map(key => model.fields[key]!.relation?.fields || key).flat())
+
     this.prepareTasks[name] = this.prepare(name)
     ;(this.ctx as Context).emit('model', name)
   }

packages/core/src/driver.ts

@@ -1,4 +1,4 @@
-import { Awaitable, defineProperty, Dict, mapValues, remove } from 'cosmokit'
+import { Awaitable, deepEqual, defineProperty, Dict, mapValues, remove } from 'cosmokit'
 import { Context, Logger, Service } from 'cordis'
 import { Eval, Update } from './eval.ts'
 import { Direction, Modifier, Selection } from './selection.ts'
@@ -35,6 +35,15 @@ export namespace Driver {
     removed?: number
   }
 
+  export interface IndexDef<K extends string = string> {
+    name?: string
+    keys: { [P in K]?: 'asc' | 'desc' }
+  }
+
+  export interface Index<K extends string = string> extends IndexDef<K> {
+    unique?: boolean
+  }
+
   export interface Transformer<S = any, T = any> {
     types: Field.Type<S>[]
     dump: (value: S | null) => T | null | void
@@ -62,6 +71,9 @@ export abstract class Driver<T = any, C extends Context = Context> {
   abstract create(sel: Selection.Mutable, data: any): Promise<any>
   abstract upsert(sel: Selection.Mutable, data: any[], keys: string[]): Promise<Driver.WriteResult>
   abstract withTransaction(callback: (session?: any) => Promise<void>): Promise<void>
+  abstract getIndexes(table: string): Promise<Driver.Index[]>
+  abstract createIndex(table: string, index: Driver.Index): Promise<void>
+  abstract dropIndex(table: string, name: string): Promise<void>
 
   public database: Database<any, any, C>
   public logger: Logger
@@ -151,6 +163,20 @@ export abstract class Driver<T = any, C extends Context = Context> {
   }
 
   async _ensureSession() {}
+
+  async prepareIndexes(table: string) {
+    const oldIndexes = await this.getIndexes(table)
+    const { indexes } = this.model(table)
+    for (const index of indexes) {
+      const oldIndex = oldIndexes.find(info => info.name === index.name)
+      if (!oldIndex) {
+        await this.createIndex(table, index)
+      } else if (!deepEqual(oldIndex, index)) {
+        await this.dropIndex(table, index.name!)
+        await this.createIndex(table, index)
+      }
+    }
+  }
 }
 
 export interface MigrationHooks {

packages/core/src/model.ts

@@ -1,4 +1,4 @@
-import { clone, filterKeys, isNullable, makeArray, mapValues, MaybeArray } from 'cosmokit'
+import { clone, deepEqual, filterKeys, isNullable, makeArray, mapValues, MaybeArray } from 'cosmokit'
 import { Context } from 'cordis'
 import { Eval, Update } from './eval.ts'
 import { DeepPartial, FlatKeys, Flatten, isFlat, Keys, Row, unravel } from './utils.ts'
@@ -71,8 +71,8 @@ export namespace Relation {
       : typeof def.shared === 'string' ? { [def.shared]: def.shared }
         : Array.isArray(def.shared) ? Object.fromEntries(def.shared.map(x => [x, x]))
           : def.shared
-    const fields = def.fields ?? ((subprimary || model.name === relmodel.name || def.type === 'manyToOne'
-      || (def.type === 'oneToOne' && !makeArray(relmodel.primary).every(key => !relmodel.fields[key]?.nullable)))
+    const fields = def.fields ?? ((subprimary || def.type === 'manyToOne'
+      || (def.type === 'oneToOne' && (model.name === relmodel.name || !makeArray(relmodel.primary).every(key => !relmodel.fields[key]?.nullable))))
       ? makeArray(relmodel.primary).map(x => `${key}.${x}`) : model.primary)
     const relation: Config = {
       type: def.type,
@@ -247,6 +247,7 @@ export namespace Model {
     autoInc: boolean
     primary: MaybeArray<K>
     unique: MaybeArray<K>[]
+    indexes: (MaybeArray<K> | Driver.IndexDef<K>)[]
     foreign: {
       [P in K]?: [string, string]
     }
@@ -257,6 +258,7 @@ export interface Model extends Model.Config {}
 
 export class Model<S = any> {
   declare ctx?: Context
+  declare indexes: Driver.Index<FlatKeys<S>>[]
   fields: Field.Config<S> = {}
   migrations = new Map<Model.Migration, string[]>()
 
@@ -266,16 +268,18 @@ export class Model<S = any> {
     this.autoInc = false
     this.primary = 'id' as never
     this.unique = []
+    this.indexes = []
     this.foreign = {}
   }
 
   extend(fields: Field.Extension<S>, config?: Partial<Model.Config>): void
   extend(fields = {}, config: Partial<Model.Config> = {}) {
-    const { primary, autoInc, unique = [] as [], foreign, callback } = config
+    const { primary, autoInc, unique = [], indexes = [], foreign, callback } = config
 
     this.primary = primary || this.primary
     this.autoInc = autoInc || this.autoInc
     unique.forEach(key => this.unique.includes(key) || this.unique.push(key))
+    indexes.map(x => this.parseIndex(x)).forEach(index => (this.indexes.some(ind => deepEqual(ind, index))) || this.indexes.push(index))
     Object.assign(this.foreign, foreign)
 
     if (callback) this.migrations.set(callback, Object.keys(fields))
@@ -292,10 +296,27 @@ export class Model<S = any> {
     // check index
     this.checkIndex(this.primary)
     this.unique.forEach(index => this.checkIndex(index))
+    this.indexes.forEach(index => this.checkIndex(index))
   }
 
-  private checkIndex(index: MaybeArray<string>) {
-    for (const key of makeArray(index)) {
+  private parseIndex(index: MaybeArray<string> | Driver.Index): Driver.Index {
+    if (typeof index === 'string' || Array.isArray(index)) {
+      return {
+        name: `index:${this.name}:` + makeArray(index).join('+'),
+        unique: false,
+        keys: Object.fromEntries(makeArray(index).map(key => [key, 'asc'])),
+      }
+    } else {
+      return {
+        name: index.name ?? `index:${this.name}:` + Object.keys(index.keys).join('+'),
+        unique: index.unique ?? false,
+        keys: index.keys,
+      }
+    }
+  }
+
+  private checkIndex(index: MaybeArray<string> | Driver.Index) {
+    for (const key of typeof index === 'string' || Array.isArray(index) ? makeArray(index) : Object.keys(index.keys)) {
       if (!this.fields[key]) {
         throw new TypeError(`missing field definition for index key "${key}"`)
       }

packages/memory/src/index.ts

@@ -8,6 +8,8 @@ export class MemoryDriver extends Driver<MemoryDriver.Config> {
     _fields: [],
   }
 
+  _indexes: Dict<Dict<any>> = {}
+
   async prepare(name: string) {}
 
   async start() {
@@ -197,6 +199,21 @@ export class MemoryDriver extends Driver<MemoryDriver.Config> {
       throw e
     })
   }
+
+  async getIndexes(table: string) {
+    return Object.values(this._indexes[table] ?? {})
+  }
+
+  async createIndex(table: string, index: Driver.Index) {
+    const name = index.name ?? 'index:' + Object.entries(index.keys).map(([key, direction]) => `${key}_${direction}`).join('+')
+    this._indexes[table] ??= {}
+    this._indexes[table][name] = { name, unique: false, ...index }
+  }
+
+  async dropIndex(table: string, name: string) {
+    this._indexes[table] ??= {}
+    delete this._indexes[table][name]
+  }
 }
 
 export namespace MemoryDriver {

packages/memory/tests/index.spec.ts

@@ -16,6 +16,9 @@ describe('@minatojs/driver-memory', () => {
 
   test(database, {
     migration: false,
+    update: {
+      index: false,
+    },
     model: {
       fields: {
         cast: false,

packages/mongo/src/index.ts

@@ -82,7 +82,7 @@ export class MongoDriver extends Driver<MongoDriver.Config> {
    * https://www.mongodb.com/docs/manual/indexes/
    */
   private async _createIndexes(table: string) {
-    const { primary, unique } = this.model(table)
+    const { fields, primary, unique } = this.model(table)
     const coll = this.db.collection(table)
     const newSpecs: IndexDescription[] = []
     const oldSpecs = await coll.indexes()
@@ -96,6 +96,7 @@ export class MongoDriver extends Driver<MongoDriver.Config> {
       const name = (index ? 'unique:' : 'primary:') + keys.join('+')
       if (oldSpecs.find(spec => spec.name === name)) return
 
+      const nullable = Object.entries(fields).filter(([key]) => keys.includes(key)).every(([, field]) => field?.nullable)
       newSpecs.push({
         name,
         key: Object.fromEntries(keys.map(key => [key, 1])),
@@ -104,9 +105,11 @@ export class MongoDriver extends Driver<MongoDriver.Config> {
         // mongodb seems to not support $ne in partialFilterExpression
         // so we cannot simply use `{ $ne: null }` to filter out null values
         // below is a workaround for https://github.com/koishijs/koishi/issues/893
-        partialFilterExpression: Object.fromEntries(keys.map((key) => [key, {
-          $type: [BSONType.date, BSONType.int, BSONType.long, BSONType.string, BSONType.objectId],
-        }])),
+        ...(nullable || index > unique.length) ? {} : {
+          partialFilterExpression: Object.fromEntries(keys.map((key) => [key, {
+            $type: [BSONType.date, BSONType.int, BSONType.long, BSONType.string, BSONType.objectId],
+          }])),
+        },
       })
     })
 
@@ -208,6 +211,7 @@ export class MongoDriver extends Driver<MongoDriver.Config> {
     const [latest] = await coll.find().sort(this.getVirtualKey(table) ? '_id' : primary, -1).limit(1).toArray()
     await fields.updateOne(meta, {
       $set: { autoInc: latest ? +latest[this.getVirtualKey(table) ? '_id' : primary] : 0 },
+      $setOnInsert: { virtual: !!this.getVirtualKey(table) },
     }, { upsert: true })
   }
 
@@ -492,6 +496,34 @@ export class MongoDriver extends Driver<MongoDriver.Config> {
     }
   }
 
+  async getIndexes(table: string) {
+    const indexes = await this.db.collection(table).listIndexes().toArray()
+    return indexes.map(({ name, key, unique }) => ({
+      name,
+      unique: !!unique,
+      keys: mapValues(key, value => value === 1 ? 'asc' : value === -1 ? 'desc' : value),
+    } as Driver.Index))
+  }
+
+  async createIndex(table: string, index: Driver.Index) {
+    const keys = mapValues(index.keys, (value) => value === 'asc' ? 1 : value === 'desc' ? -1 : isNullable(value) ? 1 : value)
+    const { fields } = this.model(table)
+    const nullable = Object.keys(index.keys).every(key => fields[key]?.nullable)
+    await this.db.collection(table).createIndex(keys, {
+      name: index.name,
+      unique: !!index.unique,
+      ...nullable ? {} : {
+        partialFilterExpression: Object.fromEntries(Object.keys(index.keys).map((key) => [key, {
+          $type: [BSONType.date, BSONType.int, BSONType.long, BSONType.string, BSONType.objectId],
+        }])),
+      },
+    })
+  }
+
+  async dropIndex(table: string, name: string) {
+    await this.db.collection(table).dropIndex(name)
+  }
+
   logPipeline(table: string, pipeline: any) {
     this.logger.debug('%s %s', table, JSON.stringify(pipeline, (_, value) => typeof value === 'bigint' ? `${value}n` : value))
   }

packages/mysql/src/index.ts

@@ -32,6 +32,10 @@ interface ColumnInfo {
 interface IndexInfo {
   INDEX_NAME: string
   COLUMN_NAME: string
+  SEQ_IN_INDEX: string
+  COLLATION: 'A' | 'D'
+  NULLABLE: string
+  NON_UNIQUE: string
 }
 
 interface QueryTask {
@@ -181,11 +185,11 @@ export class MySQLDriver extends Driver<MySQLDriver.Config> {
           def += (nullable ? ' ' : ' not ') + 'null'
         }
         // blob, text, geometry or json columns cannot have default values
-        if (initial && !typedef.startsWith('text') && !typedef.endsWith('blob')) {
+        if (!isNullable(initial) && !typedef.startsWith('text') && !typedef.endsWith('blob')) {
           def += ' default ' + this.sql.escape(initial, fields[key])
         }
 
-        if (!column && initial && (typedef.startsWith('text') || typedef.endsWith('blob'))) {
+        if (!column && !isNullable(initial) && (typedef.startsWith('text') || typedef.endsWith('blob'))) {
           alterInit[key] = this.sql.escape(initial, fields[key])
         }
       }
@@ -507,6 +511,31 @@ INSERT INTO mtt VALUES(json_extract(j, concat('$[', i, ']'))); SET i=i+1; END WH
     })
   }
 
+  async getIndexes(table: string) {
+    const indexes = await this.queue<IndexInfo[]>([
+      `SELECT *`,
+      `FROM information_schema.statistics`,
+      `WHERE TABLE_SCHEMA = ? && TABLE_NAME = ?`,
+    ].join(' '), [this.config.database, table])
+    const result: Dict<Driver.Index> = {}
+    for (const { INDEX_NAME: name, COLUMN_NAME: key, COLLATION: direction, NON_UNIQUE: unique } of indexes) {
+      if (!result[name]) result[name] = { name, unique: unique !== '1', keys: {} }
+      result[name].keys[key] = direction === 'A' ? 'asc' : direction === 'D' ? 'desc' : direction
+    }
+    return Object.values(result)
+  }
+
+  async createIndex(table: string, index: Driver.Index) {
+    const keyFields = Object.entries(index.keys).map(([key, direction]) => `${escapeId(key)} ${direction ?? 'asc'}`).join(', ')
+    await this.query(
+      `ALTER TABLE ${escapeId(table)} ADD ${index.unique ? 'UNIQUE' : ''} INDEX ${index.name ? escapeId(index.name) : ''} (${keyFields})`,
+    )
+  }
+
+  async dropIndex(table: string, name: string) {
+    await this.query(`DROP INDEX ${escapeId(name)} ON ${escapeId(table)}`)
+  }
+
   private getTypeDef({ deftype: type, length, precision, scale }: Field) {
     const getIntegerType = (length = 4) => {
       if (length <= 1) return 'tinyint'