Add unit tests for partition versioning enabled Metastore

prestodb · Sep 18, 2020 · 4b8021f · 4b8021f
1 parent 584fd8c
commit 4b8021f
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 3 deletions.
diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/MockHiveMetastore.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/MockHiveMetastore.java
@@ -0,0 +1,51 @@
+/*
+ * 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 com.facebook.presto.hive;
+
+import com.facebook.presto.common.predicate.Domain;
+import com.facebook.presto.hive.metastore.Column;
+import com.facebook.presto.hive.metastore.PartitionNameWithVersion;
+import com.facebook.presto.hive.metastore.thrift.ThriftHiveMetastore;
+import com.facebook.presto.spi.PrestoException;
+import org.apache.thrift.TException;
+
+import java.util.List;
+import java.util.Map;
+
+import static com.facebook.presto.hive.HiveErrorCode.HIVE_METASTORE_ERROR;
+import static com.facebook.presto.hive.metastore.TestCachingHiveMetastore.MockHiveCluster;
+import static java.util.Objects.requireNonNull;
+
+public class MockHiveMetastore
+        extends ThriftHiveMetastore
+{
+    private final MockHiveCluster clientProvider;
+
+    public MockHiveMetastore(MockHiveCluster mockHiveCluster)
+    {
+        super(mockHiveCluster);
+        this.clientProvider = requireNonNull(mockHiveCluster, "mockHiveCluster is null");
+    }
+
+    @Override
+    public List<PartitionNameWithVersion> getPartitionNamesWithVersionByFilter(String databaseName, String tableName, Map<Column, Domain> partitionPredicates)
+    {
+        try {
+            return clientProvider.createPartitionVersionSupportedMetastoreClient().getPartitionNamesWithVersionByFilter(databaseName, tableName, partitionPredicates);
+        }
+        catch (TException e) {
+            throw new PrestoException(HIVE_METASTORE_ERROR, e);
+        }
+    }
+}
diff --git a/...-metastore/src/test/java/com/facebook/presto/hive/metastore/TestCachingHiveMetastore.java b/...-metastore/src/test/java/com/facebook/presto/hive/metastore/TestCachingHiveMetastore.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive.metastore;
 
+import com.facebook.presto.hive.MockHiveMetastore;
 import com.facebook.presto.hive.PartitionVersionFetcher;
 import com.facebook.presto.hive.metastore.CachingHiveMetastore.MetastoreCacheScope;
 import com.facebook.presto.hive.metastore.thrift.BridgingHiveMetastore;
@@ -170,6 +171,58 @@ public void testGetPartitionNamesByParts()
         assertEquals(mockClient.getAccessCount(), 2);
     }
 
+    @Test
+    public void testCachingWithPartitionVersioning()
+    {
+        MockHiveMetastoreClient mockClient = new MockHiveMetastoreClient();
+        MockHiveCluster mockHiveCluster = new MockHiveCluster(mockClient);
+        ListeningExecutorService executor = listeningDecorator(newCachedThreadPool(daemonThreadsNamed("partition-versioning-test-%s")));
+        MockHiveMetastore mockHiveMetastore = new MockHiveMetastore(mockHiveCluster);
+        PartitionVersionFetcher hivePartitionVersionFetcher = new HivePartitionVersionFetcher();
+        boolean partitionVersioningEnabled = true;
+        CachingHiveMetastore partitionCachingEnabledmetastore = new CachingHiveMetastore(
+                new BridgingHiveMetastore(mockHiveMetastore, hivePartitionVersionFetcher),
+                executor,
+                new Duration(5, TimeUnit.MINUTES),
+                new Duration(1, TimeUnit.MINUTES),
+                1000,
+                partitionVersioningEnabled,
+                MetastoreCacheScope.PARTITION);
+
+        ImmutableList<String> expectedPartitions = ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2);
+
+        assertEquals(mockClient.getAccessCount(), 0);
+        assertEquals(partitionCachingEnabledmetastore.getPartitionNamesByFilter(TEST_DATABASE, TEST_TABLE, ImmutableMap.of()), expectedPartitions);
+        assertEquals(mockClient.getAccessCount(), 1);
+        assertEquals(partitionCachingEnabledmetastore.getPartitionNamesByFilter(TEST_DATABASE, TEST_TABLE, ImmutableMap.of()), expectedPartitions);
+        // Assert that we did not hit cache
+        assertEquals(mockClient.getAccessCount(), 2);
+
+        // Select all of the available partitions and load them into the cache
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2)).size(), 2);
+        assertEquals(mockClient.getAccessCount(), 3);
+
+        // Now if we fetch any or both of them, they should hit the cache
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION1)).size(), 1);
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION2)).size(), 1);
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2)).size(), 2);
+        assertEquals(mockClient.getAccessCount(), 3);
+
+        // This call should invalidate the partition cache because when partition versioning is enabled
+        // partitions with older or empty version will be deleted from cache. And the hivePartitionVersionFetcher used in creating
+        // partitionCachingEnabledmetastore will set the version as Optional.empty() for all the partitions fetched.
+        assertEquals(partitionCachingEnabledmetastore.getPartitionNamesByFilter(TEST_DATABASE, TEST_TABLE, ImmutableMap.of()), expectedPartitions);
+        assertEquals(mockClient.getAccessCount(), 4);
+
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2)).size(), 2);
+        // Assert that the previous cache entry is invalidated and that we did not hit the cache
+        assertEquals(mockClient.getAccessCount(), 5);
+
+        // Assert we hit the partition cache
+        assertEquals(partitionCachingEnabledmetastore.getPartitionsByNames(TEST_DATABASE, TEST_TABLE, ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2)).size(), 2);
+        assertEquals(mockClient.getAccessCount(), 5);
+    }
+
     public void testInvalidGetPartitionNamesByParts()
     {
         assertTrue(metastore.getPartitionNamesByFilter(BAD_DATABASE, TEST_TABLE, ImmutableMap.of()).isEmpty());
@@ -261,12 +314,12 @@ public void testNoCacheExceptions()
         assertEquals(mockClient.getAccessCount(), 2);
     }
 
-    private static class MockHiveCluster
+    public static class MockHiveCluster
             implements HiveCluster
     {
-        private final HiveMetastoreClient client;
+        private final MockHiveMetastoreClient client;
 
-        private MockHiveCluster(HiveMetastoreClient client)
+        private MockHiveCluster(MockHiveMetastoreClient client)
         {
             this.client = client;
         }
@@ -276,5 +329,10 @@ public HiveMetastoreClient createMetastoreClient()
         {
             return client;
         }
+
+        public MockHiveMetastoreClient createPartitionVersionSupportedMetastoreClient()
+        {
+            return client;
+        }
     }
 }
diff --git a/...tore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java b/...tore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java
@@ -13,6 +13,9 @@
  */
 package com.facebook.presto.hive.metastore.thrift;
 
