Class ImpalaShell

tests/shell/util.py:226–325

A single instance of the Impala shell. The process is started when this object is constructed, and then users should repeatedly call send_cmd(), followed eventually by get_result() to retrieve the process output. This constructor will wait until Impala shell is connected for the speci

class ImpalaShell(object):
  """A single instance of the Impala shell. The process is started when this object is
     constructed, and then users should repeatedly call send_cmd(), followed eventually by
     get_result() to retrieve the process output. This constructor will wait until
     Impala shell is connected for the specified timeout unless wait_until_connected is
     set to False or --quiet is passed into the args."""
  def __init__(self, vector, args=None, env=None, wait_until_connected=True, timeout=60,
               stdout_file=None, stderr_file=None):
    self.shell_process = self._start_new_shell_process(vector, args, env=env,
                                                       stdout_file=stdout_file,
                                                       stderr_file=stderr_file)
    # When --quiet option is passed to Impala shell, we should not wait until we see
    # "Connected to" because it will never be printed to stderr. The same is true
    # if stderr is redirected.
    if wait_until_connected and (args is None or "--quiet" not in args) and \
       stderr_file is None:
      # We don't want to hang waiting for input. So, here are the scenarios
      # we need to handle:
      # 1. Shell process exits
      # 2. Shell fails to connect. This can lead to an interactive prompt
      #    that blocks for input forever. The two messages to look for are
      #    "Error connecting" and "Socket error"
      # 3. Process successfully connecting: "Connected to"
      # Cases 1 and 2 should lead to an assert.
      start_time = time.time()
      process_status = None
      connection_err = None
      connected = False
      while time.time() - start_time < timeout:
        # Condition 1: check if the shell process has exited
        # poll() returns None until the process exits
        process_status = self.shell_process.poll()
        if process_status is not None:
          break
        # readline() can block forever, so the timeout logic may not be effective
        # if something gets stuck here.
        line = self.shell_process.stderr.readline()
        # Condition 2: check for errors connecting
        if ImpalaShellClass.ERROR_CONNECTING_MESSAGE in line or \
           ImpalaShellClass.SOCKET_ERROR_MESSAGE in line:
          connection_err = line
          break

        # Condition 3: check if the connection is successful
        connected = ImpalaShellClass.CONNECTED_TO_MESSAGE in line
        if connected:
          break

      assert process_status is None, \
          "Impala shell exited with return code {0}".format(process_status)
      assert connection_err is None, connection_err
      assert connected, "Impala shell is not connected"

  def pid(self):
    return self.shell_process.pid

  def send_cmd(self, cmd):
    """Send a single command to the shell. This method adds the end-of-query