#!/usr/bin/env impala-python
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import logging
from multiprocessing import Value
import os
import re
from textwrap import dedent
from time import sleep, time
from sys import maxint

from tests.stress.queries import QueryType
from tests.stress.util import create_and_start_daemon_thread, increment
from tests.util.thrift_util import op_handle_to_query_id

LOG = logging.getLogger(os.path.splitext(os.path.basename(__file__))[0])

# Metrics collected during the stress running process.
NUM_QUERIES_DEQUEUED = "num_queries_dequeued"
# The number of queries that were submitted to a query runner.
NUM_QUERIES_SUBMITTED = "num_queries_submitted"
# The number of queries that have entered the RUNNING state (i.e. got through Impala's
# admission control and started executing) or were cancelled or hit an error.
NUM_QUERIES_STARTED_RUNNING_OR_CANCELLED = "num_queries_started_running_or_cancelled"
NUM_QUERIES_FINISHED = "num_queries_finished"
NUM_QUERIES_EXCEEDED_MEM_LIMIT = "num_queries_exceeded_mem_limit"
NUM_QUERIES_AC_REJECTED = "num_queries_ac_rejected"
NUM_QUERIES_AC_TIMEDOUT = "num_queries_ac_timedout"
NUM_QUERIES_CANCELLED = "num_queries_cancelled"
NUM_RESULT_MISMATCHES = "num_result_mismatches"
NUM_OTHER_ERRORS = "num_other_errors"


class CancelMechanism:
  """Pseudo-enum to specify how a query should be cancelled."""
  # Cancel via the stress test client issuing an RPC to the server.
  VIA_CLIENT = "CANCEL_VIA_CLIENT"
  # Cancel by setting the EXEC_TIME_LIMIT_S query option.
  VIA_OPTION = "CANCEL_VIA_OPTION"

  ALL_MECHS = [VIA_CLIENT, VIA_OPTION]

RESULT_HASHES_DIR = "result_hashes"


class QueryTimeout(Exception):
  pass


