Repository: incubator-impala
Updated Branches:
  refs/heads/master 87fc463e0 -> 6769220e2


IMPALA-2235: Fix current db when shell auto-reconnects

The ImpalaShell didn't issue the 'USE <current-db>' command after
reconnecting to the Impala daemon. Therefore the client session
used the default DB after reconnection, not the previously selected DB.

Setting the current DB is done by the _validate_database method.
Before this commit it appended the "use <db>" command to the
command queue of the Cmd class. But, at this point we might already
have commands in the command queue that will run before the
"use <db>" command. In case of reconnection, we want to invoke
the USE command right away.

Also, the command processed by the precmd() method can entirely skip
the command queue, therefore it is not enough to insert the USE
command to the front of the command queue. We need to issue the
USE command with the onecmd() method to execute it immediately.

I extended the _validate_database method with an "immediately" flag.
If this flag is true, _validate_database will use the onecmd() method.
Otherwise, it will append the USE command to the command queue to
maintain the previous behaviour.

I added a new automated test suite named test_shell_interactive_reconnect.py
to the "custom cluster" tests. It sets the default database, and after
reconnection it checks if the shell set it again automatically.

One test case checks if the shell set the DB after manually reconnecting
to the impala daemon by issuing the CONNECT command.
The other test case checks if the shell set the DB after automatic
reconnection due to cluster restart.

I needed to backup the impala shell history file because I didn't
want to pollute it by the test cases (just like the way it is done in
tests/shell/test_shell_interactive.py). I created utility functions for
this in tests/shell/util.py and now test_shell_interactive.py and
the newly created test suite are using these utility functions.

Change-Id: I40dfa00ba0314d356fe8617446f516505c925e5e
Reviewed-on: http://gerrit.cloudera.org:8080/8368
Reviewed-by: Tim Armstrong <tarmstr...@cloudera.com>
Tested-by: Impala Public Jenkins


Project: http://git-wip-us.apache.org/repos/asf/incubator-impala/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-impala/commit/6539e89c
Tree: http://git-wip-us.apache.org/repos/asf/incubator-impala/tree/6539e89c
Diff: http://git-wip-us.apache.org/repos/asf/incubator-impala/diff/6539e89c

Branch: refs/heads/master
Commit: 6539e89c81b070f2587d00334c71312e0cd87706
Parents: 87fc463
Author: Zoltan Borok-Nagy <borokna...@cloudera.com>
Authored: Fri Nov 3 15:54:58 2017 +0100
Committer: Impala Public Jenkins <impala-public-jenk...@gerrit.cloudera.org>
Committed: Wed Nov 15 22:42:22 2017 +0000

----------------------------------------------------------------------
 shell/impala_shell.py                           | 16 ++++-
 .../test_shell_interactive_reconnect.py         | 74 ++++++++++++++++++++
 tests/shell/test_shell_interactive.py           | 10 +--
 tests/shell/util.py                             | 13 ++++
 4 files changed, 106 insertions(+), 7 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-impala/blob/6539e89c/shell/impala_shell.py
----------------------------------------------------------------------
diff --git a/shell/impala_shell.py b/shell/impala_shell.py
index 641060f..069e875 100755
--- a/shell/impala_shell.py
+++ b/shell/impala_shell.py
@@ -487,6 +487,7 @@ class ImpalaShell(object, cmd.Cmd):
     except TException:
       print_to_stderr("Connection lost, reconnecting...")
       self._connect()
+      self._validate_database(immediately=True)
     return args.encode('utf-8')
 
   def onecmd(self, line):
@@ -730,10 +731,21 @@ class ImpalaShell(object, cmd.Cmd):
     self._connect()
     self._validate_database()
 
-  def _validate_database(self):
+  def _validate_database(self, immediately=False):
+    """ Issues a "USE <db>" command where <db> is the current database.
+    It is typically needed after the connection is (re-)established to the 
Impala daemon.
+
+    If immediately is False, it appends the USE command to self.cmdqueue.
+    If immediately is True, it executes the USE command right away.
+    """
     if self.current_db:
       self.current_db = self.current_db.strip('`')
-      self.cmdqueue.append(('use `%s`' % self.current_db) + 
ImpalaShell.CMD_DELIM)
+      use_current_db = ('use `%s`' % self.current_db)
+
+      if immediately:
+        self.onecmd(use_current_db)
+      else:
+        self.cmdqueue.append(use_current_db + ImpalaShell.CMD_DELIM)
 
   def _print_if_verbose(self, message):
     if self.verbose:

http://git-wip-us.apache.org/repos/asf/incubator-impala/blob/6539e89c/tests/custom_cluster/test_shell_interactive_reconnect.py
----------------------------------------------------------------------
diff --git a/tests/custom_cluster/test_shell_interactive_reconnect.py 
b/tests/custom_cluster/test_shell_interactive_reconnect.py
new file mode 100644
index 0000000..1f82468
--- /dev/null
+++ b/tests/custom_cluster/test_shell_interactive_reconnect.py
@@ -0,0 +1,74 @@
+# 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.
+
+import pytest
+import tempfile
+import socket
+
+from tests.common.custom_cluster_test_suite import CustomClusterTestSuite
+from tests.common.impala_service import ImpaladService
+from tests.common.skip import SkipIfBuildType
+from tests.shell.util import ImpalaShell, move_shell_history, 
restore_shell_history
+
+class TestShellInteractiveReconnect(CustomClusterTestSuite):
+  """ Check if interactive shell is using the current DB after reconnecting """
+  @classmethod
+  def get_workload(cls):
+    return 'functional-query'
+
+  @classmethod
+  def setup_class(cls):
+    super(TestShellInteractiveReconnect, cls).setup_class()
+
+    cls.tempfile_name = tempfile.mktemp()
+    move_shell_history(cls.tempfile_name)
+
+  @classmethod
+  def teardown_class(cls):
+    restore_shell_history(cls.tempfile_name)
+    super(TestShellInteractiveReconnect, cls).teardown_class()
+
+  @pytest.mark.execute_serially
+  def test_manual_reconnect(self):
+    p = ImpalaShell()
+    p.send_cmd("USE functional")
+    p.send_cmd("CONNECT")
+    p.send_cmd("SHOW TABLES")
+
+    result = p.get_result()
+    assert "alltypesaggmultifilesnopart" in result.stdout
+
+  @pytest.mark.execute_serially
+  def test_auto_reconnect(self):
+    NUM_QUERIES = 'impala-server.num-queries'
+
+    impalad = ImpaladService(socket.getfqdn())
+    start_num_queries = impalad.get_metric_value(NUM_QUERIES)
+
+    p = ImpalaShell()
+    p.send_cmd("USE functional")
+
+    # wait for the USE command to finish
+    impalad.wait_for_metric_value(NUM_QUERIES, start_num_queries + 1)
+    impalad.wait_for_num_in_flight_queries(0)
+
+    self._start_impala_cluster([])
+
+    p.send_cmd("SHOW TABLES")
+    result = p.get_result()
+    assert "alltypesaggmultifilesnopart" in result.stdout
+

http://git-wip-us.apache.org/repos/asf/incubator-impala/blob/6539e89c/tests/shell/test_shell_interactive.py
----------------------------------------------------------------------
diff --git a/tests/shell/test_shell_interactive.py 
b/tests/shell/test_shell_interactive.py
index 316e8c0..7c0d5ca 100755
--- a/tests/shell/test_shell_interactive.py
+++ b/tests/shell/test_shell_interactive.py
@@ -22,19 +22,19 @@ import os
 import pexpect
 import pytest
 import re
-import shutil
 import signal
 import socket
 import sys
+import tempfile
 from time import sleep
 
 from tests.common.impala_service import ImpaladService
 from tests.common.skip import SkipIfLocal
 from util import assert_var_substitution, ImpalaShell
+from util import move_shell_history, restore_shell_history
 
 SHELL_CMD = "%s/bin/impala-shell.sh" % os.environ['IMPALA_HOME']
 SHELL_HISTORY_FILE = os.path.expanduser("~/.impalahistory")
-TMP_HISTORY_FILE = os.path.expanduser("~/.impalahistorytmp")
 QUERY_FILE_PATH = os.path.join(os.environ['IMPALA_HOME'], 'tests', 'shell')
 
 class TestImpalaShellInteractive(object):
@@ -42,12 +42,12 @@ class TestImpalaShellInteractive(object):
 
   @classmethod
   def setup_class(cls):
-    if os.path.exists(SHELL_HISTORY_FILE):
-      shutil.move(SHELL_HISTORY_FILE, TMP_HISTORY_FILE)
+    cls.tempfile_name = tempfile.mktemp()
+    move_shell_history(cls.tempfile_name)
 
   @classmethod
   def teardown_class(cls):
-    if os.path.exists(TMP_HISTORY_FILE): shutil.move(TMP_HISTORY_FILE, 
SHELL_HISTORY_FILE)
+    restore_shell_history(cls.tempfile_name)
 
   def _expect_with_cmd(self, proc, cmd, expectations=()):
     """Executes a command on the expect process instance and verifies a set of

http://git-wip-us.apache.org/repos/asf/incubator-impala/blob/6539e89c/tests/shell/util.py
----------------------------------------------------------------------
diff --git a/tests/shell/util.py b/tests/shell/util.py
index 4507706..22ffa1a 100755
--- a/tests/shell/util.py
+++ b/tests/shell/util.py
@@ -22,12 +22,14 @@ import os
 import pytest
 import re
 import shlex
+import shutil
 from subprocess import Popen, PIPE
 
 IMPALAD_HOST_PORT_LIST = pytest.config.option.impalad.split(',')
 assert len(IMPALAD_HOST_PORT_LIST) > 0, 'Must specify at least 1 impalad to 
target'
 IMPALAD = IMPALAD_HOST_PORT_LIST[0]
 SHELL_CMD = "%s/bin/impala-shell.sh -i %s" % (os.environ['IMPALA_HOME'], 
IMPALAD)
+SHELL_HISTORY_FILE = os.path.expanduser("~/.impalahistory")
 
 def assert_var_substitution(result):
   assert_pattern(r'\bfoo_number=.*$', 'foo_number= 123123', result.stdout, \
@@ -110,6 +112,17 @@ def run_impala_shell_cmd_no_expect(shell_args, 
stdin_input=None):
   cmd = "%s %s" % (SHELL_CMD, shell_args)
   return result
 
+def move_shell_history(filepath):
+  """ Moves history file to given filepath.
+      If there is no history file, this function has no effect. """
+  if os.path.exists(SHELL_HISTORY_FILE):
+    shutil.move(SHELL_HISTORY_FILE, filepath)
+
+def restore_shell_history(filepath):
+  """ Moves back history file from given filepath.
+      If 'filepath' doesn't exist in the filesystem, this function has no 
effect. """
+  if os.path.exists(filepath): shutil.move(filepath, SHELL_HISTORY_FILE)
+
 class ImpalaShellResult(object):
   def __init__(self):
     self.rc = 0

Reply via email to