+import com.facebook.presto.common.predicate.Domain;
+import com.facebook.presto.hive.metastore.Column;
+import com.facebook.presto.hive.metastore.PartitionNameWithVersion;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -53,6 +56,8 @@ public class MockHiveMetastoreClient
     public static final String TEST_PARTITION2 = "key=testpartition2";
     public static final List<String> TEST_PARTITION_VALUES1 = ImmutableList.of("testpartition1");
     public static final List<String> TEST_PARTITION_VALUES2 = ImmutableList.of("testpartition2");
+    public static final PartitionNameWithVersion TEST_PARTITION_NAME_WITH_VERSION1 = new PartitionNameWithVersion(TEST_PARTITION1, 1);
+    public static final PartitionNameWithVersion TEST_PARTITION_NAME_WITH_VERSION2 = new PartitionNameWithVersion(TEST_PARTITION2, 2);
     public static final List<String> TEST_ROLES = ImmutableList.of("testrole");
     public static final List<RolePrincipalGrant> TEST_ROLE_GRANTS = ImmutableList.of(
             new RolePrincipalGrant("role1", "user", USER, false, 0, "grantor1", USER),
@@ -213,6 +218,19 @@ public List<String> getPartitionNamesFiltered(String dbName, String tableName, L
         return ImmutableList.of(TEST_PARTITION1, TEST_PARTITION2);
     }
 
+    public List<PartitionNameWithVersion> getPartitionNamesWithVersionByFilter(String dbName, String tableName, Map<Column, Domain> partitionPredicates)
+            throws TException
+    {
+        accessCount.incrementAndGet();
+        if (throwException) {
+            throw new RuntimeException();
+        }
+        if (!dbName.equals(TEST_DATABASE) || !tableName.equals(TEST_TABLE)) {
+            throw new NoSuchObjectException();
+        }
+        return ImmutableList.of(TEST_PARTITION_NAME_WITH_VERSION1, TEST_PARTITION_NAME_WITH_VERSION2);
+    }
+
     @Override
     public Partition getPartition(String dbName, String tableName, List<String> partitionValues)
             throws TException