add 'fields' parameter in doc level query object. (opensearch-project…

…#546) (opensearch-project#550) (cherry picked from commit 2ef47f9) Signed-off-by: Surya Sashank Nistala <snistala@amazon.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: AWSHurneyt <hurneyt@amazon.com>
AWSHurneyt · Apr 12, 2024 · 8e21bac · 8e21bac
1 parent 17fb0e7
commit 8e21bac
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 3 deletions.
diff --git a/src/main/kotlin/org/opensearch/commons/alerting/model/DocLevelQuery.kt b/src/main/kotlin/org/opensearch/commons/alerting/model/DocLevelQuery.kt
@@ -14,6 +14,7 @@ import java.util.UUID
 data class DocLevelQuery(
     val id: String = UUID.randomUUID().toString(),
     val name: String,
+    val fields: List<String>,
     val query: String,
     val tags: List<String> = mutableListOf()
 ) : BaseModel {
@@ -30,6 +31,7 @@ data class DocLevelQuery(
     constructor(sin: StreamInput) : this(
         sin.readString(), // id
         sin.readString(), // name
+        sin.readStringList(), // fields
         sin.readString(), // query
         sin.readStringList() // tags
     )
@@ -38,6 +40,7 @@ data class DocLevelQuery(
         return mapOf(
             QUERY_ID_FIELD to id,
             NAME_FIELD to name,
+            FIELDS_FIELD to fields,
             QUERY_FIELD to query,
             TAGS_FIELD to tags
         )
@@ -47,6 +50,7 @@ data class DocLevelQuery(
     override fun writeTo(out: StreamOutput) {
         out.writeString(id)
         out.writeString(name)
+        out.writeStringCollection(fields)
         out.writeString(query)
         out.writeStringCollection(tags)
     }
@@ -64,6 +68,7 @@ data class DocLevelQuery(
     companion object {
         const val QUERY_ID_FIELD = "id"
         const val NAME_FIELD = "name"
+        const val FIELDS_FIELD = "fields"
         const val QUERY_FIELD = "query"
         const val TAGS_FIELD = "tags"
         const val NO_ID = ""
@@ -76,6 +81,7 @@ data class DocLevelQuery(
             lateinit var query: String
             lateinit var name: String
             val tags: MutableList<String> = mutableListOf()
+            val fields: MutableList<String> = mutableListOf()
 
             XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, xcp.currentToken(), xcp)
             while (xcp.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -101,12 +107,24 @@ data class DocLevelQuery(
                             tags.add(tag)
                         }
                     }
+                    FIELDS_FIELD -> {
+                        XContentParserUtils.ensureExpectedToken(
+                            XContentParser.Token.START_ARRAY,
+                            xcp.currentToken(),
+                            xcp
+                        )
+                        while (xcp.nextToken() != XContentParser.Token.END_ARRAY) {
+                            val field = xcp.text()
+                            fields.add(field)
+                        }
+                    }
                 }
             }
 
             return DocLevelQuery(
                 id = id,
                 name = name,
+                fields = fields,
                 query = query,
                 tags = tags
             )

diff --git a/src/test/kotlin/org/opensearch/commons/alerting/TestHelpers.kt b/src/test/kotlin/org/opensearch/commons/alerting/TestHelpers.kt
@@ -383,7 +383,7 @@ fun randomDocLevelQuery(
     name: String = "${RandomNumbers.randomIntBetween(Random(), 0, 5)}",
     tags: List<String> = mutableListOf(0..RandomNumbers.randomIntBetween(Random(), 0, 10)).map { RandomStrings.randomAsciiLettersOfLength(Random(), 10) }
 ): DocLevelQuery {
-    return DocLevelQuery(id = id, query = query, name = name, tags = tags)
+    return DocLevelQuery(id = id, query = query, name = name, tags = tags, fields = listOf("*"))
 }
 
 fun randomDocLevelMonitorInput(

diff --git a/src/test/kotlin/org/opensearch/commons/alerting/action/GetFindingsResponseTests.kt b/src/test/kotlin/org/opensearch/commons/alerting/action/GetFindingsResponseTests.kt
@@ -24,7 +24,7 @@ internal class GetFindingsResponseTests {
             "monitor_id1",
             "monitor_name1",
             "test_index1",
-            listOf(DocLevelQuery("1", "myQuery", "fieldA:valABC", listOf())),
+            listOf(DocLevelQuery("1", "myQuery", listOf(), "fieldA:valABC", listOf())),
             Instant.now()
         )
         val findingDocument1 = FindingDocument("test_index1", "doc1", true, "document 1 payload")
@@ -43,7 +43,7 @@ internal class GetFindingsResponseTests {
             "monitor_id2",
             "monitor_name2",
             "test_index2",
-            listOf(DocLevelQuery("1", "myQuery", "fieldA:valABC", listOf())),
+            listOf(DocLevelQuery("1", "myQuery", listOf(), "fieldA:valABC", listOf())),
             Instant.now()
         )
         val findingDocument21 = FindingDocument("test_index2", "doc21", true, "document 21 payload")

diff --git a/src/test/kotlin/org/opensearch/commons/alerting/model/WriteableTests.kt b/src/test/kotlin/org/opensearch/commons/alerting/model/WriteableTests.kt
@@ -10,6 +10,7 @@ import org.opensearch.commons.alerting.randomAction
 import org.opensearch.commons.alerting.randomActionExecutionPolicy
 import org.opensearch.commons.alerting.randomBucketLevelTrigger
 import org.opensearch.commons.alerting.randomChainedAlertTrigger
+import org.opensearch.commons.alerting.randomDocLevelQuery
 import org.opensearch.commons.alerting.randomDocumentLevelTrigger
 import org.opensearch.commons.alerting.randomQueryLevelMonitor
 import org.opensearch.commons.alerting.randomQueryLevelTrigger
@@ -112,6 +113,16 @@ class WriteableTests {
         Assertions.assertEquals(trigger, newTrigger, "Round tripping DocumentLevelTrigger doesn't work")
     }
 
+    @Test
+    fun `test doc-level query as stream`() {
+        val dlq = randomDocLevelQuery()
+        val out = BytesStreamOutput()
+        dlq.writeTo(out)
+        val sin = StreamInput.wrap(out.bytes().toBytesRef().bytes)
+        val newDlq = DocLevelQuery.readFrom(sin)
+        Assertions.assertEquals(dlq, newDlq, "Round tripping DocLevelQuery doesn't work")
+    }
+
     @Test
     fun `test chained alert trigger as stream`() {
         val trigger = randomChainedAlertTrigger()