Remove data for drop table with InMemoryHiveMetastore

prestodb · Aug 5, 2015 · b334442 · b334442
1 parent be5bfe3
commit b334442
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 8 deletions.
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveBenchmarkQueryRunner.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveBenchmarkQueryRunner.java
@@ -65,9 +65,9 @@ public static LocalQueryRunner createLocalQueryRunner(File tempDir)
         localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(nodeManager, 1), ImmutableMap.<String, String>of());
 
         // add hive
-        InMemoryHiveMetastore metastore = new InMemoryHiveMetastore();
+        File hiveDir = new File(tempDir, "hive_data");
-        File tpchDataDir = new File(tempDir, "tpch");
+        InMemoryHiveMetastore metastore = new InMemoryHiveMetastore(hiveDir);
-        tpchDataDir.mkdir();
+        File tpchDataDir = new File(hiveDir, "tpch");
         metastore.createDatabase(new Database("tpch", null, tpchDataDir.toURI().toString(), null));
 
         HiveConnectorFactory hiveConnectorFactory = new HiveConnectorFactory(

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java
@@ -63,8 +63,8 @@ public static QueryRunner createQueryRunner(Iterable<TpchTable<?>> tables)
             queryRunner.installPlugin(new SampledTpchPlugin());
             queryRunner.createCatalog("tpch_sampled", "tpch_sampled");
 
-            File baseDir = queryRunner.getCoordinator().getBaseDataDir().toFile();
+            File baseDir = new File(queryRunner.getCoordinator().getBaseDataDir().toFile(), "hive_data");
-            InMemoryHiveMetastore metastore = new InMemoryHiveMetastore();
+            InMemoryHiveMetastore metastore = new InMemoryHiveMetastore(baseDir);
             metastore.createDatabase(new Database("tpch", null, new File(baseDir, "tpch").toURI().toString(), null));
             metastore.createDatabase(new Database("tpch_sampled", null, new File(baseDir, "tpch_sampled").toURI().toString(), null));
 

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/metastore/InMemoryHiveMetastore.java b/presto-hive/src/test/java/com/facebook/presto/hive/metastore/InMemoryHiveMetastore.java
@@ -30,7 +30,9 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
+import static io.airlift.testing.FileUtils.deleteRecursively;
 
 public class InMemoryHiveMetastore
         implements HiveMetastore
@@ -39,12 +41,23 @@ public class InMemoryHiveMetastore
     private final ConcurrentHashMap<SchemaTableName, Table> relations = new ConcurrentHashMap<>();
     private final ConcurrentHashMap<SchemaTableName, Table> views = new ConcurrentHashMap<>();
 
+    private final File baseDirectory;
+
+    public InMemoryHiveMetastore(File baseDirectory)
+    {
+        this.baseDirectory = checkNotNull(baseDirectory, "baseDirectory is null");
+        checkArgument(!baseDirectory.exists(), "Base directory already exists");
+        checkArgument(baseDirectory.mkdirs(), "Could not create base directory");
+    }
+
     public void createDatabase(Database database)
     {
         checkNotNull(database, "database is null");
 
-        File file = new File(URI.create(database.getLocationUri()));
+        File directory = new File(URI.create(database.getLocationUri()));
-        file.mkdirs();
+        checkArgument(!directory.exists(), "Database directory already exists");
+        checkArgument(isParentDir(directory, baseDirectory), "Database directory must be inside of the metastore base directory");
+        checkArgument(directory.mkdirs(), "Could not create database directory");
 
         if (databases.putIfAbsent(database.getName(), database) != null) {
             throw new IllegalArgumentException("Database " + database.getName() + " already exists");
@@ -65,6 +78,10 @@ public void createTable(Table table)
         if (tableCopy.getSd() == null) {
             tableCopy.setSd(new StorageDescriptor());
         }
+        else if (tableCopy.getSd().getLocation() != null) {
+            File directory = new File(URI.create(tableCopy.getSd().getLocation()));
+            checkArgument(directory.exists(), "Table directory does not exist");
+        }
 
         if (relations.putIfAbsent(schemaTableName, tableCopy) != null) {
             throw new TableAlreadyExistsException(schemaTableName);
@@ -79,10 +96,19 @@ public void createTable(Table table)
     public void dropTable(String databaseName, String tableName)
     {
         SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName);
-        if (relations.remove(schemaTableName) == null) {
+        Table table = relations.remove(schemaTableName);
+        if (table == null) {
             throw new TableNotFoundException(schemaTableName);
         }
         views.remove(schemaTableName);
+
+        // remove data
+        String location = table.getSd().getLocation();
+        if (location != null) {
+            File directory = new File(URI.create(location));
+            checkArgument(isParentDir(directory, baseDirectory), "Table directory must be inside of the metastore base directory");
+            deleteRecursively(directory);
+        }
     }
 
     @Override
@@ -161,4 +187,14 @@ public Optional<Table> getTable(String databaseName, String tableName)
     public void flushCache()
     {
     }
+
+    private static boolean isParentDir(File directory, File baseDirectory)
+    {
+        for (File parent = directory.getParentFile(); parent != null; parent = parent.getParentFile()) {
+            if (parent.equals(baseDirectory)) {
+                return true;
+            }
+        }
+        return false;
+    }
 }