Ignore internal TableNotFoundException in DROP TABLE

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java

@@ -103,6 +103,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
@@ -1062,13 +1063,27 @@ public void createTable(Table table)
     @Override
     public void dropTable(String databaseName, String tableName, boolean deleteData)
     {
+        AtomicInteger attemptCount = new AtomicInteger();
         try {
             retry()
                     .stopOn(NoSuchObjectException.class)
                     .stopOnIllegalExceptions()
                     .run("dropTable", stats.getDropTable().wrap(() -> {
                         try (ThriftMetastoreClient client = createMetastoreClient()) {
-                            Table table = client.getTable(databaseName, tableName);
+                            attemptCount.incrementAndGet();
+                            Table table;
+                            try {
+                                table = client.getTable(databaseName, tableName);
+                            }
+                            catch (NoSuchObjectException e) {
+                                if (attemptCount.get() == 1) {
+                                    // Throw exception only on first attempt.
+                                    throw e;
+                                }
+                                // If table is not found on consecutive attempts it was probably dropped on first attempt and timeout occurred.
+                                // Exception in such case can be safely ignored and dropping table is finished.
+                                return null;
+                            }
                             client.dropTable(databaseName, tableName, deleteData);
                             String tableLocation = table.getSd().getLocation();
                             if (deleteFilesOnDrop && deleteData && isManagedTable(table) && !isNullOrEmpty(tableLocation)) {

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");

plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/AbstractTestHiveMetastore.java

@@ -15,6 +15,7 @@
 
 import io.trino.plugin.hive.SchemaAlreadyExistsException;
 import io.trino.plugin.hive.TableAlreadyExistsException;
+import io.trino.spi.connector.TableNotFoundException;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.parallel.Execution;
@@ -30,6 +31,7 @@
 import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE;
 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;
@@ -45,6 +47,11 @@ public void setMetastore(HiveMetastore metastore)
         this.metastore = metastore;
     }
 
+    protected HiveMetastore getMetastore()
+    {
+        return metastore;
+    }
+
     @Test
     void testCreateDatabase()
     {
@@ -97,4 +104,20 @@ void testCreateTable()
         metastore.dropTable(databaseName, tableName, false);
         metastore.dropDatabase(databaseName, false);
     }
+
+    @Test
+    public void testDropNotExistingTable()
+    {
+        String databaseName = "test_database_" + randomNameSuffix();
+        Database.Builder database = Database.builder()
+                .setDatabaseName(databaseName)
+                .setOwnerName(Optional.empty())
+                .setOwnerType(Optional.empty());
+        getMetastore().createDatabase(database.build());
+
+        assertThatThrownBy(() -> getMetastore().dropTable(databaseName, "not_existing", false))
+                .isInstanceOf(TableNotFoundException.class);
+
+        getMetastore().dropDatabase(databaseName, false);
+    }
 }

plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java

@@ -13,14 +13,32 @@
  */
 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.assertThat;
+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 +54,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") || method.getName().equals("dropTable")) {
+                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 +79,80 @@ 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);
+    }
+
+    @Test
+    public void testDropTableWithRetries()
+    {
+        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)
+                .setTableType(EXTERNAL_TABLE.name())
+                .setOwner(Optional.empty());
+        table.getStorageBuilder()
+                .setStorageFormat(fromHiveStorageFormat(PARQUET));
+        getMetastore().createTable(table.build(), NO_PRIVILEGES);
+
+        assertThat(getMetastore().getTable(databaseName, tableName)).isPresent();
+        getMetastore().dropTable(databaseName, tableName, false);
+        assertThat(getMetastore().getTable(databaseName, tableName)).isEmpty();
+
+        getMetastore().dropDatabase(databaseName, false);
+    }
 }

plugin/trino-hive/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);
+}

plugin/trino-hive/src/test/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));
     }
 }