def test_reconnect(self):
    """Regression Test for IMPALA-1235

    Verifies that a connect command by the user is honoured.
    """

    def get_num_open_sessions(impala_service):
      """Helper method to retrieve the number of open sessions"""
      return impala_service.get_metric_value('impala-server.num-open-beeswax-sessions')

    hostname = socket.getfqdn()
    initial_impala_service = ImpaladService(hostname)
    target_impala_service = ImpaladService(hostname, webserver_port=25001,
        beeswax_port=21001, be_port=22001)
    # Get the initial state for the number of sessions.
    num_sessions_initial = get_num_open_sessions(initial_impala_service)
    num_sessions_target = get_num_open_sessions(target_impala_service)
    # Connect to localhost:21000 (default)
    p = ImpalaShell()
    sleep(2)
    # Make sure we're connected <hostname>:21000
    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
    # 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

 def test_malformed_query(self):
   """Test the handling of malformed query without closing quotation"""
   shell = ImpalaShell()
   query = "with v as (select 1) \nselect foo('\\\\'), ('bar \n;"
   shell.send_cmd(query)
   result = shell.get_result()
   assert "ERROR: ParseException: Unmatched string literal" in result.stderr

  def test_multiline_queries_in_history(self, tmp_history_file):
    """Test to ensure that multiline queries with comments are preserved in history

    Ensure that multiline queries are preserved when they're read back from history.
    Additionally, also test that comments are preserved.
    """
    # readline gets its input from tty, so using stdin does not work.
    child_proc = pexpect.spawn(SHELL_CMD)
    # List of (input query, expected text in output).
    # The expected output is usually the same as the input with a number prefix, except
    # where the shell strips newlines before a semicolon.
    queries = [
        ("select\n1;--comment", "[1]: select\n1;--comment"),
        ("select 1 --comment\n;", "[2]: select 1 --comment;"),
        ("select 1 --comment\n\n\n;", "[3]: select 1 --comment;"),
        ("select /*comment*/\n1;", "[4]: select /*comment*/\n1;"),
        ("select\n/*comm\nent*/\n1;", "[5]: select\n/*comm\nent*/\n1;")]
    for query, _ in queries:
      child_proc.expect(PROMPT_REGEX)
      child_proc.sendline(query)
      child_proc.expect("Fetched 1 row\(s\) in [0-9]+\.?[0-9]*s")
    child_proc.expect(PROMPT_REGEX)
    child_proc.sendline('quit;')
    p = ImpalaShell()
    p.send_cmd('history')
    result = p.get_result()
    for _, history_entry in queries:
      assert history_entry in result.stderr, "'%s' not in '%s'" % (history_entry,
                                                                   result.stderr)

 def test_change_delimiter(self):
   """Test change output delimiter if delimited mode is enabled"""
   p = ImpalaShell()
   p.send_cmd("use tpch")
   p.send_cmd("set write_delimited=true")
   p.send_cmd("set delimiter=,")
   p.send_cmd("select * from nation")
   result = p.get_result()
   assert "21,VIETNAM,2" in result.stdout

  def test_cancellation(self):
    """Test cancellation (Ctrl+C event)."""
    args = '-q "select sleep(10000)"'
    p = ImpalaShell(args)
    sleep(3)
    os.kill(p.pid(), signal.SIGINT)
    result = p.get_result()

    assert "Cancelling Query" in result.stderr, result.stderr

 def test_cancellation(self, vector):
   """Test cancellation (Ctrl+C event). Run a query that sleeps 10ms per row so will run
   for 110s if not cancelled, but will detect cancellation quickly because of the small
   batch size."""
   query = "set num_nodes=1; set mt_dop=1; set batch_size=1; \
            select sleep(10) from functional_parquet.alltypesagg"
   p = ImpalaShell(vector, ['-q', query])
   p.wait_for_query_start()
   os.kill(p.pid(), signal.SIGINT)
   result = p.get_result()
   assert "Cancelling Query" in result.stderr, result.stderr

 def test_with_clause(self):
   # IMPALA-7939: Fix issue where CTE that contains "insert", "upsert", "update", or
   # "delete" is categorized as a DML statement.
   for keyword in ["insert", "upsert", "update", "delete", "\\'insert\\'",
                   "\\'upsert\\'", "\\'update\\'", "\\'delete\\'"]:
     p = ImpalaShell()
     p.send_cmd("with foo as "
                "(select * from functional.alltypestiny where string_col='%s') "
                "select * from foo limit 1" % keyword)
     result = p.get_result()
     assert "Fetched 0 row" in result.stderr

 def test_cancellation(self):
   impalad = ImpaladService(socket.getfqdn())
   impalad.wait_for_num_in_flight_queries(0)
   command = "select sleep(10000);"
   p = ImpalaShell()
   p.send_cmd(command)
   sleep(3)
   os.kill(p.pid(), signal.SIGINT)
   result = p.get_result()
   assert "Cancelled" not in result.stderr
   assert impalad.wait_for_num_in_flight_queries(0)

  def run_and_verify_query_cancellation_test(self, vector, stmt, cancel_at_state):
    """Starts the execution of the received query, waits until the query
    execution in fact starts and then cancels it. Expects the query
    cancellation to succeed."""
    p = ImpalaShell(vector, ['-q', stmt])

    self.wait_for_query_state(vector, stmt, cancel_at_state)

    os.kill(p.pid(), signal.SIGINT)
    result = p.get_result()
    assert "Cancelling Query" in result.stderr
    assert "Invalid query handle" not in result.stderr

  def test_queries_closed(self, vector):
    """Regression test for IMPALA-897."""
    args = ['-f', '{0}/test_close_queries.sql'.format(QUERY_FILE_PATH), '--quiet', '-B']
    # Execute the shell command async
    p = ImpalaShell(vector, args)

    impalad_service = ImpaladService(get_impalad_host_port(vector).split(':')[0])
    # The last query in the test SQL script will sleep for 10 seconds, so sleep
    # here for 5 seconds and verify the number of in-flight queries is 1.
    sleep(5)
    assert 1 == impalad_service.get_num_in_flight_queries()
    assert p.get_result().rc == 0
    assert 0 == impalad_service.get_num_in_flight_queries()

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
  if type(input_lines) is str:
    input_lines = [input_lines]
  # workaround to make Popen environment 'utf-8' compatible
  # since piping defaults to ascii
  my_env = os.environ
  my_env['PYTHONIOENCODING'] = 'utf-8'
  p = ImpalaShell(shell_args, env=my_env)
  for line in input_lines:
    p.send_cmd(line)
  return p.get_result()

  def test_queries_closed(self):
    """Regression test for IMPALA-897."""
    args = '-f %s/test_close_queries.sql --quiet -B' % QUERY_FILE_PATH
    cmd = "%s %s" % (SHELL_CMD, args)
    # Execute the shell command async
    p = ImpalaShell(args)

    impalad_service = ImpaladService(IMPALAD.split(':')[0])
    # The last query in the test SQL script will sleep for 10 seconds, so sleep
    # here for 5 seconds and verify the number of in-flight queries is 1.
    sleep(5)
    assert 1 == impalad_service.get_num_in_flight_queries()
    assert p.get_result().rc == 0
    assert 0 == impalad_service.get_num_in_flight_queries()

 def test_write_delimited(self):
   """Test output rows in delimited mode"""
   p = ImpalaShell()
   p.send_cmd("use tpch")
   p.send_cmd("set write_delimited=true")
   p.send_cmd("select * from nation")
   result = p.get_result()
   assert "+----------------+" not in result.stdout
   assert "21\tVIETNAM\t2" in result.stdout

  def test_reconnect(self, vector):
    """Regression Test for IMPALA-1235

    Verifies that a connect command by the user is honoured.
    """
    # Disconnect existing clients so there are no open sessions.
    self.client.close()
    self.hs2_client.close()

    def wait_for_num_open_sessions(impala_service, num, err):
      """Helper method to wait for the number of open sessions to reach 'num'."""
      metric_name = get_open_sessions_metric(vector)
      assert impala_service.wait_for_metric_value(metric_name, num) == num, err

    hostname = socket.getfqdn()
    initial_impala_service = ImpaladService(hostname)
    target_impala_service = ImpaladService(hostname, webserver_port=25001,
        beeswax_port=21001, be_port=22001, hs2_port=21051)
    if vector.get_value("protocol") == "hs2":
      target_port = 21051
    else:
      target_port = 21001
    # This test is running serially, so there shouldn't be any open sessions, but wait
    # here in case a session from a previous test hasn't been fully closed yet.
    wait_for_num_open_sessions(initial_impala_service, 0,
        "first impalad should not have any remaining open sessions.")
    wait_for_num_open_sessions(target_impala_service, 0,
        "second impalad should not have any remaining open sessions.")
    # Connect to the first impalad
    p = ImpalaShell(vector)

    # Make sure we're connected <hostname>:<port>
    wait_for_num_open_sessions(initial_impala_service, 1,
        "Not connected to %s:%d" % (hostname, get_impalad_port(vector)))
    p.send_cmd("connect %s:%d" % (hostname, target_port))

    # The number of sessions on the target impalad should have been incremented.
    wait_for_num_open_sessions(
        target_impala_service, 1, "Not connected to %s:%d" % (hostname, target_port))
    assert "[%s:%d] default>" % (hostname, target_port) in p.get_result().stdout

    # The number of sessions on the initial impalad should have been decremented.
    wait_for_num_open_sessions(initial_impala_service, 0,
        "Connection to %s:%d should have been closed" % (
          hostname, get_impalad_port(vector)))

  def test_multiline_queries_in_history(self):
    """Test to ensure that multiline queries with comments are preserved in history

    Ensure that multiline queries are preserved when they're read back from history.
    Additionally, also test that comments are preserved.
    """
    # regex for pexpect, a shell prompt is expected after each command..
    prompt_regex = '.*%s:2100.*' % socket.getfqdn()
    # readline gets its input from tty, so using stdin does not work.
    child_proc = pexpect.spawn(SHELL_CMD)
    queries = ["select\n1--comment;",
        "select /*comment*/\n1;",
        "select\n/*comm\nent*/\n1;"]
    for query in queries:
      child_proc.expect(prompt_regex)
      child_proc.sendline(query)
    child_proc.expect(prompt_regex)
    child_proc.sendline('quit;')
    p = ImpalaShell()
    p.send_cmd('history')
    result = p.get_result()
    for query in queries:
      assert query in result.stderr, "'%s' not in '%s'" % (query, result.stderr)

  def test_reconnect(self):
    """Regression Test for IMPALA-1235

    Verifies that a connect command by the user is honoured.
    """

    def wait_for_num_open_sessions(impala_service, num, err):
      """Helper method to wait for the number of open sessions to reach 'num'."""
      assert impala_service.wait_for_metric_value(
          'impala-server.num-open-beeswax-sessions', num) == num, err

    hostname = socket.getfqdn()
    initial_impala_service = ImpaladService(hostname)
    target_impala_service = ImpaladService(hostname, webserver_port=25001,
        beeswax_port=21001, be_port=22001)
    # This test is running serially, so there shouldn't be any open sessions, but wait
    # here in case a session from a previous test hasn't been fully closed yet.
    wait_for_num_open_sessions(
        initial_impala_service, 0, "21000 should not have any remaining open sessions.")
    wait_for_num_open_sessions(
        target_impala_service, 0, "21001 should not have any remaining open sessions.")
    # Connect to localhost:21000 (default)
    p = ImpalaShell()

    # Make sure we're connected <hostname>:21000
    wait_for_num_open_sessions(
        initial_impala_service, 1, "Not connected to %s:21000" % hostname)
    p.send_cmd("connect %s:21001" % hostname)

    # The number of sessions on the target impalad should have been incremented.
    wait_for_num_open_sessions(
        target_impala_service, 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.
    wait_for_num_open_sessions(initial_impala_service, 0,
        "Connection to %s:21000 should have been closed" % hostname)

  def test_timezone_validation(self):
    """Test that query option TIMEZONE is validated when executing a query.

       Query options are not sent to the coordinator immediately, so the error checking
       will only happen when running a query.
    """
    p = ImpalaShell()
    p.send_cmd('set timezone=BLA;')
    p.send_cmd('select 1;')
    results = p.get_result()
    assert "Fetched 1 row" not in results.stderr
    assert "ERROR: Errors parsing query options" in results.stderr
    assert "Invalid timezone name 'BLA'" in results.stderr

  def test_set_and_set_all(self):
    """IMPALA-2181. Tests the outputs of SET and SET ALL commands. SET should contain the
    REGULAR and ADVANCED options only. SET ALL should contain all the options grouped by
    display level."""
    shell1 = ImpalaShell()
    shell1.send_cmd("set")
    result = shell1.get_result()
    assert "Query options (defaults shown in []):" in result.stdout
    assert "ABORT_ON_ERROR" in result.stdout
    assert "Advanced Query Options:" in result.stdout
    assert "APPX_COUNT_DISTINCT" in result.stdout
    assert "SUPPORT_START_OVER" in result.stdout
    # Development, deprecated and removed options should not be shown.
    # Note: there are currently no deprecated options
    assert "Development Query Options:" not in result.stdout
    assert "DEBUG_ACTION" not in result.stdout  # Development option.
    assert "MAX_IO_BUFFERS" not in result.stdout  # Removed option.

    shell2 = ImpalaShell()
    shell2.send_cmd("set all")
    result = shell2.get_result()
    assert "Query options (defaults shown in []):" in result.stdout
    assert "Advanced Query Options:" in result.stdout
    assert "Development Query Options:" in result.stdout
    assert "Deprecated Query Options:" not in result.stdout
    advanced_part_start_idx = result.stdout.find("Advanced Query Options")
    development_part_start_idx = result.stdout.find("Development Query Options")
    deprecated_part_start_idx = result.stdout.find("Deprecated Query Options")
    advanced_part = result.stdout[advanced_part_start_idx:development_part_start_idx]
    development_part = result.stdout[development_part_start_idx:deprecated_part_start_idx]
    assert "ABORT_ON_ERROR" in result.stdout[:advanced_part_start_idx]
    assert "APPX_COUNT_DISTINCT" in advanced_part
    assert "SUPPORT_START_OVER" in advanced_part
    assert "DEBUG_ACTION" in development_part
    # Removed options should not be shown.
    assert "MAX_IO_BUFFERS" not in result.stdout

 def test_compute_stats_with_live_progress_options(self):
   """Test that setting LIVE_PROGRESS options won't cause COMPUTE STATS query fail"""
   impalad = ImpaladService(socket.getfqdn())
   p = ImpalaShell()
   p.send_cmd("set live_progress=True")
   p.send_cmd("set live_summary=True")
   p.send_cmd('create table test_live_progress_option(col int);')
   try:
     p.send_cmd('compute stats test_live_progress_option;')
   finally:
     p.send_cmd('drop table if exists test_live_progress_option;')
   result = p.get_result()
   assert "Updated 1 partition(s) and 1 column(s)" in result.stdout

  def test_ddl_queries_are_closed(self):
    """Regression test for IMPALA-1317

    The shell does not call close() for alter, use and drop queries, leaving them in
    flight. This test issues those queries in interactive mode, and checks the debug
    webpage to confirm that they've been closed.
    TODO: Add every statement type.
    """

    TMP_DB = 'inflight_test_db'
    TMP_TBL = 'tmp_tbl'
    MSG = '%s query should be closed'
    NUM_QUERIES = 'impala-server.num-queries'

    impalad = ImpaladService(socket.getfqdn())
    p = ImpalaShell()
    try:
      start_num_queries = impalad.get_metric_value(NUM_QUERIES)
      p.send_cmd('create database if not exists %s' % TMP_DB)
      p.send_cmd('use %s' % TMP_DB)
      impalad.wait_for_metric_value(NUM_QUERIES, start_num_queries + 2)
      assert impalad.wait_for_num_in_flight_queries(0), MSG % 'use'
      p.send_cmd('create table %s(i int)' % TMP_TBL)
      p.send_cmd('alter table %s add columns (j int)' % TMP_TBL)
      impalad.wait_for_metric_value(NUM_QUERIES, start_num_queries + 4)
      assert impalad.wait_for_num_in_flight_queries(0), MSG % 'alter'
      p.send_cmd('drop table %s' % TMP_TBL)
      impalad.wait_for_metric_value(NUM_QUERIES, start_num_queries + 5)
      assert impalad.wait_for_num_in_flight_queries(0), MSG % 'drop'
    finally:
      run_impala_shell_interactive("drop table if exists %s.%s;" % (TMP_DB, TMP_TBL))
      run_impala_shell_interactive("drop database if exists foo;")