fix(minato): support flat keys in extend, project, etc, fix #81

cordiverse · Apr 15, 2024 · 9fe00a0 · 9fe00a0
1 parent 40b5647
commit 9fe00a0
Show file tree

Hide file tree

Showing 6 changed files with 71 additions and 64 deletions.
diff --git a/packages/core/src/database.ts b/packages/core/src/database.ts
@@ -1,6 +1,6 @@
-import { Dict, Intersect, makeArray, mapValues, MaybeArray, omit, valueMap } from 'cosmokit'
+import { Dict, makeArray, mapValues, MaybeArray, omit, valueMap } from 'cosmokit'
 import { Context, Service, Spread } from 'cordis'
-import { Flatten, Indexable, Keys, randomId, Row, unravel } from './utils.ts'
+import { FlatKeys, FlatPick, Indexable, Keys, randomId, Row, unravel } from './utils.ts'
 import { Selection } from './selection.ts'
 import { Field, Model } from './model.ts'
 import { Driver } from './driver.ts'
@@ -18,11 +18,9 @@ type TableType<S, T extends TableLike<S>> =
 export namespace Join1 {
   export type Input<S> = readonly Keys<S>[]
 
-  export type Output<S, U extends Input<S>> = Intersect<
-    | U extends readonly (infer K extends Keys<S>)[]
-    ? { [P in K]: TableType<S, P> }
-    : never
-  >
+  export type Output<S, U extends Input<S>> = {
+    [P in U[number]]: TableType<S, P>
+  }
 
   type Parameters<S, U extends Input<S>> =
     | U extends readonly [infer K extends Keys<S>, ...infer R]
@@ -111,7 +109,7 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     await driver.prepare(name)
   }
 
