Test hive metastore operations involving retries

trinodb · Jan 10, 2024 · 5113d9b · 5113d9b
1 parent 41687ef
commit 5113d9b
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 3 deletions.
diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java
@@ -18,6 +18,7 @@
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.filesystem.hdfs.HdfsFileSystemFactory;
 import io.trino.plugin.hive.metastore.HiveMetastoreConfig;
+import io.trino.plugin.hive.metastore.thrift.MetastoreClientAdapterProvider;
 import io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory;
 import io.trino.plugin.hive.metastore.thrift.ThriftHiveMetastoreFactory;
 import io.trino.plugin.hive.metastore.thrift.ThriftMetastore;
@@ -32,6 +33,7 @@
 import static io.trino.plugin.base.security.UserNameProvider.SIMPLE_USER_NAME_PROVIDER;
 import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT;
 import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS;
+import static io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory.TIMEOUT;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.Executors.newFixedThreadPool;
 
@@ -58,6 +60,14 @@ public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address, Du
         return this;
     }
 
+    public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address, MetastoreClientAdapterProvider metastoreClientAdapterProvider)
+    {
+        requireNonNull(address, "address is null");
+        checkState(tokenAwareMetastoreClientFactory == null, "Metastore client already set");
+        tokenAwareMetastoreClientFactory = new TestingTokenAwareMetastoreClientFactory(HiveTestUtils.SOCKS_PROXY, address, TIMEOUT, metastoreClientAdapterProvider);
+        return this;
+    }
+
     public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address)
     {
         requireNonNull(address, "address is null");

diff --git a/...in/trino-hive/src/test/java/io/trino/plugin/hive/metastore/AbstractTestHiveMetastore.java b/...in/trino-hive/src/test/java/io/trino/plugin/hive/metastore/AbstractTestHiveMetastore.java
@@ -45,6 +45,11 @@ public void setMetastore(HiveMetastore metastore)
         this.metastore = metastore;
     }
 
+    protected HiveMetastore getMetastore()
+    {
+        return metastore;
+    }
+
     @Test
     void testCreateDatabase()
     {

diff --git a/...in/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java b/...in/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java
@@ -13,14 +13,31 @@
  */
 package io.trino.plugin.hive.metastore;
 
+import io.trino.plugin.hive.SchemaAlreadyExistsException;
+import io.trino.plugin.hive.TableAlreadyExistsException;
 import io.trino.plugin.hive.containers.HiveHadoop;
 import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore;
+import io.trino.plugin.hive.metastore.thrift.MetastoreClientAdapterProvider;
+import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreClient;
 import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.parallel.Execution;
 
+import java.lang.reflect.InvocationTargetException;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.google.common.reflect.Reflection.newProxy;
+import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME;
+import static io.trino.plugin.hive.HiveStorageFormat.PARQUET;
+import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE;
 import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder;
+import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES;
+import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat;
+import static io.trino.testing.TestingNames.randomNameSuffix;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
 import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD;
 
@@ -36,8 +53,22 @@ final class TestBridgingHiveMetastore
         hiveHadoop = HiveHadoop.builder().build();
         hiveHadoop.start();
 
+        MetastoreClientAdapterProvider metastoreClientAdapterProvider = delegate -> newProxy(ThriftMetastoreClient.class, (proxy, method, methodArgs) -> {
+            Object result;
+            try {
+                result = method.invoke(delegate, methodArgs);
+            }
+            catch (InvocationTargetException e) {
+                throw e.getCause();
+            }
+            if (method.getName().equals("createDatabase") || method.getName().equals("createTable")) {
+                throw new RuntimeException("Test-simulated Hive Metastore timeout exception");
+            }
+            return result;
+        });
+
         setMetastore(new BridgingHiveMetastore(testingThriftHiveMetastoreBuilder()
-                .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint())
+                .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint(), metastoreClientAdapterProvider)
                 .thriftMetastoreConfig(new ThriftMetastoreConfig().setDeleteFilesOnDrop(true))
                 .build()));
     }
@@ -47,4 +78,53 @@ void afterAll()
     {
         hiveHadoop.stop();
     }
