Merge pull request #5 from brokenhandsio/complex-search

Complex Search API
brokenhandsio · May 18, 2021 · 4ef4e9a · 4ef4e9a
2 parents 0ddebc2 + 470e45d
commit 4ef4e9a
Show file tree

Hide file tree

Showing 3 changed files with 138 additions and 0 deletions.
diff --git a/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift b/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift
@@ -176,6 +176,43 @@ extension ElasticsearchClient {
         }
     }
 
+    public func searchDocumentsCount<Query: Encodable>(from indexName: String, query: Query) -> EventLoopFuture<ESCountResponse> {
+        do {
+            let url = try buildURL(path: "/\(indexName)/_count")
+            let body = try ByteBuffer(data: self.jsonEncoder.encode(query))
+            var headers = HTTPHeaders()
+            headers.add(name: "content-type", value: "application/json")
+            return sendRequest(url: url, method: .GET, headers: headers, body: body)
+        } catch {
+            return self.eventLoop.makeFailedFuture(error)
+        }
+    }
+
+    public func searchDocumentsPaginated<Document: Decodable, QueryBody: Encodable>(from indexName: String, queryBody: QueryBody, size: Int = 10, offset: Int = 0, type: Document.Type = Document.self) -> EventLoopFuture<ESGetMultipleDocumentsResponse<Document>> {
+        do {
+            let url = try buildURL(path: "/\(indexName)/_search")
+            let queryBody = ESComplexSearchRequest(from: offset, size: size, query: queryBody)
+            let body = try ByteBuffer(data: self.jsonEncoder.encode(queryBody))
+            var headers = HTTPHeaders()
+            headers.add(name: "content-type", value: "application/json")
+            return sendRequest(url: url, method: .GET, headers: headers, body: body)
+        } catch {
+            return self.eventLoop.makeFailedFuture(error)
+        }
+    }
+
+    public func customSearch<Document: Decodable, Query: Encodable>(from indexName: String, query: Query, type: Document.Type = Document.self) -> EventLoopFuture<ESGetMultipleDocumentsResponse<Document>> {
+        do {
+            let url = try buildURL(path: "/\(indexName)/_search")
+            let body = try ByteBuffer(data: self.jsonEncoder.encode(query))
+            var headers = HTTPHeaders()
+            headers.add(name: "content-type", value: "application/json")
+            return sendRequest(url: url, method: .GET, headers: headers, body: body)
+        } catch {
+            return self.eventLoop.makeFailedFuture(error)
+        }
+    }
+
     public func deleteIndex(_ name: String) -> EventLoopFuture<ESDeleteIndexResponse> {
         do {
             let url = try buildURL(path: "/\(name)")

diff --git a/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift b/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift
@@ -12,6 +12,12 @@ struct ESSearchRequest: Codable {
     }
 }
 
+struct ESComplexSearchRequest<Query: Encodable>: Encodable {
+    let from: Int
+    let size: Int
+    let query: Query
+}
+
 struct ESSearchQueryString: Codable {
     let queryString: ESSearchQuery
 

diff --git a/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift b/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift
@@ -328,6 +328,101 @@ class ElasticSearchIntegrationTests: XCTestCase {
         XCTAssertEqual(retrievedItem.source.count, 1)
     }
 
+    func testCountWithQueryBody() throws {
+        try setupItems()
+
+        struct SearchQuery: Encodable {
+            let query: QueryBody
+        }
+
+        struct QueryBody: Encodable {
+            let queryString: QueryString
+
+            enum CodingKeys: String, CodingKey {
+                case queryString = "query_string"
+            }
+        }
+
+        struct QueryString: Encodable {
+            let query: String
+        }
+
+        let queryString = QueryString(query: "Apples")
+        let queryBody = QueryBody(queryString: queryString)
+        let searchQuery = SearchQuery(query: queryBody)
+        let results = try client.searchDocumentsCount(from: indexName, query: searchQuery).wait()
+        XCTAssertEqual(results.count, 5)
+    }
+
+    func testPaginationQueryWithQueryBody() throws {
+        for index in 1...100 {
+            let name = "Some \(index) Apples"
+            let item = SomeItem(id: UUID(), name: name)
+            _ = try client.createDocument(item, in: self.indexName).wait()
+        }
+
+        // This is required for ES to settle and load the indexes to return the right results
+        Thread.sleep(forTimeInterval: 1.0)
+
+        struct QueryBody: Encodable {
+            let queryString: QueryString
+
+            enum CodingKeys: String, CodingKey {
+                case queryString = "query_string"
+            }
+        }
+
+        struct QueryString: Encodable {
+            let query: String
+        }
+
+        let queryString = QueryString(query: "Apples")
+        let queryBody = QueryBody(queryString: queryString)
+
+        let results: ESGetMultipleDocumentsResponse<SomeItem> = try client.searchDocumentsPaginated(from: indexName, queryBody: queryBody, size: 20, offset: 10).wait()
+        XCTAssertEqual(results.hits.hits.count, 20)
+        XCTAssertTrue(results.hits.hits.contains(where: { $0.source.name == "Some 11 Apples" }))
+        XCTAssertTrue(results.hits.hits.contains(where: { $0.source.name == "Some 29 Apples" }))
+    }
+
+    func testCustomSearch() throws {
+        for index in 1...100 {
+            let name = "Some \(index) Apples"
+            let item = SomeItem(id: UUID(), name: name)
+            _ = try client.createDocument(item, in: self.indexName).wait()
+        }
+
+        // This is required for ES to settle and load the indexes to return the right results
+        Thread.sleep(forTimeInterval: 1.0)
+
+        struct Query: Encodable {
+            let query: QueryBody
+            let from: Int
+            let size: Int
+        }
+
+        struct QueryBody: Encodable {
+            let queryString: QueryString
+
+            enum CodingKeys: String, CodingKey {
+                case queryString = "query_string"
+            }
+        }
+
+        struct QueryString: Encodable {
+            let query: String
+        }
+
+        let queryString = QueryString(query: "Apples")
+        let queryBody = QueryBody(queryString: queryString)
+        let query = Query(query: queryBody, from: 10, size: 20)
+
+        let results: ESGetMultipleDocumentsResponse<SomeItem> = try client.customSearch(from: indexName, query: query).wait()
+        XCTAssertEqual(results.hits.hits.count, 20)
+        XCTAssertTrue(results.hits.hits.contains(where: { $0.source.name == "Some 11 Apples" }))
+        XCTAssertTrue(results.hits.hits.contains(where: { $0.source.name == "Some 29 Apples" }))
+    }
+
     // MARK: - Private
     private func setupItems() throws {
         for index in 1...10 {