add Aggs types

rcarver · Jan 13, 2025 · 26d978b · 26d978b
1 parent 68813e2
commit 26d978b
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 16 deletions.
diff --git a/Sources/ElasticsearchQueryBuilder/Components.swift b/Sources/ElasticsearchQueryBuilder/Components.swift
@@ -60,22 +60,6 @@ extension esb {
         }
     }
 
-    /// Adds `query` block to the query syntax.
-    public struct Query<Component: DictComponent>: DictComponent {
-        var component: Component
-        public init(@QueryDictBuilder component: () -> Component) {
-            self.component = component()
-        }
-        public func makeDict() -> QueryDict {
-            let dict = self.component.makeDict()
-            if dict.isEmpty {
-                return [:]
-            } else {
-                return [ "query" : .dict(self.component.makeDict()) ]
-            }
-        }
-    }
-
     /// Adds a `key` with any type of value to the query syntax.
     public struct Key: DictComponent {
         var key: String
@@ -108,6 +92,73 @@ extension esb {
         }
     }
 
+    /// Adds a block to the syntax.
+    public struct Dict<Component: DictComponent>: DictComponent {
+        var key: String
+        var component: Component
+        public init(_ key: String, @QueryDictBuilder component: () -> Component) {
+            self.key = key
+            self.component = component()
+        }
+        public func makeDict() -> QueryDict {
+            let dict = self.component.makeDict()
+            if dict.isEmpty {
+                return [:]
+            } else {
+                return [ key : .dict(self.component.makeDict()) ]
+            }
+        }
+    }
+
+    /// Adds a `query` block to the syntax.
+    public struct Query<Component: DictComponent>: DictComponent {
+        var component: Component
+        public init(@QueryDictBuilder component: () -> Component) {
+            self.component = component()
+        }
+        public func makeDict() -> QueryDict {
+            let dict = self.component.makeDict()
+            if dict.isEmpty {
+                return [:]
+            } else {
+                return [ "query" : .dict(self.component.makeDict()) ]
+            }
+        }
+    }
+
+    /// Adds an `aggs` block to the syntax.
+    public struct Aggs<Component: DictComponent>: DictComponent {
+        var component: Component
+        public init(@QueryDictBuilder component: () -> Component) {
+            self.component = component()
+        }
+        public func makeDict() -> QueryDict {
+            let dict = self.component.makeDict()
+            if dict.isEmpty {
+                return [:]
+            } else {
+                return [ "aggs" : .dict(self.component.makeDict()) ]
+            }
+        }
+    }
+
+    /// Defines and named aggregate within `Aggs`
+    public struct Agg: DictComponent {
+        var name: String
+        var term: QueryDict
+        public init(_ name: String, field: String) {
+            self.name = name
+            self.term = [ "field" : .string(field) ]
+        }
+        public init(_ name: String, term: QueryDict) {
+            self.name = name
+            self.term = term
+        }
+        public func makeDict() -> QueryDict {
+            return [ self.name : [ "terms" : .dict(self.term) ] ]
+        }
+    }
+
     /// Adds `minimum_should_match` to the query syntax.
     public struct MinimumShouldMatch: DictComponent {
         var count: Int

diff --git a/Tests/ElasticsearchQueryBuilderTests/ComponentTests.swift b/Tests/ElasticsearchQueryBuilderTests/ComponentTests.swift
@@ -43,6 +43,35 @@ final class KeyTests: XCTestCase {
     }
 }
 
+final class DictTests: XCTestCase {
+    func testBuildDict() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.Dict("query") {
+                esb.Key("match_bool_prefix") {
+                    [
+                        "message": "quick brown f"
+                    ]
+                }
+            }
+        }
+        expectNoDifference(build().makeQuery(), [
+            "query": [
+                "match_bool_prefix": [
+                    "message": "quick brown f"
+                ]
+            ]
+        ])
+    }
+    func testBuildDictEmpty() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.Dict("query") {
+                esb.Key("match", .dict([:]))
+            }
+        }
+        expectNoDifference(build().makeQuery(), [:])
+    }
+}
+
 final class ComposableBuilderTests: XCTestCase {
     func testBuildDict() throws {
         @QueryDictBuilder func makeKey() -> some DictComponent {
@@ -153,6 +182,31 @@ final class QueryTests: XCTestCase {
     }
 }
 
+final class AggsTests: XCTestCase {
+    func testBuild() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.Aggs {
+                esb.Agg("name", field: "name")
+            }
+        }
+        expectNoDifference(build().makeQuery(), [
+            "aggs": [
+                "name": [
+                    "terms": [ "field": "name" ]
+                ]
+            ]
+        ])
+    }
+    func testBuildEmpty() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.Aggs {
+                esb.Key("name", .dict([:]))
+            }
+        }
+        expectNoDifference(build().makeQuery(), [:])
+    }
+}
+
 final class BoolTests: XCTestCase {
     func testBuild() throws {
         @ElasticsearchQueryBuilder func build(_ enabled: Bool) -> some esb.QueryDSL {