def test_var_substitution(self):
    args = '--var=foo=123 --var=BAR=456 --delimited --output_delimiter=" " -c -f %s' \
           % (os.path.join(QUERY_FILE_PATH, 'test_var_substitution.sql'))
    result = run_impala_shell_cmd(args, expect_success=True)
    assert_var_substitution(result)
    args = '--var=foo'
    result = run_impala_shell_cmd(args, expect_success=False)
    assert ("Error: Could not parse key-value \"foo\". It must follow the pattern "
             "\"KEY=VALUE\".") in result.stderr

    # IMPALA-7673: Test that variable substitution in command line can accept values
    # from other variables just like the one in interactive shell.
    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:msg1}2" '
                                  '--var="msg3=${var:msg1}${var:msg2}" '
                                  '--query="select ${var:msg3}"')
    self._validate_shell_messages(result.stderr, ['112', 'Fetched 1 row(s)'],
                                  should_exist=True)

    # Test with an escaped variable.
    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:msg1}2" '
                                  '--var="msg3=\${var:msg1}${var:msg2}" '
                                  '--query="select \'${var:msg3}\'"')
    self._validate_shell_messages(result.stderr, ['${var:msg1}12', 'Fetched 1 row(s)'],
                                  should_exist=True)

    # Referencing a non-existent variable will result in an error.
    result = run_impala_shell_cmd('--var="msg1=1" --var="msg2=${var:doesnotexist}2" '
                                  '--var="msg3=\${var:msg1}${var:msg2}" '
                                  '--query="select \'${var:msg3}\'"',
                                  expect_success=False)
    self._validate_shell_messages(result.stderr,
                                  ['Error: Unknown variable DOESNOTEXIST',
                                   'Could not execute command: select \'${var:msg3}\''],
                                  should_exist=True)

  def test_large_sql(self, unique_database):
    # In this test, we are only interested in the performance of Impala shell and not
    # the performance of Impala in general. So, this test will execute a large query
    # from a non-existent table since this will make the query execution time negligible.
    sql_file, sql_path = tempfile.mkstemp()
    num_cols = 10000
    os.write(sql_file, "select \n")
    for i in xrange(num_cols):
      if i < num_cols:
        os.write(sql_file, "col_{0} as a{1},\n".format(i, i))
        os.write(sql_file, "col_{0} as b{1},\n".format(i, i))
        os.write(sql_file, "col_{0} as c{1}{2}\n".format(
            i, i, "," if i < num_cols - 1 else ""))
    os.write(sql_file, "from non_existence_large_table;")
    os.close(sql_file)

    try:
      args = "-q -f {0} -d {1}".format(sql_path, unique_database)
      start_time = time()
      run_impala_shell_cmd(args, False)
      end_time = time()
      time_limit_s = 10
      actual_time_s = end_time - start_time
      assert actual_time_s <= time_limit_s, (
          "It took {0} seconds to execute the query. Time limit is {1} seconds.".format(
              actual_time_s, time_limit_s))
    finally:
      os.remove(sql_path)

  def test_kerberos_option(self):
    args = "-k"

    # If you have a valid kerberos ticket in your cache, this test fails - so
    # here we set a bogus KRB5CCNAME in the environment so that klist (and other
    # kerberos commands) won't find the normal ticket cache.
    # KERBEROS TODO: add kerberized cluster test case
    os.environ["KRB5CCNAME"] = "/tmp/this/file/hopefully/does/not/exist"

    # The command will fail because we're trying to connect to a kerberized impalad.
    results = run_impala_shell_cmd(args, expect_success=False)
    # Check that impala is using the right service name.
    assert "Using service name 'impala'" in results.stderr
    assert "Starting Impala Shell using Kerberos authentication" in results.stderr
    # Check that Impala warns the user if klist does not exist on the system, or if
    # no kerberos tickets are initialized.
    try:
      call(["klist"])
      expected_error_msg = ("-k requires a valid kerberos ticket but no valid kerberos "
                            "ticket found.")
      assert expected_error_msg in results.stderr
    except OSError:
      assert 'klist not found on the system' in results.stderr
    # Make sure we don't try to re-connect
    assert "retrying the connection with a secure transport" not in results.stderr
    # Change the service name
    args += " -s foobar"
    results = run_impala_shell_cmd(args, expect_success=False)
    assert "Using service name 'foobar'" in results.stderr

 def test_execute_queries_from_file(self):
   args = '-f %s/test_file_comments.sql --quiet -B' % QUERY_FILE_PATH
   result = run_impala_shell_cmd(args)
   output = result.stdout
   args = '-f %s/test_file_no_comments.sql --quiet -B' % QUERY_FILE_PATH
   result = run_impala_shell_cmd(args)
   assert output == result.stdout, "Queries with comments not parsed correctly"

  def test_kudu_dml_reporting(self, vector, unique_database):
    db = unique_database
    run_impala_shell_cmd(vector, [
        '--query=create table %s.dml_test (id int primary key, '
        'age int null) partition by hash(id) partitions 2 stored as kudu' % db])

    self._validate_dml_stmt(
        vector, "insert into %s.dml_test (id) values (7), (7)" % db, 1, 1)
    self._validate_dml_stmt(vector, "insert into %s.dml_test (id) values (7)" % db, 0, 1)
    self._validate_dml_stmt(
        vector, "upsert into %s.dml_test (id) values (7), (7)" % db, 2, 0)
    self._validate_dml_stmt(
        vector, "update %s.dml_test set age = 1 where id = 7" % db, 1, 0)
    self._validate_dml_stmt(vector, "delete from %s.dml_test where id = 7" % db, 1, 0)

    # UPDATE/DELETE where there are no matching rows; there are no errors because the
    # scan produced no rows.
    self._validate_dml_stmt(
        vector, "update %s.dml_test set age = 1 where id = 8" % db, 0, 0)
    self._validate_dml_stmt(vector, "delete from %s.dml_test where id = 7" % db, 0, 0)

    # WITH clauses, only apply to INSERT and UPSERT
    self._validate_dml_stmt(vector,
        "with y as (values(7)) insert into %s.dml_test (id) select * from y" % db, 1, 0)
    self._validate_dml_stmt(vector,
        "with y as (values(7)) insert into %s.dml_test (id) select * from y" % db, 0, 1)
    self._validate_dml_stmt(vector,
        "with y as (values(7)) upsert into %s.dml_test (id) select * from y" % db, 1, 0)

 def test_malformed_query(self):
   """Test that malformed queries are handled by the commandline shell"""
   args = " -q \"with foo as (select bar from temp where temp.a='\""
   result = run_impala_shell_cmd(args, expect_success=False)
   assert "Encountered: EOF" in result.stderr
   args = "-q \"with v as (select 1) \;\""
   result = run_impala_shell_cmd(args, expect_success=False)
   assert "Encountered: Unexpected character" in result.stderr

 def test_ldap_password_from_shell(self):
   args = "-l --auth_creds_ok_in_clear --ldap_password_cmd='%s'"
   result = run_impala_shell_cmd(args % 'cmddoesntexist', expect_success=False)
   assert ("Error retrieving LDAP password (command was: 'cmddoesntexist', exception "
           "was: '[Errno 2] No such file or directory')") in result.stderr
   result = run_impala_shell_cmd(args % 'cat filedoesntexist', expect_success=False)
   assert ("Error retrieving LDAP password (command was 'cat filedoesntexist', error "
           "was: 'cat: filedoesntexist: No such file or directory')") in result.stderr

 def test_execute_queries_from_file(self, vector):
   args = ['-f', '{0}/test_file_comments.sql'.format(QUERY_FILE_PATH), '--quiet', '-B']
   result = run_impala_shell_cmd(vector, args)
   output = result.stdout
   args = ['-f', '{0}/test_file_no_comments.sql'.format(QUERY_FILE_PATH), '--quiet',
           '-B']
   result = run_impala_shell_cmd(vector, args)
   assert output == result.stdout, "Queries with comments not parsed correctly"

  def test_query_time_and_link_message(self, unique_database):
    shell_messages = ["Query submitted at: ", "(Coordinator: ",
        "Query progress can be monitored at: "]
    # CREATE statements should not print query time and webserver address.
    results = run_impala_shell_cmd('--query="create table %s.shell_msg_test (id int)"' %
        unique_database)
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=False)

    # SELECT, INSERT and CTAS queries should print the query time message and webserver
    # address.
    results = run_impala_shell_cmd('--query="insert into %s.shell_msg_test values (1)"' %
        unique_database)
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=True)
    results = run_impala_shell_cmd('--query="select * from %s.shell_msg_test"' %
        unique_database)
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=True)
    results = run_impala_shell_cmd('--query="create table %s.shell_msg_ctas_test as \
        select * from %s.shell_msg_test"' % (unique_database, unique_database))
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=True)

    # DROP statements should not print query time and webserver address.
    results = run_impala_shell_cmd('--query="drop table %s.shell_msg_test"' %
        unique_database)
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=False)
    run_impala_shell_cmd('--query="drop table %s.shell_msg_ctas_test"' %
        unique_database)

    # Simple queries should not print query time and webserver address.
    results = run_impala_shell_cmd('--query="use default"')
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=False)
    results = run_impala_shell_cmd('--query="show tables"')
    self._validate_shell_messages(results.stderr, shell_messages, should_exist=False)

  def test_config_file(self):
    """Test the optional configuration file."""
    # Positive tests
    args = '--config_file=%s/good_impalarc' % QUERY_FILE_PATH
    result = run_impala_shell_cmd(args)
    assert 'WARNING:' not in result.stderr
    assert 'Query: select 1' in result.stderr

    # override option in config file through command line
    args = '--config_file=%s/good_impalarc --query="select 2"' % QUERY_FILE_PATH
    result = run_impala_shell_cmd(args)
    assert 'Query: select 2' in result.stderr

    # IMPALA-8317: Add support for list-type, i.e. action=append in config file.
    args = '--config_file=%s/good_impalarc ' \
           '--query="select \'${VAR:msg1}\'; select \'${VAR:msg2}\'; ' \
           'select \'${VAR:msg3}\'; select \'${VAR:msg4}\'; set"' % \
           QUERY_FILE_PATH
    result = run_impala_shell_cmd(args)
    assert 'Query: select \'hello\'' in result.stderr
    assert 'Query: select \'world\'' in result.stderr
    assert 'Query: select \'foo\'' in result.stderr
    assert 'Query: select \'bar\'' in result.stderr
    assert 'DEFAULT_FILE_FORMAT: parquet' in result.stdout

    # Override the variables in the config file with the ones passed via --var.
    args = '--config_file=%s/good_impalarc --var=msg1=foo --var=msg2=bar ' \
           '--query="select \'${VAR:msg1}\'; select \'${VAR:msg2}\'"' % \
           QUERY_FILE_PATH
    result = run_impala_shell_cmd(args)
    assert 'Query: select \'foo\'' in result.stderr
    assert 'Query: select \'bar\'' in result.stderr

    # Negative Tests
    # specified config file does not exist
    args = '--config_file=%s/does_not_exist' % QUERY_FILE_PATH
    run_impala_shell_cmd(args, expect_success=False)
    # bad formatting of config file
    args = '--config_file=%s/bad_impalarc' % QUERY_FILE_PATH
    run_impala_shell_cmd(args, expect_success=False)

    # Testing config file related warning and error messages
    args = '--config_file=%s/impalarc_with_warnings' % QUERY_FILE_PATH
    result = run_impala_shell_cmd(args, expect_success=True, wait_until_connected=False)
    assert "WARNING: Option 'config_file' can be only set from shell." in result.stderr
    err_msg = ("WARNING: Unable to read configuration file correctly. "
               "Ignoring unrecognized config option: 'invalid_option'\n")
    assert  err_msg in result.stderr

    args = '--config_file=%s/impalarc_with_error' % QUERY_FILE_PATH
    result = run_impala_shell_cmd(args, expect_success=False)
    err_msg = ("Unexpected value in configuration file. "
               "'maybe' is not a valid value for a boolean option.")
    assert  err_msg in result.stderr

  def test_print_header(self, populated_table):
    args = '--print_header -B --output_delim="," -q "select * from %s"' % populated_table
    result = run_impala_shell_cmd(args)
    result_rows = result.stdout.strip().split('\n')
    assert len(result_rows) == 4
    assert result_rows[0].split(',') == ['i', 's']

    args = '-B --output_delim="," -q "select * from %s"' % populated_table
    result = run_impala_shell_cmd(args)
    result_rows = result.stdout.strip().split('\n')
    assert len(result_rows) == 3

 def test_output_format(self, vector):
   expected_output = ['1'] * 3
   args = ['-q', 'select 1,1,1', '-B', '--quiet']
   result = run_impala_shell_cmd(vector, args)
   actual_output = [r.strip() for r in result.stdout.split('\t')]
   assert actual_output == expected_output
   result = run_impala_shell_cmd(vector, args + ['--output_delim=|'])
   actual_output = [r.strip() for r in result.stdout.split('|')]
   assert actual_output == expected_output
   result = run_impala_shell_cmd(vector, args + ['--output_delim=||'],
                                 expect_success=False)
   assert "Illegal delimiter" in result.stderr

 def test_output_format(self):
   expected_output = ['1'] * 3
   args = '-q "select 1,1,1" -B --quiet'
   result = run_impala_shell_cmd(args)
   actual_output = [r.strip() for r in result.stdout.split('\t')]
   assert actual_output == expected_output
   result = run_impala_shell_cmd(args + ' --output_delim="|"')
   actual_output = [r.strip() for r in result.stdout.split('|')]
   assert actual_output == expected_output
   result = run_impala_shell_cmd(args + ' --output_delim="||"',
                                 expect_success=False)
   assert "Illegal delimiter" in result.stderr

  def test_print_header(self, vector, populated_table):
    args = ['--print_header', '-B', '--output_delim=,', '-q',
            'select * from {0}'.format(populated_table)]
    result = run_impala_shell_cmd(vector, args)
    result_rows = result.stdout.strip().split('\n')
    assert len(result_rows) == 4
    assert result_rows[0].split(',') == ['i', 's']

    args = ['-B', '--output_delim=,', '-q', 'select * from {0}'.format(populated_table)]
    result = run_impala_shell_cmd(vector, args)
    result_rows = result.stdout.strip().split('\n')
    assert len(result_rows) == 3

 def test_impala_shell_timeout(self):
   """Tests that impala shell times out during connect.
      This creates a random listening socket and we try to connect to this
      socket through the impala-shell. The impala-shell should timeout and not hang
      indefinitely while connecting
   """
   with closing(socket.socket()) as s:
     s.bind(("", 0))
     # maximum number of queued connections on this socket is 1.
     s.listen(1)
     test_port = s.getsockname()[1]
     args = '-q "select foo; select bar;" --ssl -t 2000 -i localhost:%d' % (test_port)
     run_impala_shell_cmd(args, expect_success=False)

  def test_allow_creds_in_clear(self):
    args = '-l'
    result = run_impala_shell_cmd(args, expect_success=False)
    err_msg = ("LDAP credentials may not be sent over insecure connections. "
               "Enable SSL or set --auth_creds_ok_in_clear")

    assert err_msg in result.stderr

 def test_completed_query_errors_2(self):
   args = ('-q "set abort_on_error=true;'
           ' select id, cnt from functional_parquet.bad_column_metadata t,'
           ' (select 1 cnt) u"')
   result = run_impala_shell_cmd(args, expect_success=False)
   assert 'ERROR: Column metadata states there are 11 values, ' in result.stderr
   assert 'but read 10 values from column id.' in result.stderr

 def test_international_characters_prettyprint(self):
   """IMPALA-2717: ensure we can handle international characters in pretty-printed
   output"""
   args = """-q "select '%s'" """ % RUSSIAN_CHARS
   result = run_impala_shell_cmd(args.encode('utf-8'))
   assert 'UnicodeDecodeError' not in result.stderr
   assert RUSSIAN_CHARS.encode('utf-8') in result.stdout

 def test_international_characters_prettyprint(self, vector):
   """IMPALA-2717: ensure we can handle international characters in pretty-printed
   output"""
   args = ['-q', "select '{0}'".format(RUSSIAN_CHARS.encode('utf-8'))]
   result = run_impala_shell_cmd(vector, args)
   assert 'UnicodeDecodeError' not in result.stderr
   assert RUSSIAN_CHARS.encode('utf-8') in result.stdout

 def test_international_characters(self):
   """Sanity test to ensure that the shell can read international characters."""
   russian_chars = (u"А, Б, В, Г, Д, Е, Ё, Ж, З, И, Й, К, Л, М, Н, О, П, Р,"
                    u"С, Т, У, Ф, Х, Ц,Ч, Ш, Щ, Ъ, Ы, Ь, Э, Ю, Я")
   args = """-B -q "select '%s'" """ % russian_chars
   result = run_impala_shell_cmd(args.encode('utf-8'))
   assert 'UnicodeDecodeError' not in result.stderr
   assert russian_chars.encode('utf-8') in result.stdout