#!/usr/bin/env python
# Copyright (c) 2012 Cloudera, Inc. All rights reserved.
# Validates limit on scan nodes
#
import logging
import pytest
from copy import copy
from tests.common.impala_test_suite import ImpalaTestSuite
from tests.common.test_vector import *
from tests.common.test_dimensions import create_exec_option_dimension
from tests.common.test_dimensions import create_uncompressed_text_dimension

# Tests very deep expression trees and expressions with many children. Impala defines
# a 'safe' upper bound on the expr depth and the number of expr children in the
# FE Expr.java and any changes to those limits should be reflected in this test.
# The expr limits primarily guard against stack overflows or similar problems
# causing crashes. Therefore, this tests succeeds if no Impalads crash.
class TestExprLimits(ImpalaTestSuite):
  # Keep these in sync with Expr.java
  EXPR_CHILDREN_LIMIT = 10000;
  EXPR_DEPTH_LIMIT = 1500;

  @classmethod
  def get_workload(self):
    return 'functional-query'

  @classmethod
  def add_test_dimensions(cls):
    super(TestExprLimits, cls).add_test_dimensions()
    if cls.exploration_strategy() != 'exhaustive':
      # Ensure the test runs with codegen enabled and disabled, even when the
      # exploration strategy is not exhaustive.
      cls.TestMatrix.clear_dimension('exec_option')
      cls.TestMatrix.add_dimension(create_exec_option_dimension(
          cluster_sizes=[0], disable_codegen_options=[False, True], batch_sizes=[0]))

    # There is no reason to run these tests using all dimensions.
    cls.TestMatrix.add_dimension(create_uncompressed_text_dimension(cls.get_workload()))

  def test_expr_child_limit(self, vector):
    # IN predicate
    in_query = "select 1 IN("
    for i in xrange(0, self.EXPR_CHILDREN_LIMIT - 1):
      in_query += str(i)
      if (i + 1 != self.EXPR_CHILDREN_LIMIT - 1):
        in_query += ","
    in_query += ")"
    self.__exec_query(in_query)

    # CASE expr
    case_query = "select case "
    for i in xrange(0, self.EXPR_CHILDREN_LIMIT/2):
      case_query += " when true then 1"
    case_query += " end"
    self.__exec_query(case_query)

  def test_expr_depth_limit(self, vector):
    # Compound predicates
    and_query = "select " + self.__gen_deep_infix_expr("true", " and false")
    self.__exec_query(and_query)
    or_query = "select " + self.__gen_deep_infix_expr("true", " or false")
    self.__exec_query(or_query)

    # Arithmetic expr
    arith_query = "select " + self.__gen_deep_infix_expr("1", " + 1")
    self.__exec_query(arith_query)

    func_query = "select " + self.__gen_deep_func_expr("lower(", "'abc'", ")")
    self.__exec_query(func_query)

    # Casts.
    cast_query = "select " + self.__gen_deep_func_expr("cast(", "1", " as int)")
    self.__exec_query(cast_query)

  def __gen_deep_infix_expr(self, prefix, repeat_suffix):
    expr = prefix
    for i in xrange(self.EXPR_DEPTH_LIMIT - 1):
      expr += repeat_suffix
    return expr

  def __gen_deep_func_expr(self, open_func, base_arg, close_func):
    expr = ""
    for i in xrange(self.EXPR_DEPTH_LIMIT - 1):
      expr += open_func
    expr += base_arg
    for i in xrange(self.EXPR_DEPTH_LIMIT - 1):
      expr += close_func
    return expr

  def __exec_query(self, sql_str):
    try:
      impala_ret = self.execute_query(sql_str)
      assert impala_ret.success, "Failed to execute query %s" % (sql_str)
    except: # consider any exception a failure
      assert False, "Failed to execute query %s" % (sql_str)