Enhance per bucket, and per document monitor notification message ctx.

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

@@ -43,6 +43,14 @@ data class AggregationResultBucket(
         return builder
     }
 
+    fun asTemplateArg(): Map<String, Any?> {
+        return mapOf(
+            PARENTS_BUCKET_PATH to parentBucketPath,
+            BUCKET_KEYS to bucketKeys,
+            BUCKET to bucket
+        )
+    }
+
     companion object {
         const val CONFIG_NAME = "agg_alert_content"
         const val PARENTS_BUCKET_PATH = "parent_bucket_path"

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

@@ -0,0 +1,46 @@
+/*
+ * 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
+) {
+    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 = alert.asTemplateArg().toMutableMap()
+
+        // Add the non-null custom context fields to the alert templateArgs.
+        customContextFields.forEach { (key, value) ->
+            value?.let { templateArgs[key] = it }
+        }
+        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.alert.id, "Template args id does not match")
+        assertEquals(templateArgs[Alert.ALERT_VERSION_FIELD], alertContext.alert.version, "Template args version does not match")
+        assertEquals(templateArgs[Alert.STATE_FIELD], alertContext.alert.state.toString(), "Template args state does not match")
+        assertEquals(templateArgs[Alert.ERROR_MESSAGE_FIELD], alertContext.alert.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.alert.endTime?.toEpochMilli(), "Template args end time does not")
+        assertEquals(templateArgs[Alert.START_TIME_FIELD], alertContext.alert.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.alert.severity, "Template args severity does not match")
+        assertEquals(templateArgs[Alert.CLUSTERS_FIELD], alertContext.alert.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")
+    }
+}