diff --git a/service/src/java/org/apache/hive/service/cli/operation/OperationManager.java b/service/src/java/org/apache/hive/service/cli/operation/OperationManager.java index 40515da01f..506ffe6721 100644 --- a/service/src/java/org/apache/hive/service/cli/operation/OperationManager.java +++ b/service/src/java/org/apache/hive/service/cli/operation/OperationManager.java @@ -194,7 +194,7 @@ private Operation getOperationInternal(OperationHandle operationHandle) { } private String getQueryId(Operation operation) { - return operation.getParentSession().getHiveConf().getVar(ConfVars.HIVEQUERYID); + return operation.getQueryId(); } private void addOperation(Operation operation) { diff --git a/service/src/test/org/apache/hive/service/cli/session/TestSessionCleanup.java b/service/src/test/org/apache/hive/service/cli/session/TestSessionCleanup.java index 83748c3eb1..51ce2c2426 100644 --- a/service/src/test/org/apache/hive/service/cli/session/TestSessionCleanup.java +++ b/service/src/test/org/apache/hive/service/cli/session/TestSessionCleanup.java @@ -29,6 +29,8 @@ import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.apache.hive.service.cli.CLIService; +import org.apache.hive.service.cli.OperationHandle; import org.apache.hive.service.cli.SessionHandle; import org.apache.hive.service.cli.thrift.EmbeddedThriftBinaryCLIService; import org.apache.hive.service.cli.thrift.ThriftCLIServiceClient; @@ -39,11 +41,22 @@ * TestSessionCleanup. */ public class TestSessionCleanup { + // Create subclass of EmbeddedThriftBinaryCLIService, just so we can get an accessor to the CLIService. + // Needed for access to the OperationManager. + private class MyEmbeddedThriftBinaryCLIService extends EmbeddedThriftBinaryCLIService { + public MyEmbeddedThriftBinaryCLIService() { + super(); + } + + public CLIService getCliService() { + return cliService; + } + } @Test // This is to test session temporary files are cleaned up after HIVE-11768 public void testTempSessionFileCleanup() throws Exception { - EmbeddedThriftBinaryCLIService service = new EmbeddedThriftBinaryCLIService(); + MyEmbeddedThriftBinaryCLIService service = new MyEmbeddedThriftBinaryCLIService(); HiveConf hiveConf = new HiveConf(); hiveConf .setVar(HiveConf.ConfVars.HIVE_AUTHORIZATION_MANAGER, @@ -54,7 +67,12 @@ public void testTempSessionFileCleanup() throws Exception { Set existingPipeoutFiles = new HashSet(Arrays.asList(getPipeoutFiles())); SessionHandle sessionHandle = client.openSession("user1", "foobar", Collections.emptyMap()); - client.executeStatement(sessionHandle, "set a=b", null); + OperationHandle opHandle1 = client.executeStatement(sessionHandle, "set a=b", null); + String queryId1 = service.getCliService().getQueryId(opHandle1.toTOperationHandle()); + Assert.assertNotNull(queryId1); + OperationHandle opHandle2 = client.executeStatement(sessionHandle, "set b=c", null); + String queryId2 = service.getCliService().getQueryId(opHandle2.toTOperationHandle()); + Assert.assertNotNull(queryId2); File operationLogRootDir = new File( new HiveConf().getVar(ConfVars.HIVE_SERVER2_LOGGING_OPERATION_LOG_LOCATION)); Assert.assertNotEquals(operationLogRootDir.list().length, 0); @@ -67,6 +85,13 @@ public void testTempSessionFileCleanup() throws Exception { Set finalPipeoutFiles = new HashSet(Arrays.asList(getPipeoutFiles())); finalPipeoutFiles.removeAll(existingPipeoutFiles); Assert.assertTrue(finalPipeoutFiles.isEmpty()); + + // Verify both operationHandles are no longer held by the OperationManager + Assert.assertEquals(0, service.getCliService().getSessionManager().getOperations().size()); + + // Verify both queryIds are no longer held by the OperationManager + Assert.assertNull(service.getCliService().getSessionManager().getOperationManager().getOperationByQueryId(queryId2)); + Assert.assertNull(service.getCliService().getSessionManager().getOperationManager().getOperationByQueryId(queryId1)); } private String[] getPipeoutFiles() {