Address review feedback

Author: singhpk234

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

@@ -87,7 +87,15 @@ public void writeEntity(
       @Nonnull PolarisBaseEntity entity,
       boolean nameOrParentChanged,
       PolarisBaseEntity originalEntity) {
-    persistEntity(callCtx, entity, originalEntity, datasourceOperations);
+    try {
+      datasourceOperations.runWithinTransaction(
+          statement -> {
+            persistEntity(callCtx, entity, originalEntity, statement);
+            return true;
+          });
+    } catch (SQLException e) {
+      throw new RuntimeException("Error persisting entity", e);
+    }
   }
 
   @Override
@@ -130,11 +138,12 @@ private void persistEntity(
       @Nonnull PolarisCallContext callCtx,
       @Nonnull PolarisBaseEntity entity,
       PolarisBaseEntity originalEntity,
-      Object executor) {
+      Statement statement)
+      throws SQLException {
     ModelEntity modelEntity = ModelEntity.fromEntity(entity);
     if (originalEntity == null) {
       try {
-        execute(executor, generateInsertQuery(modelEntity, realmId));
+        statement.executeUpdate(generateInsertQuery(modelEntity, realmId));
       } catch (SQLException e) {
         if (datasourceOperations.isConstraintViolation(e)) {
           PolarisBaseEntity existingEntity =
@@ -162,7 +171,7 @@ private void persistEntity(
               "realm_id",
               realmId);
       try {
-        int rowsUpdated = execute(executor, generateUpdateQuery(modelEntity, params));
+        int rowsUpdated = statement.executeUpdate(generateUpdateQuery(modelEntity, params));
         if (rowsUpdated == 0) {
           throw new RetryOnConcurrencyException(
               "Entity '%s' id '%s' concurrently modified; expected version %s",
@@ -175,16 +184,6 @@ private void persistEntity(
     }
   }
 
-  private int execute(Object executor, String query) throws SQLException {
-    if (executor instanceof Statement) {
-      return ((Statement) executor).executeUpdate(query);
-    } else if (executor instanceof DatasourceOperations) {
-      return ((DatasourceOperations) executor).executeUpdate(query);
-    } else {
-      throw new IllegalArgumentException("Unsupported executor: " + executor);
-    }
-  }
-
   @Override
   public void writeToGrantRecords(
       @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) {

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

@@ -18,15 +18,14 @@
  */
 package org.apache.polaris.extension.persistence.relational.jdbc;
 
-import com.google.common.base.Preconditions;
 import io.smallrye.common.annotation.Identifier;
 import jakarta.annotation.Nullable;
 import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.enterprise.inject.Instance;
 import jakarta.inject.Inject;
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
 import javax.sql.DataSource;
@@ -70,15 +69,12 @@ public class JdbcMetaStoreManagerFactory implements MetaStoreManagerFactory {
   final Map<String, StorageCredentialCache> storageCredentialCacheMap = new HashMap<>();
   final Map<String, EntityCache> entityCacheMap = new HashMap<>();
   final Map<String, Supplier<BasePersistence>> sessionSupplierMap = new HashMap<>();
-  final List<DataSource> dataSources;
-
   protected final PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl();
+
   @Inject PolarisStorageIntegrationProvider storageIntegrationProvider;
+  @Inject Instance<DataSource> dataSource;
 
-  @Inject
-  protected JdbcMetaStoreManagerFactory(List<DataSource> dataSources) {
-    this.dataSources = dataSources;
-  }
+  protected JdbcMetaStoreManagerFactory() {}
 
   protected PrincipalSecretsGenerator secretsGenerator(
       RealmContext realmContext, @Nullable RootCredentialsSet rootCredentialsSet) {
@@ -111,13 +107,11 @@ private void initializeForRealm(
   }
 
   private DatasourceOperations getDatasourceOperations(boolean isBootstrap) {
-    // TODO: remove when multiple dataSources are supported.
-    Preconditions.checkState(dataSources.size() == 1, "More than one dataSources configured");
-    DatasourceOperations databaseOperations = new DatasourceOperations(dataSources.getFirst());
+    DatasourceOperations databaseOperations = new DatasourceOperations(dataSource.get());
     if (isBootstrap) {
       try {
         DatabaseType databaseType;
-        try (Connection connection = dataSources.getFirst().getConnection()) {
+        try (Connection connection = dataSource.get().getConnection()) {
           String productName = connection.getMetaData().getDatabaseProductName();
           databaseType = DatabaseType.fromDisplayName(productName);
         }

quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusProducers.java

@@ -18,18 +18,13 @@
  */
 package org.apache.polaris.admintool.config;
 
-import io.quarkus.arc.All;
-import io.quarkus.arc.InstanceHandle;
 import io.smallrye.common.annotation.Identifier;
 import jakarta.annotation.Nullable;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.enterprise.inject.Any;
 import jakarta.enterprise.inject.Instance;
 import jakarta.enterprise.inject.Produces;
 import java.time.Clock;
-import java.util.ArrayList;
-import java.util.List;
-import javax.sql.DataSource;
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
@@ -81,13 +76,4 @@ public PolarisConfigurationStore configurationStore() {
     // A configuration store is not required when running the admin tool.
     return new PolarisConfigurationStore() {};
   }
-
-  @Produces
-  public List<DataSource> dataSources(@All List<InstanceHandle<DataSource>> dataSources) {
-    List<DataSource> dataSourceList = new ArrayList<>();
-    for (InstanceHandle<DataSource> handle : dataSources) {
-      dataSourceList.add(handle.get());
-    }
-    return dataSourceList;
-  }
 }

quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java

@@ -18,8 +18,6 @@
  */
 package org.apache.polaris.service.quarkus.config;
 
-import io.quarkus.arc.All;
-import io.quarkus.arc.InstanceHandle;
 import io.quarkus.runtime.StartupEvent;
 import io.smallrye.common.annotation.Identifier;
 import io.smallrye.context.SmallRyeManagedExecutor;
@@ -34,9 +32,6 @@
 import jakarta.ws.rs.container.ContainerRequestContext;
 import jakarta.ws.rs.core.Context;
 import java.time.Clock;
-import java.util.ArrayList;
-import java.util.List;
-import javax.sql.DataSource;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
@@ -266,13 +261,4 @@ public ActiveRolesProvider activeRolesProvider(
   public void closeTaskExecutor(@Disposes @Identifier("task-executor") ManagedExecutor executor) {
     executor.close();
   }
-
-  @Produces
-  public List<DataSource> dataSources(@All List<InstanceHandle<DataSource>> dataSources) {
-    List<DataSource> dataSourceList = new ArrayList<>();
-    for (InstanceHandle<DataSource> handle : dataSources) {
-      dataSourceList.add(handle.get());
-    }
-    return dataSourceList;
-  }
 }