diff --git a/README.md b/README.md index 359dc31..38b6688 100644 --- a/README.md +++ b/README.md @@ -37,9 +37,11 @@ Currently the library supports: * Document delete * Document search * Document count +* Document retrieve * Bulk create/update/delete/index * Index delete * Index exists +* Scripting If you'd like to add extra functionality, either [open an issue](https://github.com/brokenhandsio/elasticsearch-nio-client/issues/new) and raise a PR. Any contributions are gratefully accepted! diff --git a/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift b/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift index c73b650..745188d 100644 --- a/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift +++ b/Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift @@ -3,6 +3,15 @@ import NIO import NIOHTTP1 extension ElasticsearchClient { + public func get(id: String, from indexName: String) -> EventLoopFuture> { + do { + let url = try buildURL(path: "/\(indexName)/_doc/\(id)") + return sendRequest(url: url, method: .GET, headers: .init(), body: nil) + } catch { + return self.eventLoop.makeFailedFuture(error) + } + } + public func bulk(_ operations: [ESBulkOperation]) -> EventLoopFuture { guard operations.count > 0 else { return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No operations to perform for the bulk API", status: nil)) @@ -33,7 +42,7 @@ extension ElasticsearchClient { bodyString.append("\(deleteLineString)\n") case .index: guard let document = operation.document else { - return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No document provided for create bulk operation", status: nil)) + return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No document provided for index bulk operation", status: nil)) } let indexInfo = BulkIndex(index: bulkOperationBody) let indexLine = try self.jsonEncoder.encode(indexInfo) @@ -44,7 +53,7 @@ extension ElasticsearchClient { bodyString.append("\(indexLineString)\n\(dataLineString)\n") case .update: guard let document = operation.document else { - return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No document provided for create bulk operation", status: nil)) + return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No document provided for update bulk operation", status: nil)) } let updateInfo = BulkUpdate(update: bulkOperationBody) let updateLine = try self.jsonEncoder.encode(updateInfo) @@ -53,6 +62,17 @@ extension ElasticsearchClient { throw ElasticSearchClientError(message: "Failed to convert bulk data from Data to String", status: nil) } bodyString.append("\(updateLineString)\n\(dataLineString)\n") + case .updateScript: + guard let document = operation.document else { + return self.eventLoop.makeFailedFuture(ElasticSearchClientError(message: "No script provided for update script bulk operation", status: nil)) + } + let updateInfo = BulkUpdateScript(update: bulkOperationBody) + let updateLine = try self.jsonEncoder.encode(updateInfo) + let dataLine = try self.jsonEncoder.encode(BulkUpdateScriptDocument(script: document)) + guard let updateLineString = String(data: updateLine, encoding: .utf8), let dataLineString = String(data: dataLine, encoding: .utf8) else { + throw ElasticSearchClientError(message: "Failed to convert bulk data from Data to String", status: nil) + } + bodyString.append("\(updateLineString)\n\(dataLineString)\n") } } let body = ByteBuffer(string: bodyString) @@ -100,6 +120,18 @@ extension ElasticsearchClient { } } + public func updateDocumentWithScript(_ script: Script, id: String, in indexName: String) -> EventLoopFuture { + do { + let url = try buildURL(path: "/\(indexName)/_update/\(id)") + let body = try ByteBuffer(data: self.jsonEncoder.encode(script)) + var headers = HTTPHeaders() + headers.add(name: "content-type", value: "application/json") + return sendRequest(url: url, method: .POST, headers: headers, body: body) + } catch { + return self.eventLoop.makeFailedFuture(error) + } + } + public func deleteDocument(id: String, from indexName: String) -> EventLoopFuture { do { let url = try buildURL(path: "/\(indexName)/_doc/\(id)") diff --git a/Sources/ElasticsearchNIOClient/Models/BulkOperations/BulkUpdateScript.swift b/Sources/ElasticsearchNIOClient/Models/BulkOperations/BulkUpdateScript.swift new file mode 100644 index 0000000..0b4b793 --- /dev/null +++ b/Sources/ElasticsearchNIOClient/Models/BulkOperations/BulkUpdateScript.swift @@ -0,0 +1,9 @@ +import Foundation + +struct BulkUpdateScript: Codable { + let update: BulkOperationBody +} + +struct BulkUpdateScriptDocument: Encodable { + let script: Script +} diff --git a/Sources/ElasticsearchNIOClient/Models/ESBulkOperation.swift b/Sources/ElasticsearchNIOClient/Models/ESBulkOperation.swift index 570553f..da9c075 100644 --- a/Sources/ElasticsearchNIOClient/Models/ESBulkOperation.swift +++ b/Sources/ElasticsearchNIOClient/Models/ESBulkOperation.swift @@ -19,4 +19,5 @@ public enum BulkOperationType { case delete case index case update + case updateScript } diff --git a/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift b/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift index 86b708c..e326504 100644 --- a/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift +++ b/Tests/ElasticsearchNIOClientTests/ElasticsearchNIOClientTests.swift @@ -196,6 +196,138 @@ class ElasticSearchIntegrationTests: XCTestCase { XCTAssertTrue(results.hits.hits.contains(where: { $0.source.name == "Some 29 Apples" })) } + func testGetItem() throws { + let item = SomeItem(id: UUID(), name: "Some item") + _ = try client.createDocumentWithID(item, in: self.indexName).wait() + + Thread.sleep(forTimeInterval: 1.0) + + let retrievedItem: ESGetSingleDocumentResponse = try client.get(id: item.id.uuidString, from: self.indexName).wait() + XCTAssertEqual(retrievedItem.source.name, item.name) + } + + func testBulkUpdateWithScript() throws { + var items = [SomeItem]() + for index in 1...10 { + let name: String + if index % 2 == 0 { + name = "Some \(index) Apples" + } else { + name = "Some \(index) Bananas" + } + let item = SomeItem(id: UUID(), name: name, count: 0) + _ = try client.createDocumentWithID(item, in: self.indexName).wait() + items.append(item) + } + + // This is required for ES to settle and load the indexes to return the right results + Thread.sleep(forTimeInterval: 1.0) + + struct ScriptBody: Codable { + let inline: String + } + + let scriptBody = ScriptBody(inline: "ctx._source.count = ctx._source.count += 1") + + let bulkOperation = [ + ESBulkOperation(operationType: .updateScript, index: self.indexName, id: items[0].id.uuidString, document: scriptBody), + ] + + let response = try client.bulk(bulkOperation).wait() + XCTAssertEqual(response.items.count, 1) + XCTAssertNotNil(response.items.first?.update) + XCTAssertFalse(response.errors) + + let retrievedItem: ESGetSingleDocumentResponse = try client.get(id: items[0].id.uuidString, from: self.indexName).wait() + XCTAssertEqual(retrievedItem.source.count, 1) + } + + func testUpdateWithScript() throws { + let item = SomeItem(id: UUID(), name: "Some Item", count: 0) + _ = try client.createDocumentWithID(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 ScriptRequest: Codable { + let script: ScriptBody + } + + struct ScriptBody: Codable { + let inline: String + } + + let scriptBody = ScriptBody(inline: "ctx._source.count = ctx._source.count += 1") + let request = ScriptRequest(script: scriptBody) + + let response = try client.updateDocumentWithScript(request, id: item.id.uuidString, in: self.indexName).wait() + XCTAssertEqual(response.result, "updated") + + let retrievedItem: ESGetSingleDocumentResponse = try client.get(id: item.id.uuidString, from: self.indexName).wait() + XCTAssertEqual(retrievedItem.source.count, 1) + } + + func testUpdateWithNonExistentFieldScript() throws { + let item = SomeItem(id: UUID(), name: "Some Item") + _ = try client.createDocumentWithID(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 ScriptRequest: Codable { + let script: ScriptBody + } + + struct ScriptBody: Codable { + let inline: String + } + + let scriptBody = ScriptBody(inline: "if(ctx._source.containsKey('count')) { ctx._source.count += 1 } else { ctx._source.count = 1 }") + let request = ScriptRequest(script: scriptBody) + + let response = try client.updateDocumentWithScript(request, id: item.id.uuidString, in: self.indexName).wait() + XCTAssertEqual(response.result, "updated") + + let retrievedItem: ESGetSingleDocumentResponse = try client.get(id: item.id.uuidString, from: self.indexName).wait() + XCTAssertEqual(retrievedItem.source.count, 1) + } + + func testBulkUpdateWithNonExistentFieldScript() throws { + var items = [SomeItem]() + for index in 1...10 { + let name: String + if index % 2 == 0 { + name = "Some \(index) Apples" + } else { + name = "Some \(index) Bananas" + } + let item = SomeItem(id: UUID(), name: name) + _ = try client.createDocumentWithID(item, in: self.indexName).wait() + items.append(item) + } + + // This is required for ES to settle and load the indexes to return the right results + Thread.sleep(forTimeInterval: 1.0) + + struct ScriptBody: Codable { + let inline: String + } + + let scriptBody = ScriptBody(inline: "if(ctx._source.containsKey('count')) { ctx._source.count += 1 } else { ctx._source.count = 1 }") + + let bulkOperation = [ + ESBulkOperation(operationType: .updateScript, index: self.indexName, id: items[0].id.uuidString, document: scriptBody), + ] + + let response = try client.bulk(bulkOperation).wait() + XCTAssertEqual(response.items.count, 1) + XCTAssertNotNil(response.items.first?.update) + XCTAssertFalse(response.errors) + + let retrievedItem: ESGetSingleDocumentResponse = try client.get(id: items[0].id.uuidString, from: self.indexName).wait() + XCTAssertEqual(retrievedItem.source.count, 1) + } + // MARK: - Private private func setupItems() throws { for index in 1...10 { diff --git a/Tests/ElasticsearchNIOClientTests/SomeItem.swift b/Tests/ElasticsearchNIOClientTests/SomeItem.swift index c241757..8049b19 100644 --- a/Tests/ElasticsearchNIOClientTests/SomeItem.swift +++ b/Tests/ElasticsearchNIOClientTests/SomeItem.swift @@ -3,4 +3,11 @@ import Foundation struct SomeItem: Codable, Identifiable { let id: UUID let name: String + let count: Int? + + init(id: UUID, name: String, count: Int? = nil) { + self.id = id + self.name = name + self.count = count + } }