-  extend<K extends Keys<S>>(name: K, fields: Field.Extension<S[K], N>, config: Partial<Model.Config<S[K]>> = {}) {
+  extend<K extends Keys<S>, T extends Field.Extension<S[K], N>>(name: K, fields: T, config: Partial<Model.Config<Keys<T>>> = {}) {
     let model = this.tables[name]
     if (!model) {
       model = this.tables[name] = new Model(name)
@@ -234,22 +232,22 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
   }
 
   select<T>(table: Selection<T>, query?: Query<T>): Selection<T>
-  select<T extends Keys<S>>(table: T, query?: Query<S[T]>): Selection<S[T]>
+  select<K extends Keys<S>>(table: K, query?: Query<S[K]>): Selection<S[K]>
   select(table: any, query?: any) {
     return new Selection(this.getDriver(table), table, query)
   }
 
-  join<const U extends Join1.Input<S>>(
-    tables: U,
-    callback?: Join1.Predicate<S, U>,
+  join<const X extends Join1.Input<S>>(
+    tables: X,
+    callback?: Join1.Predicate<S, X>,
     optional?: boolean[],
-  ): Selection<Join1.Output<S, U>>
+  ): Selection<Join1.Output<S, X>>
 
-  join<const U extends Join2.Input<S>>(
-    tables: U,
-    callback?: Join2.Predicate<S, U>,
-    optional?: Dict<boolean, Keys<U>>,
-  ): Selection<Join2.Output<S, U>>
+  join<X extends Join2.Input<S>>(
+    tables: X,
+    callback?: Join2.Predicate<S, X>,
+    optional?: Dict<boolean, Keys<X>>,
+  ): Selection<Join2.Output<S, X>>
 
   join(tables: any, query?: any, optional?: any) {
     if (Array.isArray(tables)) {
@@ -271,22 +269,22 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     }
   }
 
-  async get<T extends Keys<S>, K extends Keys<S[T]>>(
-    table: T,
-    query: Query<S[T]>,
-    cursor?: Driver.Cursor<K>,
-  ): Promise<Pick<S[T], K>[]> {
+  async get<K extends Keys<S>, P extends FlatKeys<S[K]> = any>(
+    table: K,
+    query: Query<S[K]>,
+    cursor?: Driver.Cursor<P>,
+  ): Promise<FlatPick<S[K], P>[]> {
     return this.select(table, query).execute(cursor)
   }
 
-  async eval<T extends Keys<S>, U>(table: T, expr: Selection.Callback<S[T], U, true>, query?: Query<S[T]>): Promise<U> {
+  async eval<K extends Keys<S>, T>(table: K, expr: Selection.Callback<S[K], T, true>, query?: Query<S[K]>): Promise<T> {
     return this.select(table, query).execute(typeof expr === 'function' ? expr : () => expr)
   }
 
-  async set<T extends Keys<S>>(
-    table: T,
-    query: Query<S[T]>,
-    update: Row.Computed<S[T], Update<S[T]>>,
+  async set<K extends Keys<S>>(
+    table: K,
+    query: Query<S[K]>,
+    update: Row.Computed<S[K], Update<S[K]>>,
   ): Promise<Driver.WriteResult> {
     const sel = this.select(table, query)
     if (typeof update === 'function') update = update(sel.row)
@@ -299,12 +297,12 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     return await sel._action('set', update).execute()
   }
 
-  async remove<T extends Keys<S>>(table: T, query: Query<S[T]>): Promise<Driver.WriteResult> {
+  async remove<K extends Keys<S>>(table: K, query: Query<S[K]>): Promise<Driver.WriteResult> {
     const sel = this.select(table, query)
     return await sel._action('remove').execute()
   }
 
-  async create<T extends Keys<S>>(table: T, data: Partial<S[T]>): Promise<S[T]> {
+  async create<K extends Keys<S>>(table: K, data: Partial<S[K]>): Promise<S[K]> {
     const sel = this.select(table)
     const { primary, autoInc } = sel.model
     if (!autoInc) {
@@ -316,10 +314,10 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     return sel._action('create', sel.model.create(data)).execute()
   }
 
-  async upsert<T extends Keys<S>>(
-    table: T,
-    upsert: Row.Computed<S[T], Update<S[T]>[]>,
-    keys?: MaybeArray<Keys<Flatten<S[T]>, Indexable>>,
+  async upsert<K extends Keys<S>>(
+    table: K,
+    upsert: Row.Computed<S[K], Update<S[K]>[]>,
+    keys?: MaybeArray<FlatKeys<S[K], Indexable>>,
   ): Promise<Driver.WriteResult> {
     const sel = this.select(table)
     if (typeof upsert === 'function') upsert = upsert(sel.row)
@@ -329,7 +327,7 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
   }
 
   async withTransaction(callback: (database: this) => Promise<void>): Promise<void>
-  async withTransaction<T extends Keys<S>>(table: T, callback: (database: this) => Promise<void>): Promise<void>
+  async withTransaction<K extends Keys<S>>(table: K, callback: (database: this) => Promise<void>): Promise<void>
   async withTransaction(arg: any, ...args: any[]) {
     if (this[kTransaction]) throw new Error('nested transactions are not supported')
     const [table, callback] = typeof arg === 'string' ? [arg, ...args] : [null, arg, ...args]
@@ -352,7 +350,7 @@ export class Database<S = {}, N = {}, C extends Context = Context> extends Servi
     await Promise.all(drivers.map(driver => driver.stop()))
   }
 
-  async drop<T extends Keys<S>>(table: T) {
+  async drop<K extends Keys<S>>(table: K) {
     await this.getDriver(table).drop(table)
   }
 

diff --git a/packages/core/src/driver.ts b/packages/core/src/driver.ts
@@ -19,11 +19,11 @@ export namespace Driver {
 
   export type Cursor<K extends string = never> = K[] | CursorOptions<K>
 
-  export interface CursorOptions<K> {
+  export interface CursorOptions<K extends string> {
     limit?: number
     offset?: number
     fields?: K[]
-    sort?: Dict<Direction>
+    sort?: Dict<Direction, K>
   }
 
   export interface WriteResult {

diff --git a/packages/core/src/model.ts b/packages/core/src/model.ts
@@ -136,19 +136,18 @@ export namespace Field {
 export namespace Model {
   export type Migration<D = any> = (database: D) => Promise<void>
 
-  export interface Config<O = {}> {
+  export interface Config<K extends string = string> {
     callback?: Migration
-    // driver?: keyof any
     autoInc: boolean
-    primary: MaybeArray<Keys<O>>
-    unique: MaybeArray<Keys<O>>[]
+    primary: MaybeArray<K>
+    unique: MaybeArray<K>[]
     foreign: {
-      [K in keyof O]?: [string, string]
+      [P in K]?: [string, string]
     }
   }
 }
 
-export interface Model<S> extends Model.Config<S> {}
+export interface Model extends Model.Config {}
 
 export class Model<S = any> {
   fields: Field.Config<S> = {}
@@ -163,7 +162,7 @@ export class Model<S = any> {
     this.foreign = {}
   }
 
-  extend(fields: Field.Extension<S>, config?: Partial<Model.Config<S>>): void
+  extend(fields: Field.Extension<S>, config?: Partial<Model.Config>): void
   extend(fields = {}, config: Partial<Model.Config> = {}) {
     const { primary, autoInc, unique = [] as [], foreign, callback } = config
 

diff --git a/packages/core/src/selection.ts b/packages/core/src/selection.ts
@@ -3,7 +3,7 @@ import { Driver } from './driver.ts'
 import { Eval, executeEval } from './eval.ts'
 import { Model } from './model.ts'
 import { Query } from './query.ts'
-import { Keys, randomId, Row } from './utils.ts'
+import { FlatKeys, FlatPick, Flatten, Keys, randomId, Row } from './utils.ts'
 import { Type } from './type.ts'
 
 declare module './eval.ts' {
@@ -129,10 +129,10 @@ class Executable<S = any, T = any> {
   }
 }
 
-type FieldLike<S = any> = Keys<S> | Selection.Callback<S>
+type FieldLike<S = any> = FlatKeys<S> | Selection.Callback<S>
 
 type FieldType<S, T extends FieldLike<S>> =
-  | T extends Keys<S> ? S[T]
+  | T extends FlatKeys<S> ? Flatten<S>[T]
   : T extends Selection.Callback<S> ? Eval<ReturnType<T>>
   : never
 
@@ -199,12 +199,12 @@ export class Selection<S = any> extends Executable<S, S[]> {
     return this
   }
 
-  groupBy<K extends Keys<S>>(fields: K | K[], query?: Selection.Callback<S, boolean>): Selection<Pick<S, K>>
-  groupBy<K extends Keys<S>, U extends Dict<FieldLike<S>>>(
+  groupBy<K extends FlatKeys<S>>(fields: K | readonly K[], query?: Selection.Callback<S, boolean>): Selection<FlatPick<S, K>>
+  groupBy<K extends FlatKeys<S>, U extends Dict<FieldLike<S>>>(
     fields: K | K[],
     extra?: U,
     query?: Selection.Callback<S, boolean>,
-  ): Selection<Pick<S, K> & FieldMap<S, U>>
+  ): Selection<FlatPick<S, K> & FieldMap<S, U>>
 
   groupBy<K extends Dict<FieldLike<S>>>(fields: K, query?: Selection.Callback<S, boolean>): Selection<FieldMap<S, K>>
   groupBy<K extends Dict<FieldLike<S>>, U extends Dict<FieldLike<S>>>(
@@ -227,7 +227,7 @@ export class Selection<S = any> extends Executable<S, S[]> {
     return this
   }
 
-  project<K extends Keys<S>>(fields: K | K[]): Selection<Pick<S, K>>
+  project<K extends FlatKeys<S>>(fields: K | readonly K[]): Selection<FlatPick<S, K>>
   project<U extends Dict<FieldLike<S>>>(fields: U): Selection<FieldMap<S, U>>
   project(fields: Keys<S>[] | Dict<FieldLike<S>>) {
     this.args[0].fields = this.resolveFields(fields)
@@ -243,15 +243,15 @@ export class Selection<S = any> extends Executable<S, S[]> {
   evaluate(): Eval.Expr<S[], boolean>
   evaluate(callback?: any): any {
     const selection = new Selection(this.driver, this)
-    if (!callback) callback = (row) => Eval.array(Eval.object(row))
+    if (!callback) callback = (row: any) => Eval.array(Eval.object(row))
     const expr = this.resolveField(callback)
     if (expr['$']) defineProperty(expr, Type.kType, Type.Array(Type.fromTerm(expr)))
     return Eval.exec(selection._action('eval', expr))
   }
 
-  execute<K extends Keys<S> = Keys<S>>(cursor?: Driver.Cursor<K>): Promise<Pick<S, K>[]>
+  execute<K extends FlatKeys<S> = any>(cursor?: Driver.Cursor<K>): Promise<Extract<FlatPick<S, K>, S>[]>
   execute<T>(callback: Selection.Callback<S, T, true>): Promise<T>
-  execute(cursor?: any) {
+  async execute(cursor?: any) {
     if (typeof cursor === 'function') {
       const selection = new Selection(this.driver, this)
       return selection._action('eval', this.resolveField(cursor)).execute()
@@ -269,12 +269,13 @@ export class Selection<S = any> extends Executable<S, S[]> {
         this.orderBy(field as any, cursor.sort[field])
       }
     }
-    if (cursor.fields) {
-      return super.execute().then(
-        rows => rows.map(row => filterKeys(row as any, key => (cursor.fields as string[]).some(k => k === key || k.startsWith(`${key}.`)))),
-      )
-    }
-    return super.execute()
+    const rows = await super.execute()
+    if (!cursor.fields) return rows
+    return rows.map((row) => {
+      return filterKeys(row as any, key => {
+        return (cursor.fields as string[]).some(k => k === key || k.startsWith(`${key}.`))
+      })
+    })
   }
 }
 

diff --git a/packages/core/src/utils.ts b/packages/core/src/utils.ts
@@ -4,9 +4,18 @@ import { Eval } from './eval.ts'
 export type Values<S> = S[keyof S]
 
 export type Keys<O, T = any> = Values<{
-  [K in keyof O]: O[K] extends T | undefined ? K : never
+  [P in keyof O]: O[P] extends T | undefined ? P : never
 }> & string
 
+export type FlatKeys<O, T = any> = Keys<Flatten<O>, T>
+
+export type FlatPick<O, K extends FlatKeys<O>> = {
+  [P in string & keyof O as K extends P | `${P}.${any}` ? P : never]:
+    | P extends K
+    ? O[P]
+    : FlatPick<O[P], Extract<K extends `${any}.${infer R}` ? R : never, FlatKeys<O[P]>>>
+}
+
 export interface AtomicTypes {
   Number: number
   String: string
@@ -21,7 +30,7 @@ export interface AtomicTypes {
 }
 
 export type Indexable = string | number
-export type Comparable = string | number | boolean | Date
+export type Comparable = string | number | boolean | bigint | Date
 
 type FlatWrap<S, A extends 0[], P extends string> = { [K in P]?: S }
   // rule out atomic types

diff --git a/packages/tests/src/model.ts b/packages/tests/src/model.ts
@@ -8,7 +8,7 @@ interface DType {
   num?: number
   double?: number
   decimal?: number
-  int64?: BigInt
+  int64?: bigint
   bool?: boolean
   list?: string[]
   array?: number[]