class QueryRunner(object):
  """Encapsulates functionality to run a query and provide a runtime report."""

  SPILLED_PATTERNS = [re.compile("ExecOption:.*Spilled"), re.compile("SpilledRuns: [^0]")]
  BATCH_SIZE = 1024

  def __init__(self, impalad, results_dir, use_kerberos, common_query_options,
               test_admission_control, check_if_mem_was_spilled=False):
    """Creates a new instance, but does not start the process. """
    self.impalad = impalad
    self.use_kerberos = use_kerberos
    self.results_dir = results_dir
    self.check_if_mem_was_spilled = check_if_mem_was_spilled
    self.common_query_options = common_query_options
    self.test_admission_control = test_admission_control
    # proc is filled out by caller
    self.proc = None
    # impalad_conn is initialised in connect()
    self.impalad_conn = None

    # All these values are shared values between processes. We want these to be accessible
    # by the parent process that started this QueryRunner, for operational purposes.
    self._metrics = {
        NUM_QUERIES_DEQUEUED: Value("i", 0),
        NUM_QUERIES_SUBMITTED: Value("i", 0),
        NUM_QUERIES_STARTED_RUNNING_OR_CANCELLED: Value("i", 0),
        NUM_QUERIES_FINISHED: Value("i", 0),
        NUM_QUERIES_EXCEEDED_MEM_LIMIT: Value("i", 0),
        NUM_QUERIES_AC_REJECTED: Value("i", 0),
        NUM_QUERIES_AC_TIMEDOUT: Value("i", 0),
        NUM_QUERIES_CANCELLED: Value("i", 0),
        NUM_RESULT_MISMATCHES: Value("i", 0),
        NUM_OTHER_ERRORS: Value("i", 0)}

  def connect(self):
    """Connect to the server and start the query runner thread."""
    self.impalad_conn = self.impalad.impala.connect(impalad=self.impalad)

  def run_query(self, query, mem_limit_mb, run_set_up=False,
                timeout_secs=maxint, cancel_mech=None, retain_profile=False):
    """Run a query and return an execution report. If 'run_set_up' is True, set up sql
    will be executed before the main query. This should be the case during the binary
    search phase of the stress test. 'cancel_mech' is optionally a CancelMechanism
    value that should be used to cancel the query after timeout_secs.
    If 'cancel_mech' is provided, don't get the query profile for timed out queries
    because the query was purposely cancelled, rather than having some problem that needs
    investigation.
    """
    assert self.impalad_conn, "connect() must be called before run_query()"
    assert cancel_mech is None or cancel_mech in CancelMechanism.ALL_MECHS

    timeout_unix_time = time() + timeout_secs
    report = QueryReport(query)
    try:
      with self.impalad_conn.cursor() as cursor:
        start_time = time()
        self._set_db_and_options(cursor, query, run_set_up, mem_limit_mb, timeout_secs,
            cancel_mech)
        error = None
        try:
          cursor.execute_async(
              "/* Mem: %s MB. Coordinator: %s. */\n"
              % (mem_limit_mb, self.impalad.host_name) + query.sql)
          report.query_id = op_handle_to_query_id(cursor._last_operation.handle if
                                                  cursor._last_operation else None)
          LOG.debug("Query id is %s", report.query_id)
          self._wait_until_fetchable(cursor, report, timeout_unix_time)
          if query.query_type == QueryType.SELECT:
            report.result_hash = self._fetch_and_hash_result(cursor, timeout_unix_time,
                                                             query)
          else:
            # If query is in error state, this will raise an exception
            cursor._wait_to_finish()
          if (retain_profile or
             query.result_hash and report.result_hash != query.result_hash):
            fetch_and_set_profile(cursor, report)
        except QueryTimeout:
          if not cancel_mech:
            fetch_and_set_profile(cursor, report)
          # Cancel from this client if a) we hit the timeout unexpectedly or b) we're
          # deliberately cancelling the query via the client.
          if not cancel_mech or cancel_mech == CancelMechanism.VIA_CLIENT:
            self._cancel(cursor, report)
          else:
            # Wait until the query is cancelled by a different mechanism.
            self._wait_until_cancelled(cursor, report.query_id)
          report.timed_out = True
          return report
        except Exception as error:
          report.query_id = op_handle_to_query_id(cursor._last_operation.handle if
                                                  cursor._last_operation else None)
          LOG.debug("Error running query with id %s: %s", report.query_id, error)
          if (cancel_mech == CancelMechanism.VIA_OPTION and
                is_exec_time_limit_error(error)):
            # Timeout via query option counts as a timeout.
            report.timed_out = True
          else:
            self._check_for_memory_errors(report, cursor, error)
        if report.has_query_error():
          return report
        report.runtime_secs = time() - start_time
        if cursor.execution_failed() or self.check_if_mem_was_spilled:
          fetch_and_set_profile(cursor, report)
          report.mem_was_spilled = any([
              pattern.search(report.profile) is not None
              for pattern in QueryRunner.SPILLED_PATTERNS])
          # TODO: is this needed? Memory errors are generally caught by the try/except.
          report.not_enough_memory = "Memory limit exceeded" in report.profile
    except Exception as error:
      # A mem limit error would have been caught above, no need to check for that here.
      LOG.debug("Caught error", error)
      report.other_error = error
    return report

  def _set_db_and_options(self, cursor, query, run_set_up, mem_limit_mb, timeout_secs,
      cancel_mech):
    """Set up a new cursor for running a query by switching to the correct database and
    setting query options."""
    if query.db_name:
      LOG.debug("Using %s database", query.db_name)
      cursor.execute("USE %s" % query.db_name)
    if run_set_up and query.set_up_sql:
      LOG.debug("Running set up query:\n%s", query.set_up_sql)
      cursor.execute(query.set_up_sql)
    for query_option, value in self.common_query_options.iteritems():
      cursor.execute(
          "SET {query_option}={value}".format(query_option=query_option, value=value))
    for query_option, value in query.options.iteritems():
      cursor.execute(
          "SET {query_option}={value}".format(query_option=query_option, value=value))
    # Set a time limit if it is the expected method of cancellation, or as an additional
    # method of cancellation if the query runs for too long. This is useful if, e.g.
    # if the client is blocked in a fetch() call and can't initiate the cancellation.
    if not cancel_mech or cancel_mech == CancelMechanism.VIA_OPTION:
      # EXEC_TIME_LIMIT_S is an int32, so we need to cap the value.
      cursor.execute("SET EXEC_TIME_LIMIT_S={timeout_secs}".format(
          timeout_secs=min(timeout_secs, 2**31 - 1)))
    cursor.execute("SET ABORT_ON_ERROR=1")
    if self.test_admission_control:
      LOG.debug(
          "Running query without mem limit at %s with timeout secs %s:\n%s",
          self.impalad.host_name, timeout_secs, query.sql)
    else:
      LOG.debug("Setting mem limit to %s MB", mem_limit_mb)
      cursor.execute("SET MEM_LIMIT=%sM" % mem_limit_mb)
      LOG.debug(
          "Running query with %s MB mem limit at %s with timeout secs %s:\n%s",
          mem_limit_mb, self.impalad.host_name, timeout_secs, query.sql)

  def _wait_until_fetchable(self, cursor, report, timeout_unix_time):
    """Wait up until timeout_unix_time until the query results can be fetched (if it's
    a SELECT query) or until it has finished executing (if it's a different query type
    like DML). If the timeout expires raises a QueryTimeout exception."""
    # Loop until the query gets to the right state or a timeout expires.
    sleep_secs = 0.1
    secs_since_log = 0
    # True if we incremented num_queries_started_running_or_cancelled for this query.
    started_running_or_cancelled = False
    while True:
      query_state = cursor.status()
      # Check if the query got past the PENDING/INITIALIZED states, either because
      # it's executing or hit an error.
      if (not started_running_or_cancelled and query_state not in ('PENDING_STATE',
                                                      'INITIALIZED_STATE')):
        started_running_or_cancelled = True
        increment(self._metrics[NUM_QUERIES_STARTED_RUNNING_OR_CANCELLED])
      # Return if we're ready to fetch results (in the FINISHED state) or we are in
      # another terminal state like EXCEPTION.
      if query_state not in ('PENDING_STATE', 'INITIALIZED_STATE', 'RUNNING_STATE'):
        return

      if time() > timeout_unix_time:
        if not started_running_or_cancelled:
          increment(self._metrics[NUM_QUERIES_STARTED_RUNNING_OR_CANCELLED])
        raise QueryTimeout()
      if secs_since_log > 5:
        secs_since_log = 0
        LOG.debug("Waiting for query to execute")
      sleep(sleep_secs)
      secs_since_log += sleep_secs

  def _wait_until_cancelled(self, cursor, query_id):
    """Wait until 'cursor' is in a CANCELED or ERROR status."""
    last_log_time = 0
    status = cursor.status()
    while status not in ["CANCELED_STATE", "ERROR_STATE"]:
      # Work around IMPALA-7561: queries don't transition out of FINISHED state once they
      # hit eos.
      if status == "FINISHED_STATE" and "Last row fetched" in cursor.get_profile():
        return
      if time() - last_log_time > 5:
        LOG.debug("Waiting for query with id {query_id} to be cancelled".format(
                  query_id=query_id))
        last_log_time = time()
      sleep(0.1)
      status = cursor.status()

  def update_from_query_report(self, report):
    LOG.debug("Updating runtime stats (Query Runner PID: {0})".format(self.proc.pid))
    increment(self._metrics[NUM_QUERIES_FINISHED])
    if report.not_enough_memory:
      increment(self._metrics[NUM_QUERIES_EXCEEDED_MEM_LIMIT])
    if report.ac_rejected:
      increment(self._metrics[NUM_QUERIES_AC_REJECTED])
    if report.ac_timedout:
      increment(self._metrics[NUM_QUERIES_AC_TIMEDOUT])
    if report.was_cancelled:
      increment(self._metrics[NUM_QUERIES_CANCELLED])

  def _cancel(self, cursor, report):
    if not report.query_id:
      return

    try:
      LOG.debug("Attempting cancellation of query with id %s", report.query_id)
      cursor.cancel_operation()
      LOG.debug("Sent cancellation request for query with id %s", report.query_id)
    except Exception as e:
      LOG.debug("Error cancelling query with id %s: %s", report.query_id, e)
      try:
        LOG.debug("Attempting to cancel query through the web server.")
        self.impalad.cancel_query(report.query_id)
      except Exception as e:
        LOG.debug("Error cancelling query %s through the web server: %s",
                  report.query_id, e)

  def _check_for_memory_errors(self, report, cursor, caught_exception):
    """To be called after a query failure to check for signs of failed due to a
    mem limit or admission control rejection/timeout. The report will be updated
    accordingly.
    """
    fetch_and_set_profile(cursor, report)
    caught_msg = str(caught_exception).lower().strip()
    # Distinguish error conditions based on string fragments. The AC rejection and
    # out-of-memory conditions actually overlap (since some memory checks happen in
    # admission control) so check the out-of-memory conditions first.
    if "memory limit exceeded" in caught_msg or \
       "repartitioning did not reduce the size of a spilled partition" in caught_msg or \
       "failed to get minimum memory reservation" in caught_msg or \
       "minimum memory reservation is greater than" in caught_msg or \
       "minimum memory reservation needed is greater than" in caught_msg:
      report.not_enough_memory = True
      return
    if "rejected query from pool" in caught_msg:
      report.ac_rejected = True
      return
    if "admission for query exceeded timeout" in caught_msg:
      report.ac_timedout = True
      return

    LOG.debug("Non-mem limit error for query with id %s: %s", report.query_id,
              caught_exception, exc_info=True)
    report.other_error = caught_exception

  def _fetch_and_hash_result(self, cursor, timeout_unix_time, query):
    """Fetches results from 'cursor' and returns a hash that is independent of row order.
    Raises QueryTimeout() if we couldn't fetch all rows from the query before time()
    reaches 'timeout_unix_time'.
    'query' is only used for debug logging purposes (if the result is not as expected a
    log file will be left in RESULTS_DIR for investigation).
    """
    query_id = op_handle_to_query_id(cursor._last_operation.handle if
                                     cursor._last_operation else None)

    result_log_filename = None
    curr_hash = 1
    try:
      with self._open_result_file(query, query_id) as result_log:
        result_log.write(query.sql)
        result_log.write("\n")
        result_log_filename = result_log.name
        while True:
          # Check for timeout before each fetch call. Fetching can block indefinitely,
          # e.g. if the query is a selective scan, so we may miss the timeout, but we
          # will live with that limitation for now.
          # TODO: IMPALA-8289: use a non-blocking fetch when support is added
          if time() >= timeout_unix_time:
            LOG.debug("Hit timeout before fetch for query with id {0} {1} >= {2}".format(
                query_id, time(), timeout_unix_time))
            raise QueryTimeout()
          LOG.debug("Fetching result for query with id {0}".format(query_id))
          rows = cursor.fetchmany(self.BATCH_SIZE)
          if not rows:
            LOG.debug("No more results for query with id {0}".format(query_id))
            break
          for row in rows:
            curr_hash = _add_row_to_hash(row, curr_hash)
            if result_log:
              result_log.write(",".join([str(val) for val in row]))
              result_log.write("\thash=")
              result_log.write(str(curr_hash))
              result_log.write("\n")
    except Exception:
      # Don't retain result files for failed queries.
      if result_log_filename is not None:
        os.remove(result_log_filename)
      raise

    # Only retain result files with incorrect results.
    if result_log_filename is not None and curr_hash == query.result_hash:
      os.remove(result_log_filename)
    return curr_hash

  def _open_result_file(self, query, query_id):
    """Opens and returns a file under RESULT_HASHES_DIR to write query results to."""
    file_name = '_'.join([query.logical_query_id, query_id.replace(":", "_")])
    if query.result_hash is None:
      file_name += "_initial"
    file_name += "_results.txt"
    result_log = open(os.path.join(self.results_dir, RESULT_HASHES_DIR, file_name),
                      "w")
    return result_log

  def get_metric_val(self, name):
    """Get the current value of the metric called 'name'."""
    return self._metrics[name].value

  def get_metric_vals(self):
    """Get the current values of the all metrics as a list of (k, v) pairs."""
    return [(k, v.value) for k, v in self._metrics.iteritems()]

  def increment_metric(self, name):
    """Increment the current value of the metric called 'name'."""
    increment(self._metrics[name])


