Bug fix in "DROP TABLE" when using S3 by resetting the tiledb context

TileDB-Inc · Oct 29, 2021 · c2c5ca2 · c2c5ca2
1 parent 89f2e2a
commit c2c5ca2
Showing 1 changed file with 29 additions and 12 deletions.
diff --git a/src/main/java/com/facebook/presto/plugin/tiledb/TileDBClient.java b/src/main/java/com/facebook/presto/plugin/tiledb/TileDBClient.java
@@ -68,13 +68,7 @@ public TileDBClient(TileDBConnectorId connectorId, TileDBConfig config)
         this.config = requireNonNull(config, "config is null");
         try {
             // Create context
-            Config tileDBConfig = new Config();
-            if (config.getAwsAccessKeyId() != null && !config.getAwsAccessKeyId().isEmpty()
-                    && config.getAwsSecretAccessKey() != null && !config.getAwsSecretAccessKey().isEmpty()) {
-                tileDBConfig.set("vfs.s3.aws_access_key_id", config.getAwsAccessKeyId());
-                tileDBConfig.set("vfs.s3.aws_secret_access_key", config.getAwsSecretAccessKey());
-            }
-            ctx = new Context(tileDBConfig);
+            ctx = new Context(buildConfig());
             //tileDBConfig.close();
         }
         catch (TileDBError tileDBError) {
@@ -110,11 +104,11 @@ public TileDBConfig getConfig()
     }
 
     /**
-    * Helper function to add a table to the catalog
-    *
-    * @param schema schema to add table to
-    * @param arrayUri uri of array/table
-    */
+     * Helper function to add a table to the catalog
+     *
+     * @param schema schema to add table to
+     * @param arrayUri uri of array/table
+     */
     public TileDBTable addTableFromURI(Context localCtx, String schema, URI arrayUri)
     {
         return addTableFromURI(localCtx, schema, arrayUri, null, null);
@@ -242,6 +236,8 @@ public void rollbackCreateTable(TileDBOutputTableHandle handle)
         try {
             TileDBObject.remove(ctx, handle.getURI());
             schemas.get(handle.getSchemaName()).remove(handle.getTableName());
+            ctx.close();
+            ctx = new Context(buildConfig());
         }
         catch (TileDBError tileDBError) {
             throw new PrestoException(TileDBErrorCode.TILEDB_DROP_TABLE_ERROR, tileDBError);
@@ -254,12 +250,33 @@ public void dropTable(ConnectorSession session, TileDBTableHandle handle)
             Context localCtx = buildContext(session);
             TileDBObject.remove(localCtx, handle.getURI());
             schemas.get(handle.getSchemaName()).remove(handle.getTableName());
+            ctx.close();
+            ctx = new Context(buildConfig());
         }
         catch (TileDBError tileDBError) {
             throw new PrestoException(TileDBErrorCode.TILEDB_DROP_TABLE_ERROR, tileDBError);
         }
     }
 
+    public Config buildConfig()
+    {
+        try {
+            // Create context
+            Config tileDBConfig = new Config();
+            if (config.getAwsAccessKeyId() != null && !config.getAwsAccessKeyId().isEmpty()
+                    && config.getAwsSecretAccessKey() != null && !config.getAwsSecretAccessKey().isEmpty()) {
+                tileDBConfig.set("vfs.s3.aws_access_key_id", config.getAwsAccessKeyId());
+                tileDBConfig.set("vfs.s3.aws_secret_access_key", config.getAwsSecretAccessKey());
+            }
+            return tileDBConfig;
+            //tileDBConfig.close();
+        }
+        catch (TileDBError tileDBError) {
+            // Print stacktrace, this produces an error client side saying "internal error"
+            throw new TrinoException(TILEDB_CONTEXT_ERROR, tileDBError);
+        }
+    }
+
     public Context buildContext(ConnectorSession session) throws TileDBError
     {
         Context localCtx = ctx;