From 7a8ada07ff67262e10b575de992288c483c9cd4e Mon Sep 17 00:00:00 2001
From: Sankar Hariappan <mailtosankarh@gmail.com>
Date: Thu, 9 Feb 2017 15:59:23 +0530
Subject: [PATCH] HIVE-15858: Beeline ^C doesn't close the session

---
 .../java/org/apache/hive/beeline/BeeLine.java | 39 +++++++++++-----
 .../apache/hive/beeline/TestShutdownHook.java | 46 +++++++++++++++++++
 2 files changed, 73 insertions(+), 12 deletions(-)
 create mode 100644 beeline/src/test/org/apache/hive/beeline/TestShutdownHook.java

diff --git a/beeline/src/java/org/apache/hive/beeline/BeeLine.java b/beeline/src/java/org/apache/hive/beeline/BeeLine.java
index abeff0290ebd..3c8fcccb77bf 100644
--- a/beeline/src/java/org/apache/hive/beeline/BeeLine.java
+++ b/beeline/src/java/org/apache/hive/beeline/BeeLine.java
@@ -124,6 +124,7 @@ public class BeeLine implements Closeable {
   private static final ResourceBundle resourceBundle =
       ResourceBundle.getBundle(BeeLine.class.getSimpleName());
   private final BeeLineSignalHandler signalHandler;
+  private final Runnable shutdownHook;
   private static final String separator = System.getProperty("line.separator");
   private boolean exit = false;
   private final DatabaseConnections connections = new DatabaseConnections();
@@ -532,13 +533,26 @@ public BeeLine() {
   public BeeLine(boolean isBeeLine) {
     this.isBeeLine = isBeeLine;
     this.signalHandler = new SunSignalHandler(this);
+    this.shutdownHook = new Runnable() {
+      @Override
+      public void run() {
+        try {
+          if (history != null) {
+            history.flush();
+          }
+        } catch (IOException e) {
+          error(e);
+        } finally {
+          close();
+        }
+      }
+    };
   }
 
   DatabaseConnection getDatabaseConnection() {
     return getDatabaseConnections().current();
   }
 
-
   Connection getConnection() throws SQLException {
     if (getDatabaseConnections().current() == null
         || getDatabaseConnections().current().getConnection() == null) {
@@ -985,6 +999,9 @@ public int begin(String[] args, InputStream inputStream) throws IOException {
 
     setupHistory();
 
+    //add shutdown hook to cleanup the beeline for smooth exit
+    addBeelineShutdownHook();
+
     //this method also initializes the consoleReader which is
     //needed by initArgs for certain execution paths
     ConsoleReader reader = initializeConsoleReader(inputStream);
@@ -1184,17 +1201,11 @@ private void setupHistory() throws IOException {
     }
 
     this.history = new FileHistory(new File(getOpts().getHistoryFile()));
-    // add shutdown hook to flush the history to history file
-    ShutdownHookManager.addShutdownHook(new Runnable() {
-      @Override
-      public void run() {
-        try {
-          history.flush();
-        } catch (IOException e) {
-          error(e);
-        }
-      }
-    });
+  }
+
+  private void addBeelineShutdownHook() throws IOException {
+    // add shutdown hook to flush the history to history file and it also close all open connections
+    ShutdownHookManager.addShutdownHook(getShutdownHook());
   }
 
   public ConsoleReader initializeConsoleReader(InputStream inputStream) throws IOException {
@@ -2272,6 +2283,10 @@ DatabaseConnections getDatabaseConnections() {
     return connections;
   }
 
+  Runnable getShutdownHook() {
+    return shutdownHook;
+  }
+
   Completer getCommandCompletor() {
     return beeLineCommandCompleter;
   }
diff --git a/beeline/src/test/org/apache/hive/beeline/TestShutdownHook.java b/beeline/src/test/org/apache/hive/beeline/TestShutdownHook.java
new file mode 100644
index 000000000000..9927ef93015d
--- /dev/null
+++ b/beeline/src/test/org/apache/hive/beeline/TestShutdownHook.java
@@ -0,0 +1,46 @@
+/**
+ * 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.beeline;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.PrintStream;
+import java.io.PrintWriter;
+import java.lang.reflect.Method;
+
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class TestShutdownHook {
+  @Test
+  public void testShutdownHook() throws Exception {
+    ByteArrayOutputStream os = new ByteArrayOutputStream();
+    PrintStream ops = new PrintStream(os);
+    BeeLine beeline = new BeeLine();
+    DatabaseConnections dbConnections = beeline.getDatabaseConnections();
+    dbConnections.setConnection(new DatabaseConnection(beeline,null,null, null));
+    dbConnections.setConnection(new DatabaseConnection(beeline,null,null, null));
+    Assert.assertEquals(2, dbConnections.size());
+    beeline.setOutputStream(ops);
+    beeline.getShutdownHook().run();
+    Assert.assertEquals(0, dbConnections.size());
+  }
+}
\ No newline at end of file