Removing some unwanted code

tandonks · tandonks · commit c63652cc4479 · 2025-11-29T11:25:54.000+05:30
Signed-off-by: Kshitij Tandon &lt;tandonks@amazon.com&gt;
diff --git a/src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupMetadataService.kt b/src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupMetadataService.kt
@@ -31,6 +31,7 @@ import org.opensearch.indexmanagement.IndexManagementPlugin
 import org.opensearch.indexmanagement.common.model.dimension.DateHistogram
 import org.opensearch.indexmanagement.opensearchapi.parseWithType
 import org.opensearch.indexmanagement.opensearchapi.suspendUntil
+import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor
 import org.opensearch.indexmanagement.rollup.model.ContinuousMetadata
 import org.opensearch.indexmanagement.rollup.model.Rollup
 import org.opensearch.indexmanagement.rollup.model.RollupMetadata
@@ -41,6 +42,7 @@ import org.opensearch.indexmanagement.util.NO_ID
 import org.opensearch.search.aggregations.bucket.composite.InternalComposite
 import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder
 import org.opensearch.search.builder.SearchSourceBuilder
+import org.opensearch.search.fetch.subphase.FetchSourceContext
 import org.opensearch.search.sort.SortOrder
 import org.opensearch.transport.RemoteTransportException
 import org.opensearch.transport.client.Client
