Clean up CallContext Part 2

persistence/eclipselink/src/test/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreManagerTest.java

@@ -37,6 +37,7 @@
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
 import org.apache.polaris.core.persistence.BasePolarisMetaStoreManagerTest;
 import org.apache.polaris.core.persistence.PolarisTestMetaStoreManager;
@@ -90,6 +91,7 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
     return new PolarisTestMetaStoreManager(
         new TransactionalMetaStoreManagerImpl(),
         new PolarisCallContext(
+            Mockito.mock(RealmContext.class),
             session,
             diagServices,
             new PolarisConfigurationStore() {},

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java

@@ -157,7 +157,7 @@ public Map<String, BaseResult> purgeRealms(Iterable<String> realms) {
       PolarisMetaStoreManager metaStoreManager = getOrCreateMetaStoreManager(() -> realm);
       BasePersistence session = getOrCreateSessionSupplier(() -> realm).get();
 
-      PolarisCallContext callContext = new PolarisCallContext(session, diagServices);
+      PolarisCallContext callContext = new PolarisCallContext(() -> realm, session, diagServices);
       BaseResult result = metaStoreManager.purge(callContext);
       results.put(realm, result);
 

persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java

@@ -28,6 +28,7 @@
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.persistence.AtomicOperationMetaStoreManager;
 import org.apache.polaris.core.persistence.BasePolarisMetaStoreManagerTest;
 import org.apache.polaris.core.persistence.PolarisTestMetaStoreManager;
@@ -61,6 +62,7 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
     return new PolarisTestMetaStoreManager(
         new AtomicOperationMetaStoreManager(),
         new PolarisCallContext(
+            Mockito.mock(RealmContext.class),
             basePersistence,
             diagServices,
             new PolarisConfigurationStore() {},

polaris-core/src/main/java/org/apache/polaris/core/PolarisCallContext.java

@@ -58,18 +58,6 @@ public PolarisCallContext(
     this.clock = clock;
   }
 
-  @Deprecated
-  public PolarisCallContext(
-      @Nonnull BasePersistence metaStore,
-      @Nonnull PolarisDiagnostics diagServices,
-      @Nonnull PolarisConfigurationStore configurationStore,
-      @Nonnull Clock clock) {
-    this.metaStore = metaStore;
-    this.diagServices = diagServices;
-    this.configurationStore = configurationStore;
-    this.clock = clock;
-  }
-
   public PolarisCallContext(
       @Nonnull RealmContext realmContext,
       @Nonnull BasePersistence metaStore,
@@ -81,15 +69,6 @@ public PolarisCallContext(
     this.clock = Clock.system(ZoneId.systemDefault());
   }
 
-  @Deprecated
-  public PolarisCallContext(
-      @Nonnull BasePersistence metaStore, @Nonnull PolarisDiagnostics diagServices) {
-    this.metaStore = metaStore;
-    this.diagServices = diagServices;
-    this.configurationStore = new PolarisConfigurationStore() {};
-    this.clock = Clock.system(ZoneId.systemDefault());
-  }
-
   public BasePersistence getMetaStore() {
     return metaStore;
   }

polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java

@@ -49,22 +49,6 @@ static void unsetCurrentContext() {
     CURRENT_CONTEXT.remove();
   }
 
-  // only tests are using this method now, we can get rid of them easily in a followup
-  static CallContext of(
-      final RealmContext realmContext, final PolarisCallContext polarisCallContext) {
-    return new CallContext() {
-      @Override
-      public RealmContext getRealmContext() {
-        return realmContext;
-      }
-
-      @Override
-      public PolarisCallContext getPolarisCallContext() {
-        return polarisCallContext;
-      }
-    };
-  }
-
   /** Copy the {@link CallContext}. */
   static CallContext copyOf(CallContext base) {
     String realmId = base.getRealmContext().getRealmIdentifier();

polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java

@@ -133,7 +133,7 @@ public Map<String, BaseResult> purgeRealms(Iterable<String> realms) {
       PolarisMetaStoreManager metaStoreManager = getOrCreateMetaStoreManager(() -> realm);
       TransactionalPersistence session = getOrCreateSessionSupplier(() -> realm).get();
 
-      PolarisCallContext callContext = new PolarisCallContext(session, diagServices);
+      PolarisCallContext callContext = new PolarisCallContext(() -> realm, session, diagServices);
       BaseResult result = metaStoreManager.purge(callContext);
       results.put(realm, result);
 

polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapAtomicOperationMetaStoreManagerTest.java

@@ -25,6 +25,7 @@
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore;
 import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl;
 import org.mockito.Mockito;
@@ -37,6 +38,7 @@ public PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
     TreeMapMetaStore store = new TreeMapMetaStore(diagServices);
     PolarisCallContext callCtx =
         new PolarisCallContext(
+            Mockito.mock(RealmContext.class),
             new TreeMapTransactionalPersistenceImpl(store, Mockito.mock(), RANDOM_SECRETS),
             diagServices,
             new PolarisConfigurationStore() {},

polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreManagerTest.java

@@ -25,6 +25,7 @@
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.persistence.transactional.TransactionalMetaStoreManagerImpl;
 import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore;
 import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl;
@@ -37,6 +38,7 @@ public PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
     TreeMapMetaStore store = new TreeMapMetaStore(diagServices);
     PolarisCallContext callCtx =
         new PolarisCallContext(
+            Mockito.mock(RealmContext.class),
             new TreeMapTransactionalPersistenceImpl(store, Mockito.mock(), RANDOM_SECRETS),
             diagServices,
             new PolarisConfigurationStore() {},

polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java

@@ -40,7 +40,7 @@ protected PolarisCallContext callCtx() {
       TreeMapMetaStore store = new TreeMapMetaStore(diagServices);
       TreeMapTransactionalPersistenceImpl metaStore =
           new TreeMapTransactionalPersistenceImpl(store, Mockito.mock(), RANDOM_SECRETS);
-      callCtx = new PolarisCallContext(metaStore, diagServices);
+      callCtx = new PolarisCallContext(() -> "realm", metaStore, diagServices);
     }
     return callCtx;
   }

polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java

@@ -91,7 +91,7 @@ public InMemoryEntityCacheTest() {
     diagServices = new PolarisDefaultDiagServiceImpl();
     store = new TreeMapMetaStore(diagServices);
     metaStore = new TreeMapTransactionalPersistenceImpl(store, Mockito.mock(), RANDOM_SECRETS);
-    callCtx = new PolarisCallContext(metaStore, diagServices);
+    callCtx = new PolarisCallContext(() -> "realm", metaStore, diagServices);
     metaStoreManager = new TransactionalMetaStoreManagerImpl();
 
     // bootstrap the mata store with our test schema

polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java

@@ -96,6 +96,7 @@ public void testValidateAccessToLocationsWithWildcard() {
     Map<String, Boolean> config = Map.of("ALLOW_WILDCARD_LOCATION", true);
     PolarisCallContext polarisCallContext =
         new PolarisCallContext(
+            () -> "realm",
             Mockito.mock(),
             new PolarisDefaultDiagServiceImpl(),
             new PolarisConfigurationStore() {
@@ -106,7 +107,7 @@ public void testValidateAccessToLocationsWithWildcard() {
               }
             },
             Clock.systemUTC());
-    CallContext.setCurrentContext(CallContext.of(() -> "realm", polarisCallContext));
+    CallContext.setCurrentContext(polarisCallContext);
     Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>> result =
         storage.validateAccessToLocations(
             new FileStorageConfigurationInfo(List.of("file://", "*")),

polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java

@@ -70,7 +70,7 @@ public StorageCredentialCacheTest() {
     // to interact with the metastore
     TransactionalPersistence metaStore =
         new TreeMapTransactionalPersistenceImpl(store, Mockito.mock(), RANDOM_SECRETS);
-    callCtx = new PolarisCallContext(metaStore, diagServices);
+    callCtx = new PolarisCallContext(() -> "realm", metaStore, diagServices);
     metaStoreManager = Mockito.mock(PolarisMetaStoreManager.class);
     storageCredentialCache = new StorageCredentialCache();
   }

polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java

@@ -103,12 +103,8 @@ protected void testBrowse() {
 
   @Test
   protected void testCreateEntities() {
+    setCallContext();
     PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager;
-    CallContext callCtx =
-        CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext);
-    if (CallContext.getCurrentContext() == null) {
-      CallContext.setCurrentContext(callCtx);
-    }
     TaskEntity task1 = createTask("task1", 100L);
     TaskEntity task2 = createTask("task2", 101L);
     List<PolarisBaseEntity> createdEntities =
@@ -152,14 +148,24 @@ protected void testCreateEntities() {
                 task2.getSubTypeCode()));
   }
 
+  private void setCallContext() {
+    if (CallContext.getCurrentContext() == null) {
+      var oldCtx = polarisTestMetaStoreManager.polarisCallContext;
+      var ctx =
+          new PolarisCallContext(
+              () -> "testRealm",
+              oldCtx.getMetaStore(),
+              oldCtx.getDiagServices(),
+              oldCtx.getConfigurationStore(),
+              oldCtx.getClock());
+      CallContext.setCurrentContext(ctx);
+    }
+  }
+
   @Test
   protected void testCreateEntitiesAlreadyExisting() {
+    setCallContext();
     PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager;
-    CallContext callCtx =
-        CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext);
-    if (CallContext.getCurrentContext() == null) {
-      CallContext.setCurrentContext(callCtx);
-    }
     TaskEntity task1 = createTask("task1", 100L);
     TaskEntity task2 = createTask("task2", 101L);
     List<PolarisBaseEntity> createdEntities =
@@ -193,12 +199,8 @@ protected void testCreateEntitiesAlreadyExisting() {
 
   @Test
   protected void testCreateEntitiesWithConflict() {
+    setCallContext();
     PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager;
-    CallContext callCtx =
-        CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext);
-    if (CallContext.getCurrentContext() == null) {
-      CallContext.setCurrentContext(callCtx);
-    }
     TaskEntity task1 = createTask("task1", 100L);
     TaskEntity task2 = createTask("task2", 101L);
     TaskEntity task3 = createTask("task3", 103L);

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/ManagementServiceTest.java

@@ -68,14 +68,15 @@ public void setup() {
             .build();
     PolarisCallContext polarisCallContext =
         new PolarisCallContext(
+            fakeServices.realmContext(),
             fakeServices
                 .metaStoreManagerFactory()
                 .getOrCreateSessionSupplier(fakeServices.realmContext())
                 .get(),
             fakeServices.polarisDiagnostics(),
             fakeServices.configurationStore(),
             Mockito.mock(Clock.class));
-    CallContext.setCurrentContext(CallContext.of(fakeServices.realmContext(), polarisCallContext));
+    CallContext.setCurrentContext(polarisCallContext);
     services =
         TestServices.builder()
             .config(Map.of("SUPPORTED_CATALOG_STORAGE_TYPES", List.of("S3", "GCS", "AZURE")))
@@ -185,6 +186,7 @@ private PolarisCallContext setupCallContext(PolarisMetaStoreManager metaStoreMan
     MetaStoreManagerFactory metaStoreManagerFactory = services.metaStoreManagerFactory();
     RealmContext realmContext = services.realmContext();
     return new PolarisCallContext(
+        realmContext,
         metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(),
         services.polarisDiagnostics());
   }
@@ -193,7 +195,7 @@ private PolarisAdminService setupPolarisAdminService(
       PolarisMetaStoreManager metaStoreManager, PolarisCallContext callContext) {
     RealmContext realmContext = services.realmContext();
     return new PolarisAdminService(
-        CallContext.of(realmContext, callContext),
+        callContext,
         services.entityManagerFactory().getOrCreateEntityManager(realmContext),
         metaStoreManager,
         new UnsafeInMemorySecretsManager(),

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java

@@ -237,13 +237,14 @@ public void before(TestInfo testInfo) {
 
     polarisContext =
         new PolarisCallContext(
+            realmContext,
             managerFactory.getOrCreateSessionSupplier(realmContext).get(),
             diagServices,
             configurationStore,
             clock);
+    callContext = polarisContext;
     this.entityManager = realmEntityManagerFactory.getOrCreateEntityManager(realmContext);
 
-    callContext = CallContext.of(realmContext, polarisContext);
     CallContext.setCurrentContext(callContext);
 
     PrincipalEntity rootEntity =

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java

@@ -30,6 +30,7 @@
 import java.security.interfaces.RSAPublicKey;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
 import org.apache.polaris.core.entity.PolarisEntitySubType;
 import org.apache.polaris.core.entity.PolarisEntityType;
@@ -61,7 +62,8 @@ public void testSuccessfulTokenGeneration() throws Exception {
     final String scope = "PRINCIPAL_ROLE:TEST";
 
     PolarisCallContext polarisCallContext =
-        new PolarisCallContext(null, null, configurationStore, null);
+        new PolarisCallContext(
+            Mockito.mock(RealmContext.class), null, null, configurationStore, null);
     PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class);
     String mainSecret = "client-secret";
     PolarisPrincipalSecrets principalSecrets =

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java

@@ -45,7 +45,7 @@ public class JWTSymmetricKeyGeneratorTest {
   /** Sanity test to verify that we can generate a token */
   @Test
   public void testJWTSymmetricKeyGenerator() {
-    PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, null, null);
+    PolarisCallContext polarisCallContext = Mockito.mock(PolarisCallContext.class);
     PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class);
     String mainSecret = "test_secret";
     String clientId = "test_client_id";

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogTest.java

@@ -250,17 +250,17 @@ public void before(TestInfo testInfo) {
     userSecretsManager = userSecretsManagerFactory.getOrCreateUserSecretsManager(realmContext);
     polarisContext =
         new PolarisCallContext(
+            realmContext,
             managerFactory.getOrCreateSessionSupplier(realmContext).get(),
             diagServices,
             configurationStore,
             Clock.systemDefaultZone());
+    callContext = polarisContext;
 
     entityManager =
         new PolarisEntityManager(
             metaStoreManager, new StorageCredentialCache(), createEntityCache(metaStoreManager));
 
-    callContext = CallContext.of(realmContext, polarisContext);
-
     PrincipalEntity rootEntity =
         new PrincipalEntity(
             PolarisEntity.of(

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/IcebergCatalogViewTest.java

@@ -170,6 +170,7 @@ public void before(TestInfo testInfo) {
     userSecretsManager = userSecretsManagerFactory.getOrCreateUserSecretsManager(realmContext);
     polarisContext =
         new PolarisCallContext(
+            realmContext,
             managerFactory.getOrCreateSessionSupplier(realmContext).get(),
             diagServices,
             configurationStore,
@@ -181,8 +182,7 @@ public void before(TestInfo testInfo) {
             new StorageCredentialCache(),
             new InMemoryEntityCache(metaStoreManager));
 
-    CallContext callContext = CallContext.of(realmContext, polarisContext);
-    CallContext.setCurrentContext(callContext);
+    CallContext.setCurrentContext(polarisContext);
 
     PrincipalEntity rootEntity =
         new PrincipalEntity(
@@ -206,7 +206,7 @@ public void before(TestInfo testInfo) {
 
     PolarisAdminService adminService =
         new PolarisAdminService(
-            callContext,
+            polarisContext,
             entityManager,
             metaStoreManager,
             userSecretsManager,
@@ -232,7 +232,7 @@ public void before(TestInfo testInfo) {
 
     PolarisPassthroughResolutionView passthroughView =
         new PolarisPassthroughResolutionView(
-            callContext, entityManager, securityContext, CATALOG_NAME);
+            polarisContext, entityManager, securityContext, CATALOG_NAME);
     FileIOFactory fileIOFactory =
         new DefaultFileIOFactory(
             new RealmEntityManagerFactory(managerFactory), managerFactory, configurationStore);
@@ -242,7 +242,7 @@ public void before(TestInfo testInfo) {
         new IcebergCatalog(
             entityManager,
             metaStoreManager,
-            callContext,
+            polarisContext,
             passthroughView,
             securityContext,
             Mockito.mock(),