Class QueryToKill

tests/util/cancel_util.py:33–74

class QueryToKill:
  def __init__(self, test_suite, protocol, check_on_exit=True, user=None, nth_impalad=0):
    self.client = test_suite.create_client_for_nth_impalad(nth_impalad, protocol)
    self.sql = 'SELECT sleep(1000)'
    self.check_on_exit = check_on_exit
    self.user = user

  def poll(self):
    while True:
      try:
        results = self.client.fetch(self.sql, self.handle)
        if len(results.data) > 0:
          raise Exception("Failed to kill query within time limit.")
      except Exception as e:
        self.exc = e
        return

  def __enter__(self):
    self.handle = self.client.execute_async(self.sql, user=self.user)
    self.poll_thread = threading.Thread(target=lambda: self.poll())
    self.poll_thread.start()
    return self.client.handle_id(self.handle)

  def __exit__(self, exc_type, exc_value, traceback):  # noqa: U100
    self.poll_thread.join()
    if not self.check_on_exit:
      self.client.close()
      return
    # If ImpalaServer::UnregisterQuery() happens before the last polling, the error
    # message will be "Invalid or unknown query handle". Otherwise, the error message
    # will be "Cancelled".
    assert error_msg_startswith(
        str(self.exc),
        ["Invalid or unknown query handle", "Cancelled"],
        self.client.handle_id(self.handle),
    )
    try:
      self.client.fetch(self.sql, self.handle)
    except Exception as ex:
      assert "Invalid or unknown query handle" in str(ex) or "Cancelled" in str(ex)
    finally:
      self.client.close()


def assert_kill_ok(client, query_id, user=None):