Added support for printing sample documents for each bucket in bucket…

… level monitor notification messages. Added support for printing query/rule info in doc level monitor notification messages.

Signed-off-by: AWSHurneyt <hurneyt@amazon.com>

src/main/kotlin/org/opensearch/commons/alerting/model/AggregationResultBucket.kt

@@ -1,6 +1,5 @@
 package org.opensearch.commons.alerting.model
 
-import org.apache.logging.log4j.LogManager
 import org.opensearch.core.common.ParsingException
 import org.opensearch.core.common.io.stream.StreamInput
 import org.opensearch.core.common.io.stream.StreamOutput
@@ -14,8 +13,6 @@ import org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken
 import java.io.IOException
 import java.util.Locale
 
-private val logger = LogManager.getLogger(AggregationResultBucket::class.java)
-
 data class AggregationResultBucket(
     val parentBucketPath: String?,
     val bucketKeys: List<String>,
@@ -38,9 +35,6 @@ data class AggregationResultBucket(
     }
 
     fun innerXContent(builder: XContentBuilder): XContentBuilder {
-        logger.info("hurneyt innerXContent::PARENTS_BUCKET_PATH = {}", parentBucketPath)
-        logger.info("hurneyt innerXContent::BUCKET_KEYS = {}", bucketKeys.toTypedArray())
-        logger.info("hurneyt innerXContent::BUCKET = {}", bucket)
         builder.startObject(CONFIG_NAME)
             .field(PARENTS_BUCKET_PATH, parentBucketPath)
             .field(BUCKET_KEYS, bucketKeys.toTypedArray())
@@ -50,17 +44,11 @@ data class AggregationResultBucket(
     }
 
     fun asTemplateArg(): Map<String, Any?> {
-        logger.info("hurneyt asTemplateArg START")
-        logger.info("hurneyt asTemplateArg::PARENTS_BUCKET_PATH = {}", parentBucketPath)
-        logger.info("hurneyt asTemplateArg::BUCKET_KEYS = {}", bucketKeys.toTypedArray())
-        logger.info("hurneyt asTemplateArg::BUCKET = {}", bucket)
-        val output = mapOf(
+        return mapOf(
             PARENTS_BUCKET_PATH to parentBucketPath,
             BUCKET_KEYS to bucketKeys,
             BUCKET to bucket
         )
-        logger.info("hurneyt asTemplateArg END")
-        return output
     }
 
     companion object {

src/main/kotlin/org/opensearch/commons/alerting/model/Alert.kt

@@ -17,7 +17,7 @@ import org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken
 import java.io.IOException
 import java.time.Instant
 
-data class Alert(
+open class Alert(
     val id: String = NO_ID,
     val version: Long = NO_VERSION,
     val schemaVersion: Int = NO_SCHEMA_VERSION,
@@ -591,7 +591,7 @@ data class Alert(
         return builder
     }
 
-    fun asTemplateArg(): Map<String, Any?> {
+    open fun asTemplateArg(): Map<String, Any?> {
         return mapOf(
             ACKNOWLEDGED_TIME_FIELD to acknowledgedTime?.toEpochMilli(),
             ALERT_ID_FIELD to id,
@@ -614,4 +614,101 @@ data class Alert(
             CLUSTERS_FIELD to clusters?.joinToString(",")
         )
     }
+
+    /**
+     * Creates a copy of an [Alert] with optionally modified properties
+     */
+    fun copy(
+        id: String = this.id,
+        version: Long = this.version,
+        schemaVersion: Int = this.schemaVersion,
+        monitorId: String = this.monitorId,
+        workflowId: String = this.workflowId,
+        workflowName: String = this.workflowName,
+        monitorName: String = this.monitorName,
+        monitorVersion: Long = this.monitorVersion,
+        monitorUser: User? = this.monitorUser,
+        triggerId: String = this.triggerId,
+        triggerName: String = this.triggerName,
+        findingIds: List<String> = this.findingIds,
+        relatedDocIds: List<String> = this.relatedDocIds,
+        state: State = this.state,
+        startTime: Instant = this.startTime,
+        endTime: Instant? = this.endTime,
+        lastNotificationTime: Instant? = this.lastNotificationTime,
+        acknowledgedTime: Instant? = this.acknowledgedTime,
+        errorMessage: String? = this.errorMessage,
+        errorHistory: List<AlertError> = this.errorHistory,
+        severity: String = this.severity,
+        actionExecutionResults: List<ActionExecutionResult> = this.actionExecutionResults,
+        aggregationResultBucket: AggregationResultBucket? = this.aggregationResultBucket,
+        executionId: String? = this.executionId,
+        associatedAlertIds: List<String> = this.associatedAlertIds,
+        clusters: List<String>? = this.clusters
+    ): Alert {
+        return Alert(
+            id = id,
+            version = version,
+            schemaVersion = schemaVersion,
+            monitorId = monitorId,
+            workflowId = workflowId,
+            workflowName = workflowName,
+            monitorName = monitorName,
+            monitorVersion = monitorVersion,
+            monitorUser = monitorUser,
+            triggerId = triggerId,
+            triggerName = triggerName,
+            findingIds = findingIds,
+            relatedDocIds = relatedDocIds,
+            state = state,
+            startTime = startTime,
+            endTime = endTime,
+            lastNotificationTime = lastNotificationTime,
+            acknowledgedTime = acknowledgedTime,
+            errorMessage = errorMessage,
+            errorHistory = errorHistory,
+            severity = severity,
+            actionExecutionResults = actionExecutionResults,
+            aggregationResultBucket = aggregationResultBucket,
+            executionId = executionId,
+            associatedAlertIds = associatedAlertIds,
+            clusters = clusters
+        )
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as Alert
+
+        if (id != other.id) return false
+        if (version != other.version) return false
+        if (schemaVersion != other.schemaVersion) return false
+        if (monitorId != other.monitorId) return false
+        if (workflowId != other.workflowId) return false
+        if (workflowName != other.workflowName) return false
+        if (monitorName != other.monitorName) return false
+        if (monitorVersion != other.monitorVersion) return false
+        if (monitorUser != other.monitorUser) return false
+        if (triggerId != other.triggerId) return false
+        if (triggerName != other.triggerName) return false
+        if (findingIds != other.findingIds) return false
+        if (relatedDocIds != other.relatedDocIds) return false
+        if (state != other.state) return false
+        if (startTime != other.startTime) return false
+        if (endTime != other.endTime) return false
+        if (lastNotificationTime != other.lastNotificationTime) return false
+        if (acknowledgedTime != other.acknowledgedTime) return false
+        if (errorMessage != other.errorMessage) return false
+        if (errorHistory != other.errorHistory) return false
+        if (severity != other.severity) return false
+        if (actionExecutionResults != other.actionExecutionResults) return false
+        if (aggregationResultBucket != other.aggregationResultBucket) return false
+        if (executionId != other.executionId) return false
+        if (associatedAlertIds != other.associatedAlertIds) return false
+        if (clusters != other.clusters) return false
+
+        return true
+    }
 }

src/main/kotlin/org/opensearch/commons/alerting/model/AlertContext.kt

@@ -0,0 +1,75 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.commons.alerting.model
+
+/**
+ * This model is a wrapper for [Alert] that should only be used to create a more
+ * informative alert object to enrich mustache template notification messages.
+ */
+data class AlertContext(
+    val alert: Alert,
+    val associatedQueries: List<DocLevelQuery>? = null,
+    val sampleDocs: List<Map<String, Any?>>? = null
+) : Alert(
+    id = alert.id,
+    version = alert.version,
+    schemaVersion = alert.schemaVersion,
+    monitorId = alert.monitorId,
+    monitorName = alert.monitorName,
+    monitorVersion = alert.monitorVersion,
+    monitorUser = alert.monitorUser,
+    triggerId = alert.triggerId,
+    triggerName = alert.triggerName,
+    state = alert.state,
+    startTime = alert.startTime,
+    endTime = alert.endTime,
+    lastNotificationTime = alert.lastNotificationTime,
+    acknowledgedTime = alert.acknowledgedTime,
+    errorMessage = alert.errorMessage,
+    errorHistory = alert.errorHistory,
+    severity = alert.severity,
+    actionExecutionResults = alert.actionExecutionResults,
+    aggregationResultBucket = alert.aggregationResultBucket,
+    findingIds = alert.findingIds,
+    relatedDocIds = alert.relatedDocIds,
+    executionId = alert.executionId,
+    workflowId = alert.workflowId,
+    workflowName = alert.workflowName,
+    associatedAlertIds = alert.associatedAlertIds,
+    clusters = alert.clusters
+) {
+
+    override fun asTemplateArg(): Map<String, Any?> {
+        val queriesContext = associatedQueries?.map {
+            mapOf(
+                DocLevelQuery.QUERY_ID_FIELD to it.id,
+                DocLevelQuery.NAME_FIELD to it.name,
+                DocLevelQuery.TAGS_FIELD to it.tags
+            )
+        }
+
+        // Compile the custom context fields.
+        val customContextFields = mapOf(
+            ASSOCIATED_QUERIES_FIELD to queriesContext,
+            SAMPLE_DOCS_FIELD to sampleDocs
+        )
+
+        // Get the alert template args
+        val templateArgs = super.asTemplateArg().toMutableMap()
+
+        // Add the non-null custom context fields to the alert templateArgs.
+        customContextFields.forEach { (key, value) ->
+            if (value !== null) templateArgs[key] = value
+        }
+        return templateArgs
+    }
+
+    companion object {
+        const val ASSOCIATED_QUERIES_FIELD = "associated_queries"
+        const val SAMPLE_DOCS_FIELD = "sample_documents"
+
+    }
+}

src/test/kotlin/org/opensearch/commons/alerting/TestHelpers.kt

@@ -20,6 +20,7 @@ import org.opensearch.commons.alerting.aggregation.bucketselectorext.BucketSelec
 import org.opensearch.commons.alerting.model.ActionExecutionResult
 import org.opensearch.commons.alerting.model.AggregationResultBucket
 import org.opensearch.commons.alerting.model.Alert
+import org.opensearch.commons.alerting.model.AlertContext
 import org.opensearch.commons.alerting.model.BucketLevelTrigger
 import org.opensearch.commons.alerting.model.ChainedAlertTrigger
 import org.opensearch.commons.alerting.model.ChainedMonitorFindings
@@ -601,3 +602,22 @@ fun randomFinding(
         timestamp = timestamp
     )
 }
+
+fun randomAlertContext(
+    alert: Alert = randomAlert(),
+    associatedQueries: List<DocLevelQuery>? = (-1..2).random().takeIf { it != -1 }?.let {
+        (0..it).map { randomDocLevelQuery() }
+    },
+    sampleDocs: List<Map<String, Any?>>? = (-1..2).random().takeIf { it != -1 }?.let {
+        (0..it).map {
+            // Using 'randomFinding' to mimic documents in an index.
+            randomFinding().asTemplateArg()
+        }
+    }
+): AlertContext {
+    return AlertContext(
+        alert = alert,
+        associatedQueries = associatedQueries,
+        sampleDocs = sampleDocs
+    )
+}

src/test/kotlin/org/opensearch/commons/alerting/model/AlertContextTests.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.commons.alerting.model
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.opensearch.commons.alerting.randomAlertContext
+
+class AlertContextTests {
+    private var alertContext: AlertContext = randomAlertContext()
+
+    @BeforeEach
+    fun generateRandomData() {
+        alertContext = randomAlertContext()
+    }
+
+    @Test
+    fun `test AlertContext asTemplateArg`() {
+        val templateArgs = alertContext.asTemplateArg()
+
+        assertEquals(templateArgs[Alert.ALERT_ID_FIELD], alertContext.id, "Template args id does not match")
+        assertEquals(templateArgs[Alert.ALERT_VERSION_FIELD], alertContext.version, "Template args version does not match")
+        assertEquals(templateArgs[Alert.STATE_FIELD], alertContext.state.toString(), "Template args state does not match")
+        assertEquals(templateArgs[Alert.ERROR_MESSAGE_FIELD], alertContext.errorMessage, "Template args error message does not match")
+        assertEquals(templateArgs[Alert.ACKNOWLEDGED_TIME_FIELD], null, "Template args acknowledged time does not match")
+        assertEquals(templateArgs[Alert.END_TIME_FIELD], alertContext.endTime?.toEpochMilli(), "Template args end time does not")
+        assertEquals(templateArgs[Alert.START_TIME_FIELD], alertContext.startTime.toEpochMilli(), "Template args start time does not")
+        assertEquals(templateArgs[Alert.LAST_NOTIFICATION_TIME_FIELD], null, "Template args last notification time does not match")
+        assertEquals(templateArgs[Alert.SEVERITY_FIELD], alertContext.severity, "Template args severity does not match")
+        assertEquals(templateArgs[Alert.CLUSTERS_FIELD], alertContext.clusters?.joinToString(","), "Template args clusters does not match")
+        val formattedQueries = alertContext.associatedQueries?.map {
+            mapOf(
+                DocLevelQuery.QUERY_ID_FIELD to it.id,
+                DocLevelQuery.NAME_FIELD to it.name,
+                DocLevelQuery.TAGS_FIELD to it.tags
+            )
+        }
+        assertEquals(templateArgs[AlertContext.ASSOCIATED_QUERIES_FIELD], formattedQueries, "Template associated queries do not match")
+        assertEquals(templateArgs[AlertContext.SAMPLE_DOCS_FIELD], alertContext.sampleDocs, "Template args sample docs do not match")
+    }
+}