feat(orm): support $el and $size operator

koishijs · Aug 14, 2021 · 9a37751 · 9a37751
1 parent 719ad69
commit 9a37751
Show file tree

Hide file tree

Showing 5 changed files with 166 additions and 88 deletions.
diff --git a/packages/koishi-core/src/database.ts b/packages/koishi-core/src/database.ts
@@ -100,21 +100,27 @@ export namespace Query {
   export type Field<T extends TableType> = string & keyof Tables[T]
   export type Index<T extends TableType> = IndexKeys<Tables[T], IndexType>
 
+  type Extract<S, T, U = S> = S extends T ? U : never
+  type Primitive = string | number
+  type Comparable = Primitive | Date
+
   export interface FieldExpr<T = any> {
-    $in?: T[]
-    $nin?: T[]
-    $eq?: T
-    $ne?: T
-    $gt?: T
-    $gte?: T
-    $lt?: T
-    $lte?: T
-    $regex?: RegExp
-    $regexFor?: string
-    $bitsAllClear?: number
-    $bitsAllSet?: number
-    $bitsAnyClear?: number
-    $bitsAnySet?: number
+    $in?: Extract<T, Primitive, T[]>
+    $nin?: Extract<T, Primitive, T[]>
+    $eq?: Extract<T, Comparable>
+    $ne?: Extract<T, Comparable>
+    $gt?: Extract<T, Comparable>
+    $gte?: Extract<T, Comparable>
+    $lt?: Extract<T, Comparable>
+    $lte?: Extract<T, Comparable>
+    $el?: T extends (infer U)[] ? FieldQuery<U> : never
+    $size?: Extract<T, any[], number>
+    $regex?: Extract<T, string, RegExp>
+    $regexFor?: Extract<T, string>
+    $bitsAllClear?: Extract<T, number>
+    $bitsAllSet?: Extract<T, number>
+    $bitsAnyClear?: Extract<T, number>
+    $bitsAnySet?: Extract<T, number>
   }
 
   export interface LogicalExpr<T = any> {
@@ -123,8 +129,8 @@ export namespace Query {
     $not?: Expr<T>
   }
 
-  export type Shorthand<T = IndexType> = T | T[] | RegExp
-  export type FieldQuery<T> = FieldExpr<T> | Shorthand<T>
+  export type Shorthand<T extends Primitive = Primitive> = T | T[] | Extract<T, string, RegExp>
+  export type FieldQuery<T = any> = FieldExpr<T> | (T extends Primitive ? Shorthand<T> : never)
   export type Expr<T = any> = LogicalExpr<T> & {
     [K in keyof T]?: FieldQuery<T[K]>
   }

diff --git a/packages/koishi-test-utils/src/memory.ts b/packages/koishi-test-utils/src/memory.ts
@@ -41,47 +41,68 @@ export class MemoryDatabase {
   }
 }
 
-const queryOperators: ([string, (data: any, value: any) => boolean])[] = Object.entries({
-  $regex: (data: RegExp, value) => data.test(value),
-  $regexFor: (data, value) => new RegExp(value, 'i').test(data),
-  $in: (data: any[], value) => data.includes(value),
-  $nin: (data: any[], value) => !data.includes(value),
-  $ne: (data, value) => value !== data,
-  $eq: (data, value) => value === data,
-  $gt: (data, value) => value > data,
-  $gte: (data, value) => value >= data,
-  $lt: (data, value) => value < data,
-  $lte: (data, value) => value <= data,
-  $bitsAllSet: (data, value) => (data & value) === data,
-  $bitsAllClear: (data, value) => (data & value) === 0,
-  $bitsAnySet: (data, value) => (data & value) !== 0,
-  $bitsAnyClear: (data, value) => (data & value) !== data,
-})
+type QueryOperators = {
+  [K in keyof Query.FieldExpr]?: (query: Query.FieldExpr[K], data: any) => boolean
+}
 
-Database.extend(MemoryDatabase, {
-  async get(name, query, modifier) {
-    function executeQuery(query: Query.Expr, data: any): boolean {
-      const entries: [string, any][] = Object.entries(query)
-      return entries.every(([key, value]) => {
-        if (key === '$and') {
-          return (value as Query.Expr[]).reduce((prev, query) => prev && executeQuery(query, data), true)
-        } else if (key === '$or') {
-          return (value as Query.Expr[]).reduce((prev, query) => prev || executeQuery(query, data), false)
-        } else if (key === '$not') {
-          return !executeQuery(value, data)
-        } else if (Array.isArray(value)) {
-          return value.includes(data[key])
-        } else if (value instanceof RegExp) {
-          return value.test(data[key])
-        } else if (typeof value === 'string' || typeof value === 'number') {
-          return value === data[key]
-        }
-        return queryOperators.reduce((prev, [prop, callback]) => {
-          return prev && (prop in value ? callback(value[prop], data[key]) : true)
-        }, true)
-      })
+const queryOperators: QueryOperators = {
+  $regex: (query, data) => query.test(data),
+  $regexFor: (query, data) => new RegExp(data, 'i').test(query),
+  $in: (query, data) => query.includes(data),
+  $nin: (query, data) => !query.includes(data),
+  $ne: (query, data) => data !== query,
+  $eq: (query, data) => data === query,
+  $gt: (query, data) => data > query,
+  $gte: (query, data) => data >= query,
+  $lt: (query, data) => data < query,
+  $lte: (query, data) => data <= query,
+  $el: (query, data) => data.some(item => executeFieldQuery(query, item)),
+  $size: (query, data) => data.length === query,
+  $bitsAllSet: (query, data) => (query & data) === query,
+  $bitsAllClear: (query, data) => (query & data) === 0,
+  $bitsAnySet: (query, data) => (query & data) !== 0,
+  $bitsAnyClear: (query, data) => (query & data) !== query,
+}
+
+function executeFieldQuery(query: Query.FieldQuery, data: any) {
+  // shorthand syntax
+  if (Array.isArray(query)) {
+    return query.includes(data)
+  } else if (query instanceof RegExp) {
+    return query.test(data)
+  } else if (typeof query === 'string' || typeof query === 'number') {
+    return query === data
+  }
+
+  // query operators
+  for (const key in queryOperators) {
+    const value = query[key]
+    if (value === undefined) continue
+    if (!queryOperators[key](value, data)) return false
+  }
+
+  return true
+}
+
+function executeQuery(query: Query.Expr, data: any): boolean {
+  const entries: [string, any][] = Object.entries(query)
+  return entries.every(([key, value]) => {
+    // execute logical query
+    if (key === '$and') {
+      return (value as Query.Expr[]).reduce((prev, query) => prev && executeQuery(query, data), true)
+    } else if (key === '$or') {
+      return (value as Query.Expr[]).reduce((prev, query) => prev || executeQuery(query, data), false)
+    } else if (key === '$not') {
+      return !executeQuery(value, data)
     }
 
+    // execute field query
+    return executeFieldQuery(value, data[key])
+  })
+}
+
+Database.extend(MemoryDatabase, {
+  async get(name, query, modifier) {
     const expr = Query.resolve(name, query)
     const { fields, limit = Infinity, offset = 0 } = Query.resolveModifier(modifier)
     return this.$table(name)

diff --git a/packages/koishi-test-utils/tests/memory.spec.ts b/packages/koishi-test-utils/tests/memory.spec.ts
@@ -18,6 +18,7 @@ interface FooData {
   id?: number
   bar: string
   baz?: number
+  list?: number[]
 }
 
 Tables.extend('foo')
@@ -62,9 +63,9 @@ describe('Memory Database', () => {
 
   describe('complex expression', () => {
     before(async () => {
-      await db.createFoo({ bar: 'awesome foo', baz: 3 })
-      await db.createFoo({ bar: 'awesome bar', baz: 4 })
-      await db.createFoo({ bar: 'awesome foo bar', baz: 7 })
+      await db.createFoo({ bar: 'awesome foo', baz: 3, list: [] })
+      await db.createFoo({ bar: 'awesome bar', baz: 4, list: [1] })
+      await db.createFoo({ bar: 'awesome foo bar', baz: 7, list: [100] })
     })
 
     after(() => {
@@ -154,6 +155,20 @@ describe('Memory Database', () => {
       })).eventually.to.have.shape([{ baz: 3 }, { baz: 4 }])
     })
 
+    it('filter data by list operations', async () => {
+      await expect(db.get('foo', {
+        list: { $size: 1 },
+      })).eventually.to.have.shape([{ baz: 4 }, { baz: 7 }])
+
+      await expect(db.get('foo', {
+        list: { $el: 100 },
+      })).eventually.to.have.shape([{ baz: 7 }])
+
+      await expect(db.get('foo', {
+        list: { $el: { $lt: 50 } },
+      })).eventually.to.have.shape([{ baz: 4 }])
+    })
+
     it('should verify `$or`, `$and` and `$not`', async () => {
       await expect(db.get('foo', {
         $or: [{

diff --git a/packages/plugin-mongo/src/index.ts b/packages/plugin-mongo/src/index.ts
@@ -1,5 +1,6 @@
 import MongoDatabase, { Config } from './database'
 import { User, Tables, Database, Field, Context, Channel, Random, pick, omit, TableType, Query } from 'koishi-core'
+import { QuerySelector } from 'mongodb'
 
 export * from './database'
 export default MongoDatabase
@@ -70,40 +71,50 @@ function unescapeKey<T extends Partial<User>>(data: T) {
   return data
 }
 
+function transformFieldQuery(query: Query.FieldQuery, key: string) {
+  // shorthand syntax
+  if (typeof query === 'string' || typeof query === 'number') {
+    return { $eq: query }
+  } else if (Array.isArray(query)) {
+    if (!query.length) return
+    return { $in: query }
+  } else if (query instanceof RegExp) {
+    return { $regex: query }
+  }
+
+  // query operators
+  const result: QuerySelector<any> = {}
+  for (const prop in query) {
+    if (prop === '$el') {
+      result.$elemMatch = transformFieldQuery(query[prop], key)
+    } else if (prop === '$regexFor') {
+      result.$expr = {
+        body(data: string, value: string) {
+          return new RegExp(data, 'i').test(value)
+        },
+        args: ['$' + key, query],
+        lang: 'js',
+      }
+    } else {
+      result[prop] = query[prop]
+    }
+  }
+  return result
+}
+
 function createFilter<T extends TableType>(name: T, _query: Query<T>) {
   function transformQuery(query: Query.Expr) {
-    const filter = {}, pending = []
+    const filter = {}
     for (const key in query) {
       const value = query[key]
       if (key === '$and' || key === '$or') {
         filter[key] = value.map(transformQuery)
       } else if (key === '$not') {
         filter[key] = transformQuery(value)
-      } else if (typeof value === 'string' || typeof value === 'number') {
-        filter[key] = { $eq: value }
-      } else if (Array.isArray(value)) {
-        if (!value.length) return
-        filter[key] = { $in: value }
       } else {
-        filter[key] = {}
-        for (const prop in value) {
-          if (prop === '$regexFor') {
-            filter[key].$expr = {
-              body(data: string, value: string) {
-                return new RegExp(data, 'i').test(value)
-              },
-              args: ['$' + key, value],
-              lang: 'js',
-            }
-          } else {
-            filter[key][prop] = value[prop]
-          }
-        }
+        filter[key] = transformFieldQuery(value, key)
       }
     }
-    if (pending.length) {
-      (filter['$and'] ||= []).push(...pending)
-    }
     return filter
   }
 

diff --git a/packages/plugin-mysql/src/index.ts b/packages/plugin-mysql/src/index.ts
@@ -27,21 +27,46 @@ function createRegExpQuery(key: string, value: RegExp) {
   return `${key} REGEXP ${escape(value.source)}`
 }
 
-function createEqualQuery(key: string, value: any) {
-  return `${key} = ${escape(value)}`
+function createElementQuery(key: string, value: any) {
+  return `FIND_IN_SET(${escape(value)}, ${key})`
 }
 
-const queryOperators: Record<string, (key: string, value: any) => string> = {
+function comparator(operator: string) {
+  return function (key: string, value: any) {
+    return `${key} ${operator} ${escape(value)}`
+  }
+}
+
+const createEqualQuery = comparator('=')
+
+type QueryOperators = {
+  [K in keyof Query.FieldExpr]?: (key: string, value: Query.FieldExpr[K]) => string
+}
+
+const queryOperators: QueryOperators = {
   $in: (key, value) => createMemberQuery(key, value, ''),
   $nin: (key, value) => createMemberQuery(key, value, ' NOT'),
   $eq: createEqualQuery,
-  $ne: (key, value) => `${key} != ${escape(value)}`,
-  $gt: (key, value) => `${key} > ${escape(value)}`,
-  $gte: (key, value) => `${key} >= ${escape(value)}`,
-  $lt: (key, value) => `${key} < ${escape(value)}`,
-  $lte: (key, value) => `${key} <= ${escape(value)}`,
+  $ne: comparator('!='),
+  $gt: comparator('>'),
+  $gte: comparator('>='),
+  $lt: comparator('<'),
+  $lte: comparator('<='),
   $regex: createRegExpQuery,
   $regexFor: (key, value) => `${escape(value)} REGEXP ${key}`,
+  $el: (key, value) => {
+    if (Array.isArray(value)) {
+      return `(${value.map(value => createElementQuery(key, value)).join(' || ')})`
+    } else if (typeof value !== 'number' && typeof value !== 'string') {
+      throw new TypeError('query expr under $el is not supported')
+    } else {
+      return createElementQuery(key, value)
+    }
+  },
+  $size: (key, value) => {
+    if (!value) return `!${key}`
+    return `LENGTH(${key}) - LENGTH(REPLACE(${key}, ",", "")) = ${escape(value)}`
+  },
   $bitsAllSet: (key, value) => `${key} & ${escape(value)} = ${escape(value)}`,
   $bitsAllClear: (key, value) => `${key} & ${escape(value)} = 0`,
   $bitsAnySet: (key, value) => `${key} & ${escape(value)} != 0`,