IMPALA-6805: Show current database in Impala shell prompt Prompt format: [host:port] db_name>
Testing: - Added new shell tests - Ran end-to-end shell tests Change-Id: Ifb0ae58507321e426e5f0f16518671420974a3fc Reviewed-on: http://gerrit.cloudera.org:8080/9927 Reviewed-by: Fredy Wijaya <fwij...@cloudera.com> Reviewed-by: Michael Brown <mi...@cloudera.com> Tested-by: Impala Public Jenkins <impala-public-jenk...@cloudera.com> Reviewed-on: http://gerrit.cloudera.org:8080/10024 Reviewed-by: Philip Zeyliger <phi...@cloudera.com> Tested-by: Philip Zeyliger <phi...@cloudera.com> Project: http://git-wip-us.apache.org/repos/asf/impala/repo Commit: http://git-wip-us.apache.org/repos/asf/impala/commit/bbe53429 Tree: http://git-wip-us.apache.org/repos/asf/impala/tree/bbe53429 Diff: http://git-wip-us.apache.org/repos/asf/impala/diff/bbe53429 Branch: refs/heads/2.x Commit: bbe534298df1a23267c4109c69dc2645a0c9ecaa Parents: c2e4c43 Author: Fredy wijaya <fwij...@cloudera.com> Authored: Wed Apr 4 13:05:31 2018 -0700 Committer: Philip Zeyliger <phi...@cloudera.com> Committed: Thu Apr 12 02:59:17 2018 +0000 ---------------------------------------------------------------------- shell/impala_shell.py | 15 +++++++++-- tests/shell/test_shell_interactive.py | 42 ++++++++++++++++++++++++++---- 2 files changed, 50 insertions(+), 7 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/impala/blob/bbe53429/shell/impala_shell.py ---------------------------------------------------------------------- diff --git a/shell/impala_shell.py b/shell/impala_shell.py index f809e46..67391bf 100755 --- a/shell/impala_shell.py +++ b/shell/impala_shell.py @@ -104,6 +104,7 @@ class ImpalaShell(object, cmd.Cmd): # If not connected to an impalad, the server version is unknown. UNKNOWN_SERVER_VERSION = "Not Connected" + PROMPT_FORMAT = "[{host}:{port}] {db}> " DISCONNECTED_PROMPT = "[Not connected] > " UNKNOWN_WEBSERVER = "0.0.0.0" # Message to display in shell when cancelling a query @@ -768,7 +769,8 @@ class ImpalaShell(object, cmd.Cmd): if self.imp_client.connected: self._print_if_verbose('Connected to %s:%s' % self.impalad) self._print_if_verbose('Server version: %s' % self.server_version) - self.prompt = "[%s:%s] > " % self.impalad + self.prompt = ImpalaShell.PROMPT_FORMAT.format( + host=self.impalad[0], port=self.impalad[1], db=ImpalaShell.DEFAULT_DB) if self.refresh_after_connect: self.cmdqueue.append('invalidate metadata' + ImpalaShell.CMD_DELIM) print_to_stderr("Invalidating Metadata") @@ -1138,7 +1140,16 @@ class ImpalaShell(object, cmd.Cmd): """Executes a USE... query""" query = self._create_beeswax_query(args) if self._execute_stmt(query) is CmdStatus.SUCCESS: - self.current_db = args + self.current_db = args.strip('`').strip() + self.prompt = ImpalaShell.PROMPT_FORMAT.format(host=self.impalad[0], + port=self.impalad[1], + db=self.current_db) + elif args.strip('`') == self.current_db: + # args == current_db means -d option was passed but the "use [db]" operation failed. + # We need to set the current_db to None so that it does not show a database, which + # may not exist. + self.current_db = None + return CmdStatus.ERROR else: return CmdStatus.ERROR http://git-wip-us.apache.org/repos/asf/impala/blob/bbe53429/tests/shell/test_shell_interactive.py ---------------------------------------------------------------------- diff --git a/tests/shell/test_shell_interactive.py b/tests/shell/test_shell_interactive.py index deb75ea..5c3ee2c 100755 --- a/tests/shell/test_shell_interactive.py +++ b/tests/shell/test_shell_interactive.py @@ -58,11 +58,11 @@ class TestImpalaShellInteractive(object): def teardown_class(cls): restore_shell_history(cls.tempfile_name) - def _expect_with_cmd(self, proc, cmd, expectations=()): + def _expect_with_cmd(self, proc, cmd, expectations=(), db="default"): """Executes a command on the expect process instance and verifies a set of assertions defined by the expections.""" proc.sendline(cmd + ";") - proc.expect(":21000] >") + proc.expect(":21000] {db}>".format(db=db)) if not expectations: return for e in expectations: assert e in proc.before @@ -71,7 +71,7 @@ class TestImpalaShellInteractive(object): def test_local_shell_options(self): """Test that setting the local shell options works""" proc = pexpect.spawn(SHELL_CMD) - proc.expect(":21000] >") + proc.expect(":21000] default>") self._expect_with_cmd(proc, "set", ("LIVE_PROGRESS: False", "LIVE_SUMMARY: False")) self._expect_with_cmd(proc, "set live_progress=true") self._expect_with_cmd(proc, "set", ("LIVE_PROGRESS: True", "LIVE_SUMMARY: False")) @@ -179,11 +179,14 @@ class TestImpalaShellInteractive(object): assert get_num_open_sessions(initial_impala_service) == num_sessions_initial + 1, \ "Not connected to %s:21000" % hostname p.send_cmd("connect %s:21001" % hostname) + # Wait for a little while sleep(2) # The number of sessions on the target impalad should have been incremented. assert get_num_open_sessions(target_impala_service) == num_sessions_target + 1, \ "Not connected to %s:21001" % hostname + assert "[%s:21001] default>" % hostname in p.get_result().stdout + # The number of sessions on the initial impalad should have been decremented. assert get_num_open_sessions(initial_impala_service) == num_sessions_initial, \ "Connection to %s:21000 should have been closed" % hostname @@ -260,7 +263,7 @@ class TestImpalaShellInteractive(object): os.remove(SHELL_HISTORY_FILE) assert not os.path.exists(SHELL_HISTORY_FILE) child_proc = pexpect.spawn(SHELL_CMD) - child_proc.expect(":21000] >") + child_proc.expect(":21000] default>") self._expect_with_cmd(child_proc, "@1", ("Command index out of range")) self._expect_with_cmd(child_proc, "rerun -1", ("Command index out of range")) self._expect_with_cmd(child_proc, "select 'first_command'", ("first_command")) @@ -268,7 +271,7 @@ class TestImpalaShellInteractive(object): self._expect_with_cmd(child_proc, "@ -1", ("first_command")) self._expect_with_cmd(child_proc, "select 'second_command'", ("second_command")) child_proc.sendline('history;') - child_proc.expect(":21000] >") + child_proc.expect(":21000] default>") assert '[1]: select \'first_command\';' in child_proc.before; assert '[2]: select \'second_command\';' in child_proc.before; assert '[3]: history;' in child_proc.before; @@ -494,6 +497,35 @@ class TestImpalaShellInteractive(object): result = run_impala_shell_interactive(query) assert '| id |' in result.stdout + @pytest.mark.execute_serially + def test_shell_prompt(self): + proc = pexpect.spawn(SHELL_CMD) + proc.expect(":21000] default>") + self._expect_with_cmd(proc, "use foo", (), 'default') + self._expect_with_cmd(proc, "use functional", (), 'functional') + self._expect_with_cmd(proc, "use foo", (), 'functional') + self._expect_with_cmd(proc, 'use `tpch`', (), 'tpch') + self._expect_with_cmd(proc, 'use ` tpch `', (), 'tpch') + + proc = pexpect.spawn(SHELL_CMD, ['-d', 'functional']) + proc.expect(":21000] functional>") + self._expect_with_cmd(proc, "use foo", (), 'functional') + self._expect_with_cmd(proc, "use tpch", (), 'tpch') + self._expect_with_cmd(proc, "use foo", (), 'tpch') + + proc = pexpect.spawn(SHELL_CMD, ['-d', ' functional ']) + proc.expect(":21000] functional>") + + proc = pexpect.spawn(SHELL_CMD, ['-d', '` functional `']) + proc.expect(":21000] functional>") + + # Start an Impala shell with an invalid DB. + proc = pexpect.spawn(SHELL_CMD, ['-d', 'foo']) + proc.expect(":21000] default>") + self._expect_with_cmd(proc, "use foo", (), 'default') + self._expect_with_cmd(proc, "use functional", (), 'functional') + self._expect_with_cmd(proc, "use foo", (), 'functional') + def run_impala_shell_interactive(input_lines, shell_args=None): """Runs a command in the Impala shell interactively.""" # if argument "input_lines" is a string, makes it into a list