@@ -242,7 +244,9 @@ class RollupMetadataService(
             val dateHistogram = rollup.dimensions.first() as DateHistogram
             val dateField = dateHistogram.sourceField
 
-            val bypassMarker = "_rollup_internal_bypass_${org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK}"
+            // For multi-tier rollup, we would be querying a document on a rollup index
+            // So we set this bypassMarker in fetchSource as a flag to help bypass the validation in RollupInterceptor
+            val bypassMarker = "_rollup_internal_bypass_${RollupInterceptor.BYPASS_SIZE_CHECK}"
 
             val searchRequest = SearchRequest(rollup.sourceIndex)
                 .source(
@@ -251,16 +255,11 @@ class RollupMetadataService(
                         .query(MatchAllQueryBuilder())
                         .sort("$dateField.date_histogram", SortOrder.ASC)
                         .trackTotalHits(false)
-                        .fetchSource(org.opensearch.search.fetch.subphase.FetchSourceContext(false, arrayOf(bypassMarker), emptyArray()))
+                        .fetchSource(FetchSourceContext(false, arrayOf(bypassMarker), emptyArray()))
                         .docValueField("$dateField.date_histogram", DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT),
                 )
                 .allowPartialSearchResults(false)
 
-            // Set bypass in ThreadContext for multi-node support - skips all rollup validations
-            // This is needed because we're querying a rollup index directly without aggregations
-            // Set bypass via preference parameter for multi-node support - allows size > 0 for fetching earliest timestamp
-//            searchRequest.preference("_rollup_bypass:${org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK}")
-
             val response: SearchResponse = client.suspendUntil { search(searchRequest, it) }
 
             if (response.hits.hits.isEmpty()) {
diff --git a/src/main/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptor.kt b/src/main/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptor.kt
@@ -66,55 +66,28 @@ class RollupInterceptor(
     @Volatile private var searchRawRollupIndices = RollupSettings.ROLLUP_SEARCH_SOURCE_INDICES.get(settings)
 
     companion object {
-        /**
-         * Thread-local bypass mechanism for internal operations that need to query rollup indices
-         * without triggering interceptor validations.
-         * Kept for backward compatibility with single-node clusters.
-         */
-        private val bypassInterceptor = ThreadLocal<Int>()
-
         /**
          * Bypass level that skips all rollup search validations and rewriting.
-         * Used when the system needs to query rollup indices directly using composite aggregation
+         * Used when the system needs to query rollup indices directly using composite aggregation.
+         * Passed via FetchSourceContext marker for multi-node support.
          */
         const val BYPASS_ROLLUP_SEARCH = 1
 
         /**
          * Bypass level that allows non-zero size in rollup searches.
          * Used for internal operations like continuous rollup initialization that need to fetch
          * actual documents from rollup indices (e.g., getEarliestTimestampFromRollupIndex).
+         * Passed via FetchSourceContext marker for multi-node support.
          */
         const val BYPASS_SIZE_CHECK = 2
-
-        /**
-         * Sets the bypass level for the current thread.
-         * Must be followed by clearBypass() to avoid leaking the bypass state.
-         */
-        fun setBypass(bypassLevel: Int) {
-            bypassInterceptor.set(bypassLevel)
-        }
-
-        /**
-         * Clears the bypass level for the current thread.
-         * Should always be called in a finally block after setBypass().
-         */
-        fun clearBypass() {
-            bypassInterceptor.remove()
-        }
-
-        /**
-         * Gets the current bypass level for the thread.
-         * Returns 0 if no bypass is set.
-         */
-        fun getBypassLevel(): Int = bypassInterceptor.get() ?: 0
     }
 
     /**
      * Reads the bypass value from the request's FetchSourceContext.
      * Returns the bypass level if the marker is present in includes array, null otherwise.
      * This enables bypass mechanism to work in multi-node clusters.
      */
-    private fun getBypassFromFetchSource(request: ShardSearchRequest): Int? {
+    internal fun getBypassFromFetchSource(request: ShardSearchRequest): Int? {
         val includes = request.source()?.fetchSource()?.includes()
 
         // Look for our bypass marker in the includes array and extract the bypass level
@@ -145,35 +118,25 @@ class RollupInterceptor(
     ): TransportRequestHandler<T> = object : TransportRequestHandler<T> {
         override fun messageReceived(request: T, channel: TransportChannel, task: Task) {
             if (searchEnabled && request is ShardSearchRequest) {
-                logger.info("Idhr aaya hu 3")
                 val index = request.shardId().indexName
                 val isRollupIndex = isRollupIndex(index, clusterService.state())
                 if (isRollupIndex) {
-                    // Check bypass from FetchSourceContext (multi-node) OR ThreadLocal (single-node/backward compat)
-                    val bypassFromFetchSource = getBypassFromFetchSource(request)
-                    val bypassFromThreadLocal = getBypassLevel()
+                    // Check bypass from FetchSourceContext for multi-node support
+                    val bypassLevel = getBypassFromFetchSource(request)
 
-                    // Use FetchSourceContext value if present, otherwise fall back to ThreadLocal
-                    val effectiveBypass = bypassFromFetchSource ?: bypassFromThreadLocal
-
-                    logger.debug(
-                        "RollupInterceptor bypass check - fetchSource: $bypassFromFetchSource, " +
-                            "threadLocal: $bypassFromThreadLocal, effective: $effectiveBypass",
-                    )
+                    logger.debug("RollupInterceptor bypass check - bypassLevel: $bypassLevel")
 
                     // BYPASS_ROLLUP_SEARCH: Skip all validations and query rewriting
-                    // Used for composite aggregation query which we do for retrieving the buckets
-                    // to write to target index during rollup of rolled up index
-                    if (effectiveBypass == BYPASS_ROLLUP_SEARCH) {
+                    // Used for composite aggregation queries when rolling up rollup indices
+                    if (bypassLevel == BYPASS_ROLLUP_SEARCH) {
                         actualHandler.messageReceived(request, channel, task)
                         return
                     }
 
                     // BYPASS_SIZE_CHECK: Allow non-zero size for internal operations that need to
-                    // fetch documents (e.g., fetching earliest timestamp document for rolled index
-                    // in case of continuous rollup job). Normal rollup searches must have size=0
-                    // since they should only return aggregations
-                    if (effectiveBypass != BYPASS_SIZE_CHECK && request.source().size() != 0) {
+                    // fetch documents (e.g., fetching earliest timestamp document for continuous rollup initialization)
+                    // Normal rollup searches must have size=0 since they should only return aggregations
+                    if (bypassLevel != BYPASS_SIZE_CHECK && request.source().size() != 0) {
                         throw IllegalArgumentException("Rollup search must have size explicitly set to 0, but found ${request.source().size()}")
                     }
 
diff --git a/src/test/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptorTests.kt b/src/test/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptorTests.kt
@@ -5,60 +5,139 @@
 
 package org.opensearch.indexmanagement.rollup.interceptor
 
+import org.junit.Before
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.`when`
+import org.opensearch.cluster.metadata.IndexNameExpressionResolver
+import org.opensearch.cluster.service.ClusterService
+import org.opensearch.common.settings.ClusterSettings
+import org.opensearch.common.settings.Settings
 import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.Companion.BYPASS_ROLLUP_SEARCH
 import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.Companion.BYPASS_SIZE_CHECK
-import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.Companion.clearBypass
-import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.Companion.getBypassLevel
-import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.Companion.setBypass
+import org.opensearch.indexmanagement.rollup.settings.RollupSettings
+import org.opensearch.search.builder.SearchSourceBuilder
+import org.opensearch.search.fetch.subphase.FetchSourceContext
+import org.opensearch.search.internal.ShardSearchRequest
 import org.opensearch.test.OpenSearchTestCase
 
 class RollupInterceptorTests : OpenSearchTestCase() {
 
-    fun `test setBypass and getBypassLevel for BYPASS_ROLLUP_SEARCH`() {
-        setBypass(BYPASS_ROLLUP_SEARCH)
-        assertEquals(BYPASS_ROLLUP_SEARCH, getBypassLevel())
-        clearBypass()
+    private lateinit var interceptor: RollupInterceptor
+
+    @Before
+    fun setup() {
+        interceptor = createInterceptor()
+    }
+
+    fun `test getBypassFromFetchSource returns null when no source`() {
+        val request = mock(ShardSearchRequest::class.java)
+        `when`(request.source()).thenReturn(null)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertNull(bypassLevel)
     }
 
-    fun `test setBypass and getBypassLevel for BYPASS_SIZE_CHECK`() {
-        setBypass(BYPASS_SIZE_CHECK)
-        assertEquals(BYPASS_SIZE_CHECK, getBypassLevel())
-        clearBypass()
+    fun `test getBypassFromFetchSource returns null when no FetchSourceContext`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertNull(bypassLevel)
     }
 
-    fun `test getBypassLevel returns 0 when not set`() {
-        clearBypass()
-        assertEquals(0, getBypassLevel())
+    fun `test getBypassFromFetchSource returns null when no includes array`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(FetchSourceContext(true))
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertNull(bypassLevel)
     }
 
-    fun `test clearBypass resets bypass level`() {
-        setBypass(BYPASS_ROLLUP_SEARCH)
-        assertEquals(BYPASS_ROLLUP_SEARCH, getBypassLevel())
-        clearBypass()
-        assertEquals(0, getBypassLevel())
+    fun `test getBypassFromFetchSource returns null when no bypass marker present`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(FetchSourceContext(false, arrayOf("field1", "field2"), emptyArray()))
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertNull(bypassLevel)
     }
 
-    fun `test bypass is thread local`() {
-        setBypass(BYPASS_ROLLUP_SEARCH)
-        assertEquals(BYPASS_ROLLUP_SEARCH, getBypassLevel())
+    fun `test getBypassFromFetchSource extracts BYPASS_ROLLUP_SEARCH correctly`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(FetchSourceContext(false, arrayOf("_rollup_internal_bypass_$BYPASS_ROLLUP_SEARCH"), emptyArray()))
 
-        val thread = Thread {
-            assertEquals(0, getBypassLevel())
-            setBypass(BYPASS_SIZE_CHECK)
-            assertEquals(BYPASS_SIZE_CHECK, getBypassLevel())
-        }
-        thread.start()
-        thread.join()
+        `when`(request.source()).thenReturn(source)
 
-        assertEquals(BYPASS_ROLLUP_SEARCH, getBypassLevel())
-        clearBypass()
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertEquals(BYPASS_ROLLUP_SEARCH, bypassLevel)
     }
 
-    fun `test multiple setBypass calls overwrite previous value`() {
-        setBypass(BYPASS_ROLLUP_SEARCH)
-        assertEquals(BYPASS_ROLLUP_SEARCH, getBypassLevel())
-        setBypass(BYPASS_SIZE_CHECK)
-        assertEquals(BYPASS_SIZE_CHECK, getBypassLevel())
-        clearBypass()
+    fun `test getBypassFromFetchSource extracts BYPASS_SIZE_CHECK correctly`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(FetchSourceContext(false, arrayOf("_rollup_internal_bypass_$BYPASS_SIZE_CHECK"), emptyArray()))
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertEquals(BYPASS_SIZE_CHECK, bypassLevel)
+    }
+
+    fun `test getBypassFromFetchSource finds marker among multiple includes`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(
+            FetchSourceContext(false, arrayOf("field1", "_rollup_internal_bypass_$BYPASS_ROLLUP_SEARCH", "field2"), emptyArray()),
+        )
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertEquals(BYPASS_ROLLUP_SEARCH, bypassLevel)
+    }
+
+    fun `test getBypassFromFetchSource returns null for invalid bypass marker`() {
+        val request = mock(ShardSearchRequest::class.java)
+        val source = SearchSourceBuilder()
+        source.fetchSource(FetchSourceContext(false, arrayOf("_rollup_internal_bypass_invalid"), emptyArray()))
+
+        `when`(request.source()).thenReturn(source)
+
+        val bypassLevel = interceptor.getBypassFromFetchSource(request)
+
+        assertNull(bypassLevel)
+    }
+
+    // Helper method to create interceptor instance
+    private fun createInterceptor(): RollupInterceptor {
+        val clusterService = mock(ClusterService::class.java)
+        val clusterSettings = ClusterSettings(
+            Settings.EMPTY,
+            setOf(
+                RollupSettings.ROLLUP_SEARCH_ENABLED,
+                RollupSettings.ROLLUP_SEARCH_ALL_JOBS,
+                RollupSettings.ROLLUP_SEARCH_SOURCE_INDICES,
+            ),
+        )
+        `when`(clusterService.clusterSettings).thenReturn(clusterSettings)
+
+        val settings = Settings.EMPTY
+        val resolver = mock(IndexNameExpressionResolver::class.java)
+        return RollupInterceptor(clusterService, settings, resolver)
     }
 }
