Add TenantContext extension (#952)

* Add TenantContext extension

Signed-off-by: Chase Engelbrecht <engechas@dev-dsk-engechas-2a-28f138b0.us-west-2.amazon.com>

* Generic comment

Signed-off-by: Chase Engelbrecht <engechas@dev-dsk-engechas-2a-28f138b0.us-west-2.amazon.com>

* Fix tests

Signed-off-by: Chase Engelbrecht <engechas@dev-dsk-engechas-2a-28f138b0.us-west-2.amazon.com>

---------

Signed-off-by: Chase Engelbrecht <engechas@dev-dsk-engechas-2a-28f138b0.us-west-2.amazon.com>
Co-authored-by: Chase Engelbrecht <engechas@dev-dsk-engechas-2a-28f138b0.us-west-2.amazon.com>

Author: engechas

src/main/kotlin/org/opensearch/commons/utils/TenantContext.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.commons.utils
+
+import kotlin.coroutines.AbstractCoroutineContextElement
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.coroutineContext
+
+/**
+ * Coroutine context element that carries the tenant ID for the current request.
+ * Set when launching a coroutine and read when building SDK requests.
+ */
+data class TenantContext(val tenantId: String?) : AbstractCoroutineContextElement(TenantContext) {
+    companion object Key : CoroutineContext.Key<TenantContext>
+}
+
+/**
+ * Retrieves the tenant ID from the current coroutine context.
+ * Returns null if no [TenantContext] is present or if the tenant ID is not set.
+ */
+suspend fun currentTenantId(): String? = coroutineContext[TenantContext]?.tenantId

src/test/kotlin/org/opensearch/commons/utils/TenantContextTests.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.commons.utils
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.coroutines.startCoroutine
+
+internal class TenantContextTests {
+
+    private fun runSuspend(context: kotlin.coroutines.CoroutineContext = EmptyCoroutineContext, block: suspend () -> Unit) {
+        var result: Result<Unit>? = null
+        block.startCoroutine(object : kotlin.coroutines.Continuation<Unit> {
+            override val context = context
+            override fun resumeWith(r: Result<Unit>) { result = r }
+        })
+        result!!.getOrThrow()
+    }
+
+    @Test
+    fun `test currentTenantId returns tenant id when set`() {
+        runSuspend(TenantContext("test-tenant")) {
+            assertEquals("test-tenant", currentTenantId())
+        }
+    }
+
+    @Test
+    fun `test currentTenantId returns null for single tenant deployment`() {
+        runSuspend {
+            assertNull(currentTenantId())
+        }
+    }
+
+    @Test
+    fun `test currentTenantId returns null when tenant id header is absent`() {
+        runSuspend(TenantContext(null)) {
+            assertNull(currentTenantId())
+        }
+    }
+}