Merge pull request #4 from brokenhandsio/get-single-and-script

Add support for Scripting and Getting Single Document

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!
 

Sources/ElasticsearchNIOClient/ElasticsearchClient+Requests.swift

@@ -3,6 +3,15 @@ import NIO
 import NIOHTTP1
 
 extension ElasticsearchClient {
+    public func get<Document: Decodable>(id: String, from indexName: String) -> EventLoopFuture<ESGetSingleDocumentResponse<Document>> {
+        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<Document: Encodable>(_ operations: [ESBulkOperation<Document>]) -> EventLoopFuture<ESBulkResponse> {
         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: Encodable>(_ script: Script, id: String, in indexName: String) -> EventLoopFuture<ESUpdateDocumentResponse> {
+        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<ESDeleteDocumentResponse> {
         do {
             let url = try buildURL(path: "/\(indexName)/_doc/\(id)")

Sources/ElasticsearchNIOClient/Models/BulkOperations/BulkUpdateScript.swift

@@ -0,0 +1,9 @@
+import Foundation
+
+struct BulkUpdateScript: Codable {
+    let update: BulkOperationBody
+}
+
+struct BulkUpdateScriptDocument<Script: Encodable>: Encodable {
+    let script: Script
+}

Sources/ElasticsearchNIOClient/Models/ESBulkOperation.swift

@@ -19,4 +19,5 @@ public enum BulkOperationType {
     case delete
     case index
     case update
+    case updateScript
 }

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<SomeItem> = 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<SomeItem> = 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<SomeItem> = 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<SomeItem> = 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<SomeItem> = 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 {

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
+    }
 }