HIVE-29609: HPL/SQL leaving behind hive-staging directories (apache#6500)

Author: mdayakar

itests/hive-unit/src/test/java/org/apache/hive/beeline/TestHplSqlViaBeeLine.java

@@ -22,6 +22,7 @@
 
 import static org.apache.hive.beeline.TestBeeLineWithArgs.OutStream;
 import static org.apache.hive.beeline.TestBeeLineWithArgs.testCommandLineScript;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
 
 import java.io.File;
@@ -37,8 +38,12 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hive.UtilsForTest;
 import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.metastore.conf.MetastoreConf;
 import org.apache.hive.jdbc.miniHS2.MiniHS2;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -1480,6 +1485,30 @@ public void testERRORCODEForExecuteStatements() throws Throwable {
     testScriptFile(scriptText, args(), "First ERRORCODE: -1.*Second ERRORCODE: 0", OutStream.ERR);
   }
 
+  @Test
+  public void testHplSqlInsertRemovesStagingDirsUnderTable() throws Throwable {
+    String scriptText =
+        "DROP TABLE IF EXISTS result;\n" +
+            "CREATE TABLE result (s string);\n" +
+            "INSERT INTO result VALUES('Hello');\n" +
+            "execute 'INSERT INTO result VALUES(''World'')';\n" +
+            "SELECT * FROM result;";
+    testScriptFile(scriptText, args(), "Hello.*World");
+
+    HiveConf conf = miniHS2.getHiveConf();
+    String stagingDirPrefix = HiveConf.getVar(conf, HiveConf.ConfVars.STAGING_DIR);
+    Path wh = new Path(MetastoreConf.getVar(conf, MetastoreConf.ConfVars.WAREHOUSE));
+    FileSystem fs = wh.getFileSystem(conf);
+    Path tableDir = fs.makeQualified(new Path(wh, "result"));
+    FileStatus[] children = fs.listStatus(tableDir);
+    if (children != null) {
+      for (FileStatus child : children) {
+        assertFalse("Staging directory was not cleaned up: " + child.getPath(),
+            child.getPath().getName().startsWith(stagingDirPrefix));
+      }
+    }
+  }
+
   private static List<String> args() {
     return Arrays.asList("-d", BeeLine.BEELINE_DEFAULT_JDBC_DRIVER,
             "-u", miniHS2.getBaseJdbcURL() + ";mode=hplsql", "-n", USER_NAME);

service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java

@@ -403,11 +403,7 @@ private synchronized void cleanup(OperationState state) throws HiveSQLException
       }
     }
 
-    if (driver != null) {
-      driver.close();
-      driver.destroy();
-    }
-    driver = null;
+    closeDriver();
 
     SessionState ss = SessionState.get();
     if (ss == null) {
@@ -423,6 +419,14 @@ private synchronized void cleanup(OperationState state) throws HiveSQLException
     }
   }
 
+  private void closeDriver() {
+    if (driver != null) {
+      driver.close();
+      driver.destroy();
+    }
+    driver = null;
+  }
+
   @Override
   public void cancel(OperationState stateAfterCancel) throws HiveSQLException {
     String queryId = null;
@@ -439,7 +443,9 @@ public void cancel(OperationState stateAfterCancel) throws HiveSQLException {
 
   @Override
   public void close() throws HiveSQLException {
-    if (!embedded) {
+    if (embedded) {
+      closeDriver();
+    } else {
       cleanup(OperationState.CLOSED);
       cleanupOperationLog(0);
     }

service/src/test/org/apache/hive/service/cli/operation/TestSQLOperationDriverCleanup.java

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.hive.service.cli.operation;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.ql.IDriver;
+import org.apache.hadoop.hive.ql.session.SessionState;
+import org.apache.hive.service.cli.HandleIdentifier;
+import org.apache.hive.service.cli.OperationState;
+import org.apache.hive.service.cli.SessionHandle;
+import org.apache.hive.service.cli.session.HiveSession;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.InOrder;
+
+import java.lang.reflect.Field;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class TestSQLOperationDriverCleanup {
+
+  private HiveSession session;
+  private IDriver driver;
+
+  @Before
+  public void setUp() {
+    HiveConf conf = new HiveConf();
+    session = mock(HiveSession.class);
+    when(session.getHiveConf()).thenReturn(conf);
+    when(session.getSessionState()).thenReturn(mock(SessionState.class));
+    when(session.getUserName()).thenReturn("user");
+    SessionHandle sessionHandle = mock(SessionHandle.class);
+    when(sessionHandle.getHandleIdentifier()).thenReturn(new HandleIdentifier());
+    when(session.getSessionHandle()).thenReturn(sessionHandle);
+    driver = mock(IDriver.class);
+  }
+
+  @Test
+  public void testSQLOperationCloseReleasesDriverInHPLSQLMode() throws Exception {
+    SQLOperation operation = new SQLOperation(session, "insert into test values (1)",
+        ImmutableMap.of(), false, 0L, true);
+    setDriver(operation, driver);
+
+    operation.close();
+
+    InOrder order = inOrder(driver);
+    order.verify(driver).close();
+    order.verify(driver).destroy();
+    assertNull(getDriver(operation));
+  }
+
+  @Test
+  public void testSQLOperationCloseReleasesDriverInNonHPLSQLMode() throws Exception {
+    SQLOperation operation = new SQLOperation(session, "insert into test values (1)",
+        ImmutableMap.of(), false, 0L, false);
+    setDriver(operation, driver);
+
+    operation.close();
+
+    InOrder order = inOrder(driver);
+    order.verify(driver).close();
+    order.verify(driver).destroy();
+    assertEquals(OperationState.CLOSED, operation.getStatus().getState());
+    assertNull(getDriver(operation));
+  }
+
+  private static void setDriver(SQLOperation operation, IDriver driver) throws Exception {
+    Field driverField = SQLOperation.class.getDeclaredField("driver");
+    driverField.setAccessible(true);
+    driverField.set(operation, driver);
+  }
+
+  private static IDriver getDriver(SQLOperation operation) throws Exception {
+    Field driverField = SQLOperation.class.getDeclaredField("driver");
+    driverField.setAccessible(true);
+    return (IDriver) driverField.get(operation);
+  }
+}