From 088e7a87d823a4776a80ebc5498a8edd47cc294f Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Sat, 23 Mar 2024 12:13:26 -0500
Subject: [PATCH] fix(searchContext): fix search flag immutability (#10117)

---
 .github/actions/ci-optimization/action.yml          |  1 +
 .../metadata/context/SearchContext.java             | 12 ++++++++----
 .../metadata/context/SearchContextTest.java         | 13 +++++++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/.github/actions/ci-optimization/action.yml b/.github/actions/ci-optimization/action.yml
index cad3a03dcb..c36b739701 100644
--- a/.github/actions/ci-optimization/action.yml
+++ b/.github/actions/ci-optimization/action.yml
@@ -70,6 +70,7 @@ runs:
             - "metadata-jobs/**"
             - "metadata-service/**"
             - "metadata-utils/**"
+            - "metadata-operation-context/**"
             - "datahub-graphql-core/**"
             - "smoke-test/**"
             - "docker/**"
diff --git a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java
index d4e3712309..ac9c903fb4 100644
--- a/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java
+++ b/metadata-operation-context/src/main/java/io/datahubproject/metadata/context/SearchContext.java
@@ -21,10 +21,14 @@ public class SearchContext implements ContextInterface {
   public static SearchContext withFlagDefaults(
       @Nonnull SearchContext searchContext,
       @Nonnull Function<SearchFlags, SearchFlags> flagDefaults) {
-    return searchContext.toBuilder()
-        // update search flags
-        .searchFlags(flagDefaults.apply(searchContext.getSearchFlags()))
-        .build();
+    try {
+      return searchContext.toBuilder()
+          // update search flags
+          .searchFlags(flagDefaults.apply(searchContext.getSearchFlags().copy()))
+          .build();
+    } catch (CloneNotSupportedException e) {
+      throw new RuntimeException(e);
+    }
   }
 
   @Nonnull private final IndexConvention indexConvention;
diff --git a/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java
index 26365c283f..4858bb3422 100644
--- a/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java
+++ b/metadata-operation-context/src/test/java/io/datahubproject/metadata/context/SearchContextTest.java
@@ -76,4 +76,17 @@ public class SearchContextTest {
             .getCacheKeyComponent(),
         "Expected differences in search flags to result in different caches");
   }
+
+  @Test
+  public void testImmutableSearchFlags() {
+    SearchContext initial =
+        SearchContext.builder().indexConvention(IndexConventionImpl.NO_PREFIX).build();
+    assertEquals(initial.getSearchFlags(), new SearchFlags().setSkipCache(false));
+
+    SearchContext mutated = initial.withFlagDefaults(flags -> flags.setSkipCache(true));
+    assertEquals(mutated.getSearchFlags(), new SearchFlags().setSkipCache(true));
+
+    // ensure original is not changed
+    assertEquals(initial.getSearchFlags(), new SearchFlags().setSkipCache(false));
+  }
 }