+
+    @Test
+    public void testCreateDatabaseWithRetries()
+    {
+        // This test is similar to AbstractTestHiveMetastore#testCreateDatabase but with simulating timeout in ThriftMetastoreClient
+        String databaseName = "test_database_" + randomNameSuffix();
+        Database.Builder database = Database.builder()
+                .setDatabaseName(databaseName)
+                .setParameters(Map.of(TRINO_QUERY_ID_NAME, "query_id"))
+                .setOwnerName(Optional.empty())
+                .setOwnerType(Optional.empty());
+        getMetastore().createDatabase(database.build());
+
+        database.setParameters(Map.of(TRINO_QUERY_ID_NAME, "another_query_id"));
+        assertThatThrownBy(() -> getMetastore().createDatabase(database.build()))
+                .isInstanceOf(SchemaAlreadyExistsException.class);
+
+        getMetastore().dropDatabase(databaseName, false);
+    }
+
+    @Test
+    public void testCreateTableWithRetries()
+    {
+        // This test is similar to AbstractTestHiveMetastore#testCreateTable but with simulating timeout in ThriftMetastoreClient
+        String databaseName = "test_database_" + randomNameSuffix();
+        Database.Builder database = Database.builder()
+                .setDatabaseName(databaseName)
+                .setOwnerName(Optional.empty())
+                .setOwnerType(Optional.empty());
+        getMetastore().createDatabase(database.build());
+
+        String tableName = "test_table" + randomNameSuffix();
+        Table.Builder table = Table.builder()
+                .setDatabaseName(databaseName)
+                .setTableName(tableName)
+                .setParameters(Map.of(TRINO_QUERY_ID_NAME, "query_id"))
+                .setTableType(EXTERNAL_TABLE.name())
+                .setOwner(Optional.empty());
+        table.getStorageBuilder()
+                .setStorageFormat(fromHiveStorageFormat(PARQUET));
+        getMetastore().createTable(table.build(), NO_PRIVILEGES);
+
+        table.setParameters(Map.of(TRINO_QUERY_ID_NAME, "another_query_id"));
+        assertThatThrownBy(() -> getMetastore().createTable(table.build(), NO_PRIVILEGES))
+                .isInstanceOf(TableAlreadyExistsException.class);
+
+        getMetastore().dropTable(databaseName, tableName, false);
+        getMetastore().dropDatabase(databaseName, false);
+    }
 }
diff --git a/...e/src/test/java/io/trino/plugin/hive/metastore/thrift/MetastoreClientAdapterProvider.java b/...e/src/test/java/io/trino/plugin/hive/metastore/thrift/MetastoreClientAdapterProvider.java
@@ -0,0 +1,19 @@
+/*
+ * 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.trino.plugin.hive.metastore.thrift;
+
+public interface MetastoreClientAdapterProvider
+{
+    ThriftMetastoreClient createThriftMetastoreClientAdapter(ThriftMetastoreClient delegate);
+}
diff --git a/...t/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java b/...t/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java
@@ -31,21 +31,29 @@ public class TestingTokenAwareMetastoreClientFactory
     private final DefaultThriftMetastoreClientFactory factory;
     private final HostAndPort address;
 
+    private final MetastoreClientAdapterProvider metastoreClientAdapterProvider;
+
     public TestingTokenAwareMetastoreClientFactory(Optional<HostAndPort> socksProxy, HostAndPort address)
     {
-        this(socksProxy, address, TIMEOUT);
+        this(socksProxy, address, TIMEOUT, delegate -> delegate);
     }
 
     public TestingTokenAwareMetastoreClientFactory(Optional<HostAndPort> socksProxy, HostAndPort address, Duration timeout)
+    {
+        this(socksProxy, address, timeout, delegate -> delegate);
+    }
+
+    public TestingTokenAwareMetastoreClientFactory(Optional<HostAndPort> socksProxy, HostAndPort address, Duration timeout, MetastoreClientAdapterProvider metastoreClientAdapterProvider)
     {
         this.factory = new DefaultThriftMetastoreClientFactory(Optional.empty(), socksProxy, timeout, timeout, AUTHENTICATION, "localhost");
         this.address = requireNonNull(address, "address is null");
+        this.metastoreClientAdapterProvider = requireNonNull(metastoreClientAdapterProvider, "metastoreClientAdapterProvider is null");
     }
 
     @Override
     public ThriftMetastoreClient createMetastoreClient(Optional<String> delegationToken)
             throws TException
     {
-        return factory.create(address, delegationToken);
+        return metastoreClientAdapterProvider.createThriftMetastoreClientAdapter(factory.create(address, delegationToken));
     }
 }