class QueryReport(object):
  """Holds information about a single query run."""

  def __init__(self, query):
    self.query = query

    self.result_hash = None
    self.runtime_secs = None
    self.mem_was_spilled = False
    # not_enough_memory includes conditions like "Memory limit exceeded", admission
    # control rejecting because not enough memory, etc.
    self.not_enough_memory = False
    # ac_rejected is true if the query was rejected by admission control.
    # It is mutually exclusive with not_enough_memory - if the query is rejected by
    # admission control because the memory limit is too low, it is counted as
    # not_enough_memory.
    # TODO: reconsider whether they should be mutually exclusive
    self.ac_rejected = False
    self.ac_timedout = False
    self.other_error = None
    self.timed_out = False
    self.was_cancelled = False
    self.profile = None
    self.query_id = None

  def __str__(self):
    return dedent("""
        <QueryReport
        result_hash: %(result_hash)s
        runtime_secs: %(runtime_secs)s
        mem_was_spilled: %(mem_was_spilled)s
        not_enough_memory: %(not_enough_memory)s
        ac_rejected: %(ac_rejected)s
        ac_timedout: %(ac_timedout)s
        other_error: %(other_error)s
        timed_out: %(timed_out)s
        was_cancelled: %(was_cancelled)s
        query_id: %(query_id)s
        >
        """.strip() % self.__dict__)

  def has_query_error(self):
    """Return true if any kind of error status was returned from the query (i.e.
    the query didn't run to completion, time out or get cancelled)."""
    return (self.not_enough_memory or self.ac_rejected or self.ac_timedout
            or self.other_error)

  def write_query_profile(self, directory, prefix=None):
    """
    Write out the query profile bound to this object to a given directory.

    The file name is generated and will contain the query ID. Use the optional prefix
    parameter to set a prefix on the filename.

    Example return:
      tpcds_300_decimal_parquet_q21_00000001_a38c8331_profile.txt

    Parameters:
      directory (str): Directory to write profile.
      prefix (str): Prefix for filename.
    """
    if not (self.profile and self.query_id):
      return
    if prefix is not None:
      file_name = prefix + '_'
    else:
      file_name = ''
    file_name += self.query.logical_query_id + '_'
    file_name += self.query_id.replace(":", "_") + "_profile.txt"
    profile_log_path = os.path.join(directory, file_name)
    with open(profile_log_path, "w") as profile_log:
      profile_log.write(self.profile)


