Add test for recursive caching directory listings

Author: pettyjamesm

plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/BaseCachingDirectoryListerTest.java

@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.trino.plugin.hive.HiveQueryRunner;
+import io.trino.plugin.hive.metastore.PrincipalPrivileges;
 import io.trino.plugin.hive.metastore.Table;
 import io.trino.plugin.hive.metastore.file.FileHiveMetastore;
 import io.trino.testing.AbstractTestQueryFramework;
@@ -25,7 +26,9 @@
 
 import java.nio.file.Path;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 
 import static com.google.common.io.MoreFiles.deleteRecursively;
 import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE;
@@ -44,13 +47,18 @@ public abstract class BaseCachingDirectoryListerTest<C extends DirectoryLister>
     @Override
     protected QueryRunner createQueryRunner()
             throws Exception
+    {
+        return createQueryRunner(ImmutableMap.of("hive.allow-register-partition-procedure", "true"));
+    }
+
+    protected QueryRunner createQueryRunner(Map<String, String> properties)
+            throws Exception
     {
         Path temporaryMetastoreDirectory = createTempDirectory(null);
         closeAfterClass(() -> deleteRecursively(temporaryMetastoreDirectory, ALLOW_INSECURE));
         directoryLister = createDirectoryLister();
         return HiveQueryRunner.builder()
-                .setHiveProperties(ImmutableMap.of(
-                        "hive.allow-register-partition-procedure", "true"))
+                .setHiveProperties(properties)
                 .setMetastore(distributedQueryRunner -> fileHiveMetastore = createTestingFileHiveMetastore(temporaryMetastoreDirectory.toFile()))
                 .setDirectoryLister(directoryLister)
                 .build();
@@ -324,17 +332,32 @@ public void testDropPartitionedTable()
         assertThat(isCached(tableGroup3PartitionLocation)).isFalse();
     }
 
-    private org.apache.hadoop.fs.Path getTableLocation(String schemaName, String tableName)
+    protected Optional<Table> getTable(String schemaName, String tableName)
+    {
+        return fileHiveMetastore.getTable(schemaName, tableName);
+    }
+
+    protected void createTable(Table table, PrincipalPrivileges principalPrivileges)
+    {
+        fileHiveMetastore.createTable(table, principalPrivileges);
+    }
+
+    protected void dropTable(String schemaName, String tableName, boolean deleteData)
+    {
+        fileHiveMetastore.dropTable(schemaName, tableName, deleteData);
+    }
+
+    protected org.apache.hadoop.fs.Path getTableLocation(String schemaName, String tableName)
     {
-        return fileHiveMetastore.getTable(schemaName, tableName)
+        return getTable(schemaName, tableName)
                 .map(table -> table.getStorage().getLocation())
                 .map(tableLocation -> new org.apache.hadoop.fs.Path(tableLocation))
                 .orElseThrow(() -> new NoSuchElementException(format("The table %s.%s could not be found", schemaName, tableName)));
     }
 
-    private org.apache.hadoop.fs.Path getPartitionLocation(String schemaName, String tableName, List<String> partitionValues)
+    protected org.apache.hadoop.fs.Path getPartitionLocation(String schemaName, String tableName, List<String> partitionValues)
     {
-        Table table = fileHiveMetastore.getTable(schemaName, tableName)
+        Table table = getTable(schemaName, tableName)
                 .orElseThrow(() -> new NoSuchElementException(format("The table %s.%s could not be found", schemaName, tableName)));
 
         return fileHiveMetastore.getPartition(table, partitionValues)
@@ -343,7 +366,7 @@ private org.apache.hadoop.fs.Path getPartitionLocation(String schemaName, String
                 .orElseThrow(() -> new NoSuchElementException(format("The partition %s from the table %s.%s could not be found", partitionValues, schemaName, tableName)));
     }
 
-    private boolean isCached(org.apache.hadoop.fs.Path path)
+    protected boolean isCached(org.apache.hadoop.fs.Path path)
     {
         return isCached(directoryLister, path);
     }

plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryListerRecursiveFilesOnly.java

@@ -0,0 +1,95 @@
+/*
+ * 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.fs;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import io.airlift.units.Duration;
+import io.trino.plugin.hive.metastore.MetastoreUtil;
+import io.trino.plugin.hive.metastore.Table;
+import io.trino.testing.QueryRunner;
+import org.apache.hadoop.fs.Path;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.NoSuchElementException;
+
+import static io.trino.plugin.hive.HiveQueryRunner.TPCH_SCHEMA;
+import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES;
+import static java.lang.String.format;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+// some tests may invalidate the whole cache affecting therefore other concurrent tests
+@Test(singleThreaded = true)
+public class TestCachingDirectoryListerRecursiveFilesOnly
+        extends BaseCachingDirectoryListerTest<CachingDirectoryLister>
+{
+    @Override
+    protected CachingDirectoryLister createDirectoryLister()
+    {
+        return new CachingDirectoryLister(Duration.valueOf("5m"), 1_000_000L, List.of("tpch.*"));
+    }
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return createQueryRunner(ImmutableMap.of(
+                "hive.allow-register-partition-procedure", "true",
+                "hive.recursive-directories", "true"));
+    }
+
+    @Override
+    protected boolean isCached(CachingDirectoryLister directoryLister, Path path)
+    {
+        return directoryLister.isCached(new DirectoryListingCacheKey(path, true));
+    }
+
+    @Test
+    public void testRecursiveDirectories()
+    {
+        // Create partitioned table to force files to be inserted in sub-partition paths
+        assertUpdate("CREATE TABLE recursive_directories (clicks bigint, day date, country varchar) WITH (format = 'ORC', partitioned_by = ARRAY['day', 'country'])");
+        assertUpdate("INSERT INTO recursive_directories VALUES (1000, DATE '2022-02-01', 'US'), (2000, DATE '2022-02-01', 'US'), (4000, DATE '2022-02-02', 'US'), (1500, DATE '2022-02-01', 'AT'), (2500, DATE '2022-02-02', 'AT')", 5);
+
+        // Replace the partitioned table a new unpartitioned table with the same root location, leaving the data in place
+        Table partitionedTable = getTable(TPCH_SCHEMA, "recursive_directories")
+                .orElseThrow(() -> new NoSuchElementException(format("Failed to read table %s.%s", TPCH_SCHEMA, "recursive_directories")));
+        // Must not delete the data files when dropping the partitioned table
+        dropTable(TPCH_SCHEMA, "recursive_directories", false);
+        // Must create the table directly to bypass check that the target directory already exists
+        Table testTable = Table.builder(partitionedTable)
+                .setPartitionColumns(ImmutableList.of())
+                .build();
+        createTable(testTable, testTable.getOwner().map(MetastoreUtil::buildInitialPrivilegeSet).orElse(NO_PRIVILEGES));
+
+        // Execute a query on the new table to pull the listing into the cache
+        assertQuery("SELECT sum(clicks) FROM recursive_directories", "VALUES (11000)");
+
+        Path tableLocation = getTableLocation(TPCH_SCHEMA, "recursive_directories");
+        assertTrue(isCached(tableLocation));
+
+        // Insert should invalidate cache, even at the root directory path
+        assertUpdate("INSERT INTO recursive_directories VALUES (1000)", 1);
+        assertFalse(isCached(tableLocation));
+
+        // Results should include the new insert which is at the table location root for the unpartitioned table
+        assertQuery("SELECT sum(clicks) FROM recursive_directories", "VALUES (12000)");
+
+        assertUpdate("DROP TABLE recursive_directories");
+
+        assertFalse(isCached(tableLocation));
+    }
+}