From 7ae830b41c5092f8594c85f1b065c3a2cf2d41db Mon Sep 17 00:00:00 2001
From: Benjamin Hutchins <ben@hutchins.co>
Date: Thu, 14 Sep 2023 13:46:02 -0400
Subject: [PATCH] feat: improve searching using contains operator for sets

Fixes #667
---
 src/query/condition.ts   |  4 ++--
 src/query/filters.ts     |  8 ++++++--
 src/query/search.spec.ts | 19 ++++++++++++++++---
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/query/condition.ts b/src/query/condition.ts
index 5323186b..9ae6766a 100644
--- a/src/query/condition.ts
+++ b/src/query/condition.ts
@@ -1,5 +1,5 @@
 import { type Table } from '../dyngoose'
-import { type Filter } from './filters'
+import { type ContainsType, type Filter } from './filters'
 import { type MagicSearch } from './search'
 
 export class Condition<T extends Table, Attr, AttributeValueType> {
@@ -123,7 +123,7 @@ export class Condition<T extends Table, Attr, AttributeValueType> {
    *
    * When using `.not().contains(value)` this checks for the absence of a subsequence, or absence of a value in a set.
    */
-  contains(value: AttributeValueType extends Array<infer E> ? E : AttributeValueType): MagicSearch<T> {
+  contains(value: ContainsType<AttributeValueType>): MagicSearch<T> {
     if (this._not) {
       this.filter = ['not contains', value as any]
     } else {
diff --git a/src/query/filters.ts b/src/query/filters.ts
index 543b21c1..10c0d80b 100644
--- a/src/query/filters.ts
+++ b/src/query/filters.ts
@@ -2,6 +2,10 @@ import { type Table } from '../table'
 
 export type AttributeNames<T extends Table> = Exclude<Exclude<keyof T, keyof Table>, () => any>
 
+export type ContainsType<Type> = Type extends Array<infer E>
+  ? E
+  : Type extends Set<infer E> ? E : Type
+
 export type Filter<Type> =
   ['=', Type] |
   ['<>', Type] | // not equals
@@ -13,8 +17,8 @@ export type Filter<Type> =
   ['between', Type, Type] |
   ['includes', Type[]] |
   ['excludes', Type[]] |
-  ['contains', Type] | // contains can be used on a list or a string attribute
-  ['not contains', Type] | // not contains can be used on a list or a string attribute
+  ['contains', ContainsType<Type>] | // contains can be used on a list, string, or set attributes
+  ['not contains', ContainsType<Type>] | // not contains can be used on a list, string, or set attributes
   ['null'] |
   ['not null'] |
   ['exists'] |
diff --git a/src/query/search.spec.ts b/src/query/search.spec.ts
index a39046ce..c95b45fd 100644
--- a/src/query/search.spec.ts
+++ b/src/query/search.spec.ts
@@ -4,10 +4,11 @@ import { MagicSearch } from './search'
 
 describe('Query/Search', () => {
   before(async () => {
+    const sets = { testStringSet: new Set(['search']), testStringSetArray: ['search'] }
     await TestableTable.documentClient.batchPut([
-      TestableTable.new({ id: 500, title: 'Table.search 0', lowercaseString: 'table search 0' }),
-      TestableTable.new({ id: 501, title: 'Table.search 1', lowercaseString: 'table search 1' }),
-      TestableTable.new({ id: 502, title: 'Table.search 2', lowercaseString: 'table search 2' }),
+      TestableTable.new({ id: 500, title: 'Table.search 0', lowercaseString: 'table search 0', ...sets }),
+      TestableTable.new({ id: 501, title: 'Table.search 1', lowercaseString: 'table search 1', ...sets }),
+      TestableTable.new({ id: 502, title: 'Table.search 2', lowercaseString: 'table search 2', ...sets }),
       TestableTable.new({ id: 503, title: 'Table.search 3', lowercaseString: 'table search 3' }),
       TestableTable.new({ id: 504, title: 'Table.search 4', lowercaseString: 'reject the search 4' }),
       TestableTable.new({ id: 504, title: 'Table.search 5', lowercaseString: 'magic' }),
@@ -91,6 +92,18 @@ describe('Query/Search', () => {
     expect(result.count).to.eq(3)
   })
 
+  it('should support filtering on sets', async () => {
+    const search = new MagicSearch<TestableTable>(TestableTable)
+      .filter('testStringSet').contains('search')
+      .and()
+      .filter('testStringSetArray').contains('search')
+    const input = search.getInput()
+    expect(input.IndexName).to.be.a('undefined')
+    expect(input.FilterExpression).to.eq('contains(#a0, :v0) AND contains(#a1, :v0)')
+    const result = await search.exec()
+    expect(result.count).to.eq(3)
+  })
+
   it('ConsistentRead defaults to false', async () => {
     const search = new MagicSearch<TestableTable>(TestableTable)
     const input = search.getInput()