def _add_row_to_hash(row, curr_hash):
  """Hashes row and accumulates it with 'curr_hash'. The hash is invariant of the
  order in which rows are added to the hash, since many queries we run through
  the stress test do not have a canonical order in which they return rows."""
  for idx, val in enumerate(row):
    curr_hash += _hash_val(idx, val)
    # Modulo the result to keep it "small" otherwise the math ops can be slow
    # since python does infinite precision math.
    curr_hash %= maxint
  return curr_hash


def _hash_val(col_idx, val):
  """Compute a numeric hash of 'val', which is at column index 'col_idx'."""
  if val is None:
    # The hash() of None can change from run to run since it's based on
    # a memory address. A chosen value will be used instead.
    val = 38463209
  elif isinstance(val, float):
    # Floats returned by Impala may not be deterministic, the ending
    # insignificant digits may differ. Only the first 6 digits will be used
    # after rounding.
    sval = "%f" % val
    dot_idx = sval.find(".")
    val = round(val, 6 - dot_idx)
  return (col_idx + 1) * hash(val)


def fetch_and_set_profile(cursor, report):
  """Set the report's query profile using the given cursor.
  Producing a query profile can be somewhat expensive. A v-tune profile of
  impalad showed 10% of cpu time spent generating query profiles.
  """
  if not report.profile and cursor._last_operation:
    try:
      report.profile = cursor.get_profile()
    except Exception as e:
      LOG.debug("Error getting profile for query with id %s: %s", report.query_id, e)


def is_exec_time_limit_error(error):
  """Return true if this is an error hit as a result of hitting EXEC_TIME_LIMIT_S."""
  return "expired due to execution time limit" in str(error)