Merge pull request #40 from paulrolfe/pr/operators

Operators: `.contains`, `.startsWith`, `.endsWith`
orlandos-nl · Sep 9, 2016 · c16e2f5 · c16e2f5
2 parents 6e471c7 + 23e6173
commit c16e2f5
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/Sources/MongoKitten/QueryBuilder.swift b/Sources/MongoKitten/QueryBuilder.swift
@@ -244,6 +244,12 @@ public indirect enum AQT {
             return ["$or": .array(Document(array: expressions)) ]
         case .not(let aqt):
             return ["$not": ~aqt.document]
+        case .contains(let key, let val):
+            return [key: .regularExpression(pattern: val, options: "")]
+        case .startsWith(let key, let val):
+            return [key: .regularExpression(pattern: "^\(val)", options: "m")]
+        case .endsWith(let key, let val):
+            return [key: .regularExpression(pattern: "\(val)$", options: "m")]
         case .nothing:
             return []
         }
@@ -281,6 +287,15 @@ public indirect enum AQT {
 
     /// Whether nothing needs to be matched. Is always true and just a placeholder
     case nothing
+
+    /// Whether the String value within the `key` contains this `String`.
+    case contains(key: String, val: String)
+
+    /// Whether the String value within the `key` starts with this `String`.
+    case startsWith(key: String, val: String)
+
+    /// Whether the String value within the `key` ends with this `String`.
+    case endsWith(key: String, val: String)
 }
 
 /// The protocol all queries need to comply to
@@ -507,6 +522,27 @@ extension Document {
             return false
         case .not(let aqt):
             return !self.matches(query: Query(aqt: aqt))
+        case .contains(let key, let val):
+            switch doc[key] {
+            case .string(let stringVal):
+                return stringVal.contains(val)
+            default:
+                return false
+            }
+        case .startsWith(let key, let val):
+            switch doc[key] {
+            case .string(let stringVal):
+                return stringVal.hasPrefix(val)
+            default:
+                return false
+            }
+        case .endsWith(let key, let val):
+            switch doc[key] {
+            case .string(let stringVal):
+                return stringVal.hasSuffix(val)
+            default:
+                return false
+            }
         case .nothing:
             return true
         }

diff --git a/Tests/MongoKittenTests/CollectionTests.swift b/Tests/MongoKittenTests/CollectionTests.swift
@@ -72,6 +72,28 @@ class CollectionTests: XCTestCase {
         XCTAssertEqual(response.count, 2)
 
         XCTAssertEqual(response.first, response2)
+
+        runContainsQuery()
+        runStartsWithQuery()
+        runEndsWithQuery()
+    }
+
+    private func runContainsQuery() {
+        let query = Query(aqt: .contains(key: "username", val: "ar"))
+        let response = Array(try! TestManager.wcol.find(matching: query))
+        XCTAssert(response.count == 2)
+    }
+
+    private func runStartsWithQuery() {
+        let query = Query(aqt: .startsWith(key: "username", val: "har"))
+        let response = Array(try! TestManager.wcol.find(matching: query))
+        XCTAssert(response.count == 2)
+    }
+
+    private func runEndsWithQuery() {
+        let query = Query(aqt: .endsWith(key: "username", val: "rrie"))
+        let response = Array(try! TestManager.wcol.find(matching: query))
+        XCTAssert(response.count == 2)
     }
 
     func testAggregate() {