diff --git a/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift b/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift index 745188d..e2e2dba 100644 --- a/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift +++ b/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift @@ -176,6 +176,43 @@ extension ElasticsearchClient { } } + public func searchDocumentsCount(from indexName: String, query: Query) -> EventLoopFuture { + 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(from indexName: String, queryBody: QueryBody, size: Int = 10, offset: Int = 0, type: Document.Type = Document.self) -> EventLoopFuture> { + 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(from indexName: String, query: Query, type: Document.Type = Document.self) -> EventLoopFuture> { + 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 { do { let url = try buildURL(path: "/\(name)") diff --git a/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift b/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift index ce1499c..aa402b9 100644 --- a/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift +++ b/Sources/ElasticsearchNIOClient/Models/ESSearchRequest.swift @@ -12,6 +12,12 @@ struct ESSearchRequest: Codable { } } +struct ESComplexSearchRequest: 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 index e326504..6cdbfe2 100644 --- 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 = 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 = 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 {