Index: service/src/test/org/apache/hive/service/cli/operation/TestCloseOperations.java =================================================================== --- service/src/test/org/apache/hive/service/cli/operation/TestCloseOperations.java (revision 0) +++ service/src/test/org/apache/hive/service/cli/operation/TestCloseOperations.java (revision 0) @@ -0,0 +1,68 @@ +/** + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import org.junit.Test; + +import org.apache.hive.service.cli.session.HiveSession; +import org.apache.hive.service.cli.session.HiveSessionImpl; +import org.apache.hive.service.cli.OperationHandle; +import org.apache.hive.service.cli.OperationState; +import org.apache.hive.service.cli.HiveSQLException; + +public class TestCloseOperations { + + @Test + public void checkOperationClosing() throws Exception { + OperationManager opMgr = new OperationManager(); + HiveSession session = new HiveSessionImpl("user", "passw", null); + + session.setOperationManager(opMgr); + + OperationHandle opHandle1 = session.executeStatement("use default", null); + assertNotNull(opHandle1); + assertEquals(OperationState.FINISHED, opMgr.getOperationState(opHandle1)); + + OperationHandle opHandle2 = session.executeStatement("show databases", null); + assertNotNull(opHandle2); + assertEquals(OperationState.FINISHED, opMgr.getOperationState(opHandle2)); + + opMgr.closeOperation(opHandle1); + try { + opMgr.getOperationState(opHandle1); + fail("Shouldn't work!"); + } catch(Exception e) { + assertTrue(e instanceof HiveSQLException); + } + + session.close(); // Close all remaining associated operations + + try { + opMgr.getOperationState(opHandle2); + fail("Shouldn't work now that it's fixed!"); + } catch(Exception e) { + assertTrue(e instanceof HiveSQLException); + } + } +} Index: service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java =================================================================== --- service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java (revision 1471214) +++ service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java (working copy) @@ -270,6 +270,7 @@ if (metastoreClient != null) { metastoreClient.close(); } + operationManager.closeOperations(sessionHandle); } finally { release(); } Index: service/src/java/org/apache/hive/service/cli/operation/OperationManager.java =================================================================== --- service/src/java/org/apache/hive/service/cli/operation/OperationManager.java (revision 1471214) +++ service/src/java/org/apache/hive/service/cli/operation/OperationManager.java (working copy) @@ -18,6 +18,7 @@ package org.apache.hive.service.cli.operation; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -29,6 +30,7 @@ import org.apache.hive.service.cli.OperationHandle; import org.apache.hive.service.cli.OperationState; import org.apache.hive.service.cli.RowSet; +import org.apache.hive.service.cli.SessionHandle; import org.apache.hive.service.cli.TableSchema; import org.apache.hive.service.cli.session.HiveSession; @@ -41,6 +43,8 @@ private HiveConf hiveConf; private final Map handleToOperation = new HashMap(); + private final Map operationToSession = new HashMap(); + private final Map> sessionToOpList = new HashMap>(); public OperationManager() { super("OperationManager"); @@ -69,26 +73,26 @@ String statement, Map confOverlay) { ExecuteStatementOperation executeStatementOperation = ExecuteStatementOperation .newExecuteStatementOperation(parentSession, statement, confOverlay); - addOperation(executeStatementOperation); + addOperation(executeStatementOperation, parentSession.getSessionHandle()); return executeStatementOperation; } public GetTypeInfoOperation newGetTypeInfoOperation(HiveSession parentSession) { GetTypeInfoOperation operation = new GetTypeInfoOperation(parentSession); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } public GetCatalogsOperation newGetCatalogsOperation(HiveSession parentSession) { GetCatalogsOperation operation = new GetCatalogsOperation(parentSession); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } public GetSchemasOperation newGetSchemasOperation(HiveSession parentSession, String catalogName, String schemaName) { GetSchemasOperation operation = new GetSchemasOperation(parentSession, catalogName, schemaName); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } @@ -97,13 +101,13 @@ List tableTypes) { MetadataOperation operation = new GetTablesOperation(parentSession, catalogName, schemaName, tableName, tableTypes); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } public GetTableTypesOperation newGetTableTypesOperation(HiveSession parentSession) { GetTableTypesOperation operation = new GetTableTypesOperation(parentSession); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } @@ -111,7 +115,7 @@ String catalogName, String schemaName, String tableName, String columnName) { GetColumnsOperation operation = new GetColumnsOperation(parentSession, catalogName, schemaName, tableName, columnName); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } @@ -119,7 +123,7 @@ String catalogName, String schemaName, String functionName) { GetFunctionsOperation operation = new GetFunctionsOperation(parentSession, catalogName, schemaName, functionName); - addOperation(operation); + addOperation(operation, parentSession.getSessionHandle()); return operation; } @@ -131,11 +135,23 @@ return operation; } - private synchronized void addOperation(Operation operation) { - handleToOperation.put(operation.getHandle(), operation); + private synchronized void addOperation(Operation operation, SessionHandle sessionHandle) { + OperationHandle opHandle = operation.getHandle(); + handleToOperation.put(opHandle, operation); + operationToSession.put(opHandle, sessionHandle); + if(sessionToOpList.containsKey(sessionHandle)) { + sessionToOpList.get(sessionHandle).add(opHandle); + } else { + ArrayList opList = new ArrayList(); + opList.add(opHandle); + sessionToOpList.put(sessionHandle, opList); + } } private synchronized Operation removeOperation(OperationHandle opHandle) { + // Take care of associated sessions + SessionHandle sessHandle = operationToSession.remove(opHandle); + sessionToOpList.get(sessHandle).remove(opHandle); return handleToOperation.remove(opHandle); } @@ -169,4 +185,15 @@ throws HiveSQLException { return getOperation(opHandle).getNextRowSet(orientation, maxRows); } + + public synchronized void closeOperations(SessionHandle sessionHandle) throws HiveSQLException { + ArrayList opList = sessionToOpList.get(sessionHandle); + if(opList != null) { + for(OperationHandle opHandle : opList) { + handleToOperation.remove(opHandle).close(); + operationToSession.remove(opHandle); + } + } + sessionToOpList.remove(sessionHandle); + } }