Introduce TransactionalMetadataFactory

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveConnector.java

@@ -35,7 +35,6 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Supplier;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static io.prestosql.spi.transaction.IsolationLevel.READ_UNCOMMITTED;
@@ -46,7 +45,7 @@ public class HiveConnector
         implements Connector
 {
     private final LifeCycleManager lifeCycleManager;
-    private final Supplier<TransactionalMetadata> metadataFactory;
+    private final TransactionalMetadataFactory metadataFactory;
     private final ConnectorSplitManager splitManager;
     private final ConnectorPageSourceProvider pageSourceProvider;
     private final ConnectorPageSinkProvider pageSinkProvider;
@@ -65,7 +64,7 @@ public class HiveConnector
 
     public HiveConnector(
             LifeCycleManager lifeCycleManager,
-            Supplier<TransactionalMetadata> metadataFactory,
+            TransactionalMetadataFactory metadataFactory,
             HiveTransactionManager transactionManager,
             ConnectorSplitManager splitManager,
             ConnectorPageSourceProvider pageSourceProvider,
@@ -189,7 +188,7 @@ public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel
         checkConnectorSupports(READ_UNCOMMITTED, isolationLevel);
         ConnectorTransactionHandle transaction = new HiveTransactionHandle();
         try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
-            transactionManager.put(transaction, metadataFactory.get());
+            transactionManager.put(transaction, metadataFactory.create(transaction));
         }
         return transaction;
     }

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveMetadataFactory.java

@@ -21,6 +21,7 @@
 import io.prestosql.plugin.hive.metastore.SemiTransactionalHiveMetastore;
 import io.prestosql.plugin.hive.security.AccessControlMetadataFactory;
 import io.prestosql.plugin.hive.statistics.MetastoreHiveStatisticsProvider;
+import io.prestosql.spi.connector.ConnectorTransactionHandle;
 import io.prestosql.spi.type.TypeManager;
 import org.joda.time.DateTimeZone;
 
@@ -29,13 +30,12 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.function.Supplier;
 
 import static io.prestosql.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore;
 import static java.util.Objects.requireNonNull;
 
 public class HiveMetadataFactory
-        implements Supplier<TransactionalMetadata>
+        implements TransactionalMetadataFactory
 {
     private static final Logger log = Logger.get(HiveMetadataFactory.class);
 
@@ -164,7 +164,7 @@ public HiveMetadataFactory(
     }
 
     @Override
-    public HiveMetadata get()
+    public TransactionalMetadata create(ConnectorTransactionHandle transactionHandle)
     {
         SemiTransactionalHiveMetastore metastore = new SemiTransactionalHiveMetastore(
                 hdfsEnvironment,
@@ -192,6 +192,6 @@ public HiveMetadata get()
                 typeTranslator,
                 prestoVersion,
                 new MetastoreHiveStatisticsProvider(metastore),
-                accessControlMetadataFactory.create(metastore));
+                accessControlMetadataFactory.create(transactionHandle, metastore));
     }
 }

presto-hive/src/main/java/io/prestosql/plugin/hive/HiveModule.java

@@ -19,7 +19,6 @@
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
-import com.google.inject.TypeLiteral;
 import com.google.inject.multibindings.Multibinder;
 import io.airlift.event.client.EventClient;
 import io.prestosql.plugin.hive.metastore.SemiTransactionalHiveMetastore;
@@ -48,7 +47,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
@@ -90,7 +88,7 @@ public void configure(Binder binder)
         binder.bind(HivePartitionManager.class).in(Scopes.SINGLETON);
         binder.bind(LocationService.class).to(HiveLocationService.class).in(Scopes.SINGLETON);
         binder.bind(HiveMetadataFactory.class).in(Scopes.SINGLETON);
-        binder.bind(new TypeLiteral<Supplier<TransactionalMetadata>>() {}).to(HiveMetadataFactory.class).in(Scopes.SINGLETON);
+        binder.bind(TransactionalMetadataFactory.class).to(HiveMetadataFactory.class).in(Scopes.SINGLETON);
         binder.bind(HiveTransactionManager.class).in(Scopes.SINGLETON);
         binder.bind(ConnectorSplitManager.class).to(HiveSplitManager.class).in(Scopes.SINGLETON);
         newExporter(binder).export(ConnectorSplitManager.class).as(generator -> generator.generatedNameOf(HiveSplitManager.class));

presto-hive/src/main/java/io/prestosql/plugin/hive/TransactionalMetadataFactory.java

@@ -0,0 +1,21 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.hive;
+
+import io.prestosql.spi.connector.ConnectorTransactionHandle;
+
+public interface TransactionalMetadataFactory
+{
+    TransactionalMetadata create(ConnectorTransactionHandle transactionHandle);
+}

presto-hive/src/main/java/io/prestosql/plugin/hive/procedure/CreateEmptyPartitionProcedure.java

@@ -21,11 +21,13 @@
 import io.prestosql.plugin.hive.HiveInsertTableHandle;
 import io.prestosql.plugin.hive.HiveMetastoreClosure;
 import io.prestosql.plugin.hive.HiveTableHandle;
+import io.prestosql.plugin.hive.HiveTransactionHandle;
 import io.prestosql.plugin.hive.LocationService;
 import io.prestosql.plugin.hive.LocationService.WriteInfo;
 import io.prestosql.plugin.hive.PartitionUpdate;
 import io.prestosql.plugin.hive.PartitionUpdate.UpdateMode;
 import io.prestosql.plugin.hive.TransactionalMetadata;
+import io.prestosql.plugin.hive.TransactionalMetadataFactory;
 import io.prestosql.plugin.hive.authentication.HiveIdentity;
 import io.prestosql.plugin.hive.metastore.HiveMetastore;
 import io.prestosql.spi.PrestoException;
@@ -44,7 +46,6 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
-import java.util.function.Supplier;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.prestosql.spi.StandardErrorCode.ALREADY_EXISTS;
@@ -66,13 +67,13 @@ public class CreateEmptyPartitionProcedure
             List.class,
             List.class);
 
-    private final Supplier<TransactionalMetadata> hiveMetadataFactory;
+    private final TransactionalMetadataFactory hiveMetadataFactory;
     private final HiveMetastoreClosure metastore;
     private final LocationService locationService;
     private final JsonCodec<PartitionUpdate> partitionUpdateJsonCodec;
 
     @Inject
-    public CreateEmptyPartitionProcedure(Supplier<TransactionalMetadata> hiveMetadataFactory, HiveMetastore metastore, LocationService locationService, JsonCodec<PartitionUpdate> partitionUpdateCodec)
+    public CreateEmptyPartitionProcedure(TransactionalMetadataFactory hiveMetadataFactory, HiveMetastore metastore, LocationService locationService, JsonCodec<PartitionUpdate> partitionUpdateCodec)
     {
         this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null");
         this.metastore = new HiveMetastoreClosure(requireNonNull(metastore, "metastore is null"));
@@ -103,7 +104,7 @@ public void createEmptyPartition(ConnectorSession session, String schema, String
 
     private void doCreateEmptyPartition(ConnectorSession session, String schemaName, String tableName, List<String> partitionColumnNames, List<String> partitionValues)
     {
-        TransactionalMetadata hiveMetadata = hiveMetadataFactory.get();
+        TransactionalMetadata hiveMetadata = hiveMetadataFactory.create(new HiveTransactionHandle());
         HiveTableHandle tableHandle = (HiveTableHandle) hiveMetadata.getTableHandle(session, new SchemaTableName(schemaName, tableName));
         if (tableHandle == null) {
             throw new PrestoException(INVALID_PROCEDURE_ARGUMENT, format("Table %s does not exist", new SchemaTableName(schemaName, tableName)));

presto-hive/src/main/java/io/prestosql/plugin/hive/procedure/DropStatsProcedure.java

@@ -17,8 +17,10 @@
 import io.prestosql.plugin.hive.HiveColumnHandle;
 import io.prestosql.plugin.hive.HiveMetastoreClosure;
 import io.prestosql.plugin.hive.HiveTableHandle;
+import io.prestosql.plugin.hive.HiveTransactionHandle;
 import io.prestosql.plugin.hive.PartitionStatistics;
 import io.prestosql.plugin.hive.TransactionalMetadata;
+import io.prestosql.plugin.hive.TransactionalMetadataFactory;
 import io.prestosql.plugin.hive.authentication.HiveIdentity;
 import io.prestosql.plugin.hive.metastore.HiveMetastore;
 import io.prestosql.spi.PrestoException;
@@ -36,7 +38,6 @@
 import java.lang.invoke.MethodHandle;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Supplier;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.prestosql.spi.StandardErrorCode.INVALID_PROCEDURE_ARGUMENT;
@@ -62,11 +63,11 @@ public class DropStatsProcedure
             String.class,
             List.class);
 
-    private final Supplier<TransactionalMetadata> hiveMetadataFactory;
+    private final TransactionalMetadataFactory hiveMetadataFactory;
     private final HiveMetastoreClosure metastore;
 
     @Inject
-    public DropStatsProcedure(Supplier<TransactionalMetadata> hiveMetadataFactory, HiveMetastore metastore)
+    public DropStatsProcedure(TransactionalMetadataFactory hiveMetadataFactory, HiveMetastore metastore)
     {
         this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null");
         this.metastore = new HiveMetastoreClosure(requireNonNull(metastore, "metastore is null"));
@@ -94,7 +95,7 @@ public void dropStats(ConnectorSession session, String schema, String table, Lis
 
     private void doDropStats(ConnectorSession session, String schema, String table, List<?> partitionValues)
     {
-        TransactionalMetadata hiveMetadata = hiveMetadataFactory.get();
+        TransactionalMetadata hiveMetadata = hiveMetadataFactory.create(new HiveTransactionHandle());
         HiveTableHandle handle = (HiveTableHandle) hiveMetadata.getTableHandle(session, new SchemaTableName(schema, table));
         if (handle == null) {
             throw new PrestoException(INVALID_PROCEDURE_ARGUMENT, format("Table %s does not exist", new SchemaTableName(schema, table)));

presto-hive/src/main/java/io/prestosql/plugin/hive/procedure/RegisterPartitionProcedure.java

@@ -20,8 +20,9 @@
 import io.prestosql.plugin.hive.HiveConfig;
 import io.prestosql.plugin.hive.HiveMetadata;
 import io.prestosql.plugin.hive.HiveMetastoreClosure;
+import io.prestosql.plugin.hive.HiveTransactionHandle;
 import io.prestosql.plugin.hive.PartitionStatistics;
-import io.prestosql.plugin.hive.TransactionalMetadata;
+import io.prestosql.plugin.hive.TransactionalMetadataFactory;
 import io.prestosql.plugin.hive.authentication.HiveIdentity;
 import io.prestosql.plugin.hive.metastore.HiveMetastore;
 import io.prestosql.plugin.hive.metastore.Partition;
@@ -44,7 +45,6 @@
 import java.lang.invoke.MethodHandle;
 import java.util.List;
 import java.util.Optional;
-import java.util.function.Supplier;
 
 import static io.prestosql.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME;
 import static io.prestosql.plugin.hive.procedure.Procedures.checkIsPartitionedTable;
@@ -71,12 +71,12 @@ public class RegisterPartitionProcedure
             String.class);
 
     private final boolean allowRegisterPartition;
-    private final Supplier<TransactionalMetadata> hiveMetadataFactory;
+    private final TransactionalMetadataFactory hiveMetadataFactory;
     private final HdfsEnvironment hdfsEnvironment;
     private final HiveMetastoreClosure metastore;
 
     @Inject
-    public RegisterPartitionProcedure(HiveConfig hiveConfig, Supplier<TransactionalMetadata> hiveMetadataFactory, HiveMetastore metastore, HdfsEnvironment hdfsEnvironment)
+    public RegisterPartitionProcedure(HiveConfig hiveConfig, TransactionalMetadataFactory hiveMetadataFactory, HiveMetastore metastore, HdfsEnvironment hdfsEnvironment)
     {
         this.allowRegisterPartition = requireNonNull(hiveConfig, "hiveConfig is null").isAllowRegisterPartition();
         this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null");
@@ -134,7 +134,7 @@ private void doRegisterPartition(ConnectorSession session, String schemaName, St
             throw new PrestoException(INVALID_PROCEDURE_ARGUMENT, "Partition location does not exist: " + partitionLocation);
         }
 
-        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.get()).getMetastore();
+        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.create(new HiveTransactionHandle())).getMetastore();
 
         metastore.addPartition(
                 session,

presto-hive/src/main/java/io/prestosql/plugin/hive/procedure/SyncPartitionMetadataProcedure.java

@@ -18,8 +18,9 @@
 import com.google.common.collect.Sets;
 import io.prestosql.plugin.hive.HdfsEnvironment;
 import io.prestosql.plugin.hive.HiveMetadata;
+import io.prestosql.plugin.hive.HiveTransactionHandle;
 import io.prestosql.plugin.hive.PartitionStatistics;
-import io.prestosql.plugin.hive.TransactionalMetadata;
+import io.prestosql.plugin.hive.TransactionalMetadataFactory;
 import io.prestosql.plugin.hive.authentication.HiveIdentity;
 import io.prestosql.plugin.hive.metastore.Column;
 import io.prestosql.plugin.hive.metastore.Partition;
@@ -44,7 +45,6 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Supplier;
 import java.util.stream.Stream;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -74,12 +74,12 @@ public enum SyncMode
             String.class,
             String.class);
 
-    private final Supplier<TransactionalMetadata> hiveMetadataFactory;
+    private final TransactionalMetadataFactory hiveMetadataFactory;
     private final HdfsEnvironment hdfsEnvironment;
 
     @Inject
     public SyncPartitionMetadataProcedure(
-            Supplier<TransactionalMetadata> hiveMetadataFactory,
+            TransactionalMetadataFactory hiveMetadataFactory,
             HdfsEnvironment hdfsEnvironment)
     {
         this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null");
@@ -111,7 +111,7 @@ private void doSyncPartitionMetadata(ConnectorSession session, String schemaName
         SyncMode syncMode = toSyncMode(mode);
         HdfsContext hdfsContext = new HdfsContext(session, schemaName, tableName);
         HiveIdentity identity = new HiveIdentity(session);
-        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.get()).getMetastore();
+        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.create(new HiveTransactionHandle())).getMetastore();
         SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
 
         Table table = metastore.getTable(identity, schemaName, tableName)

presto-hive/src/main/java/io/prestosql/plugin/hive/procedure/UnregisterPartitionProcedure.java

@@ -16,7 +16,8 @@
 import com.google.common.collect.ImmutableList;
 import io.prestosql.plugin.hive.HiveMetadata;
 import io.prestosql.plugin.hive.HiveMetastoreClosure;
-import io.prestosql.plugin.hive.TransactionalMetadata;
+import io.prestosql.plugin.hive.HiveTransactionHandle;
+import io.prestosql.plugin.hive.TransactionalMetadataFactory;
 import io.prestosql.plugin.hive.authentication.HiveIdentity;
 import io.prestosql.plugin.hive.metastore.HiveMetastore;
 import io.prestosql.plugin.hive.metastore.Partition;
@@ -36,7 +37,6 @@
 
 import java.lang.invoke.MethodHandle;
 import java.util.List;
-import java.util.function.Supplier;
 
 import static io.prestosql.plugin.hive.procedure.Procedures.checkIsPartitionedTable;
 import static io.prestosql.plugin.hive.procedure.Procedures.checkPartitionColumns;
@@ -58,11 +58,11 @@ public class UnregisterPartitionProcedure
             List.class,
             List.class);
 
-    private final Supplier<TransactionalMetadata> hiveMetadataFactory;
+    private final TransactionalMetadataFactory hiveMetadataFactory;
     private final HiveMetastoreClosure metastore;
 
     @Inject
-    public UnregisterPartitionProcedure(Supplier<TransactionalMetadata> hiveMetadataFactory, HiveMetastore metastore)
+    public UnregisterPartitionProcedure(TransactionalMetadataFactory hiveMetadataFactory, HiveMetastore metastore)
     {
         this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null");
         this.metastore = new HiveMetastoreClosure(requireNonNull(metastore, "metastore is null"));
@@ -105,7 +105,7 @@ private void doUnregisterPartition(ConnectorSession session, String schemaName,
         Partition partition = metastore.getPartition(new HiveIdentity(session), schemaName, tableName, partitionValues)
                 .orElseThrow(() -> new PrestoException(NOT_FOUND, format("Partition %s does not exist", partitionName)));
 
-        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.get()).getMetastore();
+        SemiTransactionalHiveMetastore metastore = ((HiveMetadata) hiveMetadataFactory.create(new HiveTransactionHandle())).getMetastore();
 
         metastore.dropPartition(
                 session,

presto-hive/src/main/java/io/prestosql/plugin/hive/security/AccessControlMetadataFactory.java

@@ -14,8 +14,9 @@
 package io.prestosql.plugin.hive.security;
 
 import io.prestosql.plugin.hive.metastore.SemiTransactionalHiveMetastore;
+import io.prestosql.spi.connector.ConnectorTransactionHandle;
 
 public interface AccessControlMetadataFactory
 {
-    AccessControlMetadata create(SemiTransactionalHiveMetastore metastore);
+    AccessControlMetadata create(ConnectorTransactionHandle transactionHandle, SemiTransactionalHiveMetastore metastore);
 }

presto-hive/src/main/java/io/prestosql/plugin/hive/security/HiveSecurityModule.java

@@ -68,7 +68,7 @@ private static class StaticAccessControlMetadataModule
         @Override
         public void configure(Binder binder)
         {
-            binder.bind(AccessControlMetadataFactory.class).toInstance(metastore -> new AccessControlMetadata() {});
+            binder.bind(AccessControlMetadataFactory.class).toInstance((transactionHandle, metastore) -> new AccessControlMetadata() {});
         }
     }
 }

presto-hive/src/main/java/io/prestosql/plugin/hive/security/SqlStandardSecurityModule.java

@@ -18,6 +18,7 @@
 import com.google.inject.Scopes;
 import io.prestosql.plugin.hive.metastore.SemiTransactionalHiveMetastore;
 import io.prestosql.spi.connector.ConnectorAccessControl;
+import io.prestosql.spi.connector.ConnectorTransactionHandle;
 
 public class SqlStandardSecurityModule
         implements Module
@@ -35,7 +36,7 @@ private static final class SqlStandardAccessControlMetadataFactory
         public SqlStandardAccessControlMetadataFactory() {}
 
         @Override
-        public AccessControlMetadata create(SemiTransactionalHiveMetastore metastore)
+        public AccessControlMetadata create(ConnectorTransactionHandle transactionHandle, SemiTransactionalHiveMetastore metastore)
         {
             return new SqlStandardAccessControlMetadata(metastore);
         }

presto-hive/src/test/java/io/prestosql/plugin/hive/AbstractTestHive.java

@@ -759,7 +759,7 @@ protected final void setup(String databaseName, HiveConfig hiveConfig, HiveMetas
                 heartbeatService,
                 new HiveTypeTranslator(),
                 TEST_SERVER_VERSION,
-                SqlStandardAccessControlMetadata::new);
+                (transactionHandle, metastore) -> new SqlStandardAccessControlMetadata(metastore));
         transactionManager = new HiveTransactionManager();
         splitManager = new HiveSplitManager(
                 transactionHandle -> ((HiveMetadata) transactionManager.get(transactionHandle)).getMetastore(),
@@ -838,7 +838,7 @@ protected ConnectorSession newSession(Map<String, Object> propertyValues)
 
     protected Transaction newTransaction()
     {
-        return new HiveTransaction(transactionManager, metadataFactory.get());
+        return new HiveTransaction(transactionManager, (HiveMetadata) metadataFactory.create(new HiveTransactionHandle()));
     }
 
     interface Transaction