tests/common/custom_cluster_test_suite.py - impala - Git at Google

 # 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.
 #
 # Superclass for all tests that need a custom cluster.
 # TODO: Configure cluster size and other parameters.

 import logging
 import os
 import os.path
 import pipes
 import pytest
 import re
 import subprocess
 from subprocess import check_call
 from tests.common.impala_test_suite import ImpalaTestSuite
 from tests.common.impala_cluster import ImpalaCluster
 from tests.util.filesystem_utils import IS_LOCAL
 from time import sleep

 IMPALA_HOME = os.environ['IMPALA_HOME']
 CLUSTER_SIZE = 3
 NUM_COORDINATORS = CLUSTER_SIZE

 # Additional args passed to respective daemon command line.
 IMPALAD_ARGS = 'impalad_args'
 STATESTORED_ARGS = 'state_store_args'
 CATALOGD_ARGS = 'catalogd_args'
 # Additional args passed to the start-impala-cluster script.
 START_ARGS = 'start_args'
 SENTRY_CONFIG = 'sentry_config'
 SENTRY_LOG_DIR = 'sentry_log_dir'
 # Default query options passed to the impala daemon command line. Handled separately from
 # other impala daemon arguments to allow merging multiple defaults into a single list.
 DEFAULT_QUERY_OPTIONS = 'default_query_options'
 IMPALA_LOG_DIR = 'impala_log_dir'

 # Run with fast topic updates by default to reduce time to first query running.
 DEFAULT_STATESTORE_ARGS = '--statestore_update_frequency_ms=50 \
     --statestore_priority_update_frequency_ms=50 \
     --statestore_heartbeat_frequency_ms=50'

 class CustomClusterTestSuite(ImpalaTestSuite):
   """Every test in a test suite deriving from this class gets its own Impala cluster.
   Custom arguments may be passed to the cluster by using the @with_args decorator."""
   @classmethod
   def get_workload(cls):
     return 'tpch'

   @classmethod
   def add_test_dimensions(cls):
     super(CustomClusterTestSuite, cls).add_test_dimensions()
     cls.add_custom_cluster_constraints()

   @classmethod
   def add_custom_cluster_constraints(cls):
     # Defines constraints for custom cluster tests, called by add_test_dimensions.
     # By default, custom cluster tests only run on text/none and with a limited set of
     # exec options. Subclasses may override this to relax these default constraints.
     super(CustomClusterTestSuite, cls).add_test_dimensions()
     cls.ImpalaTestMatrix.add_constraint(lambda v:
         v.get_value('table_format').file_format == 'text' and
         v.get_value('table_format').compression_codec == 'none')
     cls.ImpalaTestMatrix.add_constraint(lambda v:
         v.get_value('exec_option')['batch_size'] == 0 and
         v.get_value('exec_option')['disable_codegen'] == False and
         v.get_value('exec_option')['num_nodes'] == 0)

   @classmethod
   def setup_class(cls):
     # Explicit override of ImpalaTestSuite.setup_class(). For custom cluster, the
     # ImpalaTestSuite.setup_class() procedure needs to happen on a per-method basis.
     # IMPALA-3614: @SkipIfLocal.multiple_impalad workaround
     # IMPALA-2943 TODO: When pytest is upgraded, see if this explicit skip can be
     # removed in favor of the class-level SkipifLocal.multiple_impalad decorator.
     if IS_LOCAL:
       pytest.skip("multiple impalads needed")

   @classmethod
   def teardown_class(cls):
     # Explicit override of ImpalaTestSuite.teardown_class(). For custom cluster, the
     # ImpalaTestSuite.teardown_class() procedure needs to happen on a per-method basis.
     pass

   @staticmethod
   def with_args(impalad_args=None, statestored_args=None, catalogd_args=None,
       start_args=None, sentry_config=None, default_query_options=None,
       impala_log_dir=None, sentry_log_dir=None):
     """Records arguments to be passed to a cluster by adding them to the decorated
     method's func_dict"""
     def decorate(func):
       if impalad_args is not None:
         func.func_dict[IMPALAD_ARGS] = impalad_args
       if statestored_args is None:
         func.func_dict[STATESTORED_ARGS] = DEFAULT_STATESTORE_ARGS
       else:
         func.func_dict[STATESTORED_ARGS] = \
             DEFAULT_STATESTORE_ARGS + " " + statestored_args
       if catalogd_args is not None:
         func.func_dict[CATALOGD_ARGS] = catalogd_args
       if start_args is not None:
         func.func_dict[START_ARGS] = start_args
       if sentry_config is not None:
         func.func_dict[SENTRY_CONFIG] = sentry_config
       if sentry_log_dir is not None:
         func.func_dict[SENTRY_LOG_DIR] = sentry_log_dir
       if default_query_options is not None:
         func.func_dict[DEFAULT_QUERY_OPTIONS] = default_query_options
       if impala_log_dir is not None:
         func.func_dict[IMPALA_LOG_DIR] = impala_log_dir
       return func
     return decorate

   def setup_method(self, method):
     cluster_args = list()
     for arg in [IMPALAD_ARGS, STATESTORED_ARGS, CATALOGD_ARGS]:
       if arg in method.func_dict:
         cluster_args.append("--%s=\"%s\" " % (arg, method.func_dict[arg]))
     if START_ARGS in method.func_dict:
       cluster_args.append(method.func_dict[START_ARGS])

     if SENTRY_CONFIG in method.func_dict:
       self._start_sentry_service(method.func_dict[SENTRY_CONFIG],
           method.func_dict.get(SENTRY_LOG_DIR))
     # Start a clean new cluster before each test
     if IMPALA_LOG_DIR in method.func_dict:
       self._start_impala_cluster(cluster_args,
           default_query_options=method.func_dict.get(DEFAULT_QUERY_OPTIONS),
           impala_log_dir=method.func_dict[IMPALA_LOG_DIR])
     else:
       self._start_impala_cluster(cluster_args,
           default_query_options=method.func_dict.get(DEFAULT_QUERY_OPTIONS))
     super(CustomClusterTestSuite, self).setup_class()

   def teardown_method(self, method):
     super(CustomClusterTestSuite, self).teardown_class()

   @classmethod
   def _stop_impala_cluster(cls):
     # TODO: Figure out a better way to handle case where processes are just starting
     # / cleaning up so that sleeps are not needed.
     sleep(2)
     check_call([os.path.join(IMPALA_HOME, 'bin/start-impala-cluster.py'), '--kill_only'])
     sleep(2)

   @classmethod
   def _start_sentry_service(cls, sentry_service_config, sentry_log_dir=None):
     sentry_env = dict(os.environ)
     if sentry_log_dir is not None:
         sentry_env['SENTRY_LOG_DIR'] = sentry_log_dir
     sentry_env['SENTRY_SERVICE_CONFIG'] = sentry_service_config
     call = subprocess.Popen(
         ['/bin/bash', '-c', os.path.join(IMPALA_HOME,
                                          'testdata/bin/run-sentry-service.sh')],
         env=sentry_env)
     call.wait()
     if call.returncode != 0:
       raise RuntimeError("Unable to start Sentry")

   @classmethod
   def _stop_sentry_service(cls):
     subprocess.check_call([os.path.join(os.environ["IMPALA_HOME"],
                                         "testdata/bin/kill-sentry-service.sh")],
                           close_fds=True)

   @classmethod
   def _start_impala_cluster(cls, options, impala_log_dir=os.getenv('LOG_DIR', "/tmp/"),
       cluster_size=CLUSTER_SIZE, num_coordinators=NUM_COORDINATORS,
       use_exclusive_coordinators=False, log_level=1, expected_num_executors=CLUSTER_SIZE,
       default_query_options=None):
     cls.impala_log_dir = impala_log_dir
     # We ignore TEST_START_CLUSTER_ARGS here. Custom cluster tests specifically test that
     # certain custom startup arguments work and we want to keep them independent of dev
     # environments.
     cmd = [os.path.join(IMPALA_HOME, 'bin/start-impala-cluster.py'),
            '--cluster_size=%d' % cluster_size,
            '--num_coordinators=%d' % num_coordinators,
            '--log_dir=%s' % impala_log_dir,
            '--log_level=%s' % log_level]

     if use_exclusive_coordinators:
       cmd.append("--use_exclusive_coordinators")

     if pytest.config.option.use_local_catalog:
       cmd.append("--impalad_args=--use_local_catalog=1")
       cmd.append("--catalogd_args=--catalog_topic_mode=minimal")

     if pytest.config.option.pull_incremental_statistics:
       cmd.append("--impalad_args=%s --catalogd_args=%s" %
                  ("--pull_incremental_statistics", "--pull_incremental_statistics"))

     default_query_option_kvs = []
     # Put any defaults first, then any arguments after that so they can override defaults.
     if os.environ.get("ERASURE_CODING") == "true":
       default_query_option_kvs.append(("allow_erasure_coded_files", "true"))
     if default_query_options is not None:
       default_query_option_kvs.extend(default_query_options)
     # Add the default query options after any arguments. This will override any default
     # options set in --impalad_args by design to force tests to pass default_query_options
     # into this function directly.
     options.append("--impalad_args=--default_query_options={0}".format(
         ','.join(["{0}={1}".format(k, v) for k, v in default_query_option_kvs])))

     logging.info("Starting cluster with command: %s" %
                  " ".join(pipes.quote(arg) for arg in cmd + options))
     try:
       check_call(cmd + options, close_fds=True)
     finally:
       # Failure tests expect cluster to be initialised even if start-impala-cluster fails.
       cls.cluster = ImpalaCluster()
     statestored = cls.cluster.statestored
     if statestored is None:
       raise Exception("statestored was not found")

     # The number of statestore subscribers is
     # cluster_size (# of impalad) + 1 (for catalogd).
     expected_subscribers = cluster_size + 1

     statestored.service.wait_for_live_subscribers(expected_subscribers, timeout=60)
     for impalad in cls.cluster.impalads:
       impalad.service.wait_for_num_known_live_backends(expected_num_executors, timeout=60)

   def assert_impalad_log_contains(self, level, line_regex, expected_count=1):
     """
     Convenience wrapper around assert_log_contains for impalad logs.
     """
     self.assert_log_contains("impalad", level, line_regex, expected_count)

   def assert_catalogd_log_contains(self, level, line_regex, expected_count=1):
     """
     Convenience wrapper around assert_log_contains for catalogd logs.
     """
     self.assert_log_contains("catalogd", level, line_regex, expected_count)

   def assert_log_contains(self, daemon, level, line_regex, expected_count=1):
     """
     Assert that the daemon log with specified level (e.g. ERROR, WARNING, INFO) contains
     expected_count lines with a substring matching the regex. When expected_count is -1,
     at least one match is expected.
     When using this method to check log files of running processes, the caller should
     make sure that log buffering has been disabled, for example by adding
     '-logbuflevel=-1' to the daemon startup options.
     """
     pattern = re.compile(line_regex)
     found = 0
     log_file_path = os.path.join(self.impala_log_dir, daemon + "." + level)
     # Resolve symlinks to make finding the file easier.
     log_file_path = os.path.realpath(log_file_path)
     with open(log_file_path) as log_file:
       for line in log_file:
         if pattern.search(line):
           found += 1
     if expected_count == -1:
       assert found > 0, "Expected at least one line in file %s matching regex '%s'"\
         ", but found none." % (log_file_path, line_regex)
     else:
       assert found == expected_count, "Expected %d lines in file %s matching regex '%s'"\
         ", but found %d lines. Last line was: \n%s" %\
         (expected_count, log_file_path, line_regex, found, line)
	# 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.
	#
	# Superclass for all tests that need a custom cluster.
	# TODO: Configure cluster size and other parameters.

	import logging
	import os
	import os.path
	import pipes
	import pytest
	import re
	import subprocess
	from subprocess import check_call
	from tests.common.impala_test_suite import ImpalaTestSuite
	from tests.common.impala_cluster import ImpalaCluster
	from tests.util.filesystem_utils import IS_LOCAL
	from time import sleep

	IMPALA_HOME = os.environ['IMPALA_HOME']
	CLUSTER_SIZE = 3
	NUM_COORDINATORS = CLUSTER_SIZE

	# Additional args passed to respective daemon command line.
	IMPALAD_ARGS = 'impalad_args'
	STATESTORED_ARGS = 'state_store_args'
	CATALOGD_ARGS = 'catalogd_args'
	# Additional args passed to the start-impala-cluster script.
	START_ARGS = 'start_args'
	SENTRY_CONFIG = 'sentry_config'
	SENTRY_LOG_DIR = 'sentry_log_dir'
	# Default query options passed to the impala daemon command line. Handled separately from
	# other impala daemon arguments to allow merging multiple defaults into a single list.
	DEFAULT_QUERY_OPTIONS = 'default_query_options'
	IMPALA_LOG_DIR = 'impala_log_dir'

	# Run with fast topic updates by default to reduce time to first query running.
	DEFAULT_STATESTORE_ARGS = '--statestore_update_frequency_ms=50 \
	--statestore_priority_update_frequency_ms=50 \
	--statestore_heartbeat_frequency_ms=50'

	class CustomClusterTestSuite(ImpalaTestSuite):
	"""Every test in a test suite deriving from this class gets its own Impala cluster.
	Custom arguments may be passed to the cluster by using the @with_args decorator."""
	@classmethod
	def get_workload(cls):
	return 'tpch'

	@classmethod
	def add_test_dimensions(cls):
	super(CustomClusterTestSuite, cls).add_test_dimensions()
	cls.add_custom_cluster_constraints()

	@classmethod
	def add_custom_cluster_constraints(cls):
	# Defines constraints for custom cluster tests, called by add_test_dimensions.
	# By default, custom cluster tests only run on text/none and with a limited set of
	# exec options. Subclasses may override this to relax these default constraints.
	super(CustomClusterTestSuite, cls).add_test_dimensions()
	cls.ImpalaTestMatrix.add_constraint(lambda v:
	v.get_value('table_format').file_format == 'text' and
	v.get_value('table_format').compression_codec == 'none')
	cls.ImpalaTestMatrix.add_constraint(lambda v:
	v.get_value('exec_option')['batch_size'] == 0 and
	v.get_value('exec_option')['disable_codegen'] == False and
	v.get_value('exec_option')['num_nodes'] == 0)

	@classmethod
	def setup_class(cls):
	# Explicit override of ImpalaTestSuite.setup_class(). For custom cluster, the
	# ImpalaTestSuite.setup_class() procedure needs to happen on a per-method basis.
	# IMPALA-3614: @SkipIfLocal.multiple_impalad workaround
	# IMPALA-2943 TODO: When pytest is upgraded, see if this explicit skip can be
	# removed in favor of the class-level SkipifLocal.multiple_impalad decorator.
	if IS_LOCAL:
	pytest.skip("multiple impalads needed")

	@classmethod
	def teardown_class(cls):
	# Explicit override of ImpalaTestSuite.teardown_class(). For custom cluster, the
	# ImpalaTestSuite.teardown_class() procedure needs to happen on a per-method basis.
	pass

	@staticmethod
	def with_args(impalad_args=None, statestored_args=None, catalogd_args=None,
	start_args=None, sentry_config=None, default_query_options=None,
	impala_log_dir=None, sentry_log_dir=None):
	"""Records arguments to be passed to a cluster by adding them to the decorated
	method's func_dict"""
	def decorate(func):
	if impalad_args is not None:
	func.func_dict[IMPALAD_ARGS] = impalad_args
	if statestored_args is None:
	func.func_dict[STATESTORED_ARGS] = DEFAULT_STATESTORE_ARGS
	else:
	func.func_dict[STATESTORED_ARGS] = \
	DEFAULT_STATESTORE_ARGS + " " + statestored_args
	if catalogd_args is not None:
	func.func_dict[CATALOGD_ARGS] = catalogd_args
	if start_args is not None:
	func.func_dict[START_ARGS] = start_args
	if sentry_config is not None:
	func.func_dict[SENTRY_CONFIG] = sentry_config
	if sentry_log_dir is not None:
	func.func_dict[SENTRY_LOG_DIR] = sentry_log_dir
	if default_query_options is not None:
	func.func_dict[DEFAULT_QUERY_OPTIONS] = default_query_options
	if impala_log_dir is not None:
	func.func_dict[IMPALA_LOG_DIR] = impala_log_dir
	return func
	return decorate

	def setup_method(self, method):
	cluster_args = list()
	for arg in [IMPALAD_ARGS, STATESTORED_ARGS, CATALOGD_ARGS]:
	if arg in method.func_dict:
	cluster_args.append("--%s=\"%s\" " % (arg, method.func_dict[arg]))
	if START_ARGS in method.func_dict:
	cluster_args.append(method.func_dict[START_ARGS])

	if SENTRY_CONFIG in method.func_dict:
	self._start_sentry_service(method.func_dict[SENTRY_CONFIG],
	method.func_dict.get(SENTRY_LOG_DIR))
	# Start a clean new cluster before each test
	if IMPALA_LOG_DIR in method.func_dict:
	self._start_impala_cluster(cluster_args,
	default_query_options=method.func_dict.get(DEFAULT_QUERY_OPTIONS),
	impala_log_dir=method.func_dict[IMPALA_LOG_DIR])
	else:
	self._start_impala_cluster(cluster_args,
	default_query_options=method.func_dict.get(DEFAULT_QUERY_OPTIONS))
	super(CustomClusterTestSuite, self).setup_class()

	def teardown_method(self, method):
	super(CustomClusterTestSuite, self).teardown_class()

	@classmethod
	def _stop_impala_cluster(cls):
	# TODO: Figure out a better way to handle case where processes are just starting
	# / cleaning up so that sleeps are not needed.
	sleep(2)
	check_call([os.path.join(IMPALA_HOME, 'bin/start-impala-cluster.py'), '--kill_only'])
	sleep(2)

	@classmethod
	def _start_sentry_service(cls, sentry_service_config, sentry_log_dir=None):
	sentry_env = dict(os.environ)
	if sentry_log_dir is not None:
	sentry_env['SENTRY_LOG_DIR'] = sentry_log_dir
	sentry_env['SENTRY_SERVICE_CONFIG'] = sentry_service_config
	call = subprocess.Popen(
	['/bin/bash', '-c', os.path.join(IMPALA_HOME,
	'testdata/bin/run-sentry-service.sh')],
	env=sentry_env)
	call.wait()
	if call.returncode != 0:
	raise RuntimeError("Unable to start Sentry")

	@classmethod
	def _stop_sentry_service(cls):
	subprocess.check_call([os.path.join(os.environ["IMPALA_HOME"],
	"testdata/bin/kill-sentry-service.sh")],
	close_fds=True)

	@classmethod
	def _start_impala_cluster(cls, options, impala_log_dir=os.getenv('LOG_DIR', "/tmp/"),
	cluster_size=CLUSTER_SIZE, num_coordinators=NUM_COORDINATORS,
	use_exclusive_coordinators=False, log_level=1, expected_num_executors=CLUSTER_SIZE,
	default_query_options=None):
	cls.impala_log_dir = impala_log_dir
	# We ignore TEST_START_CLUSTER_ARGS here. Custom cluster tests specifically test that
	# certain custom startup arguments work and we want to keep them independent of dev
	# environments.
	cmd = [os.path.join(IMPALA_HOME, 'bin/start-impala-cluster.py'),
	'--cluster_size=%d' % cluster_size,
	'--num_coordinators=%d' % num_coordinators,
	'--log_dir=%s' % impala_log_dir,
	'--log_level=%s' % log_level]

	if use_exclusive_coordinators:
	cmd.append("--use_exclusive_coordinators")

	if pytest.config.option.use_local_catalog:
	cmd.append("--impalad_args=--use_local_catalog=1")
	cmd.append("--catalogd_args=--catalog_topic_mode=minimal")

	if pytest.config.option.pull_incremental_statistics:
	cmd.append("--impalad_args=%s --catalogd_args=%s" %
	("--pull_incremental_statistics", "--pull_incremental_statistics"))

	default_query_option_kvs = []
	# Put any defaults first, then any arguments after that so they can override defaults.
	if os.environ.get("ERASURE_CODING") == "true":
	default_query_option_kvs.append(("allow_erasure_coded_files", "true"))
	if default_query_options is not None:
	default_query_option_kvs.extend(default_query_options)
	# Add the default query options after any arguments. This will override any default
	# options set in --impalad_args by design to force tests to pass default_query_options
	# into this function directly.
	options.append("--impalad_args=--default_query_options={0}".format(
	','.join(["{0}={1}".format(k, v) for k, v in default_query_option_kvs])))

	logging.info("Starting cluster with command: %s" %
	" ".join(pipes.quote(arg) for arg in cmd + options))
	try:
	check_call(cmd + options, close_fds=True)
	finally:
	# Failure tests expect cluster to be initialised even if start-impala-cluster fails.
	cls.cluster = ImpalaCluster()
	statestored = cls.cluster.statestored
	if statestored is None:
	raise Exception("statestored was not found")

	# The number of statestore subscribers is
	# cluster_size (# of impalad) + 1 (for catalogd).
	expected_subscribers = cluster_size + 1

	statestored.service.wait_for_live_subscribers(expected_subscribers, timeout=60)
	for impalad in cls.cluster.impalads:
	impalad.service.wait_for_num_known_live_backends(expected_num_executors, timeout=60)

	def assert_impalad_log_contains(self, level, line_regex, expected_count=1):
	"""
	Convenience wrapper around assert_log_contains for impalad logs.
	"""
	self.assert_log_contains("impalad", level, line_regex, expected_count)

	def assert_catalogd_log_contains(self, level, line_regex, expected_count=1):
	"""
	Convenience wrapper around assert_log_contains for catalogd logs.
	"""
	self.assert_log_contains("catalogd", level, line_regex, expected_count)

	def assert_log_contains(self, daemon, level, line_regex, expected_count=1):
	"""
	Assert that the daemon log with specified level (e.g. ERROR, WARNING, INFO) contains
	expected_count lines with a substring matching the regex. When expected_count is -1,
	at least one match is expected.
	When using this method to check log files of running processes, the caller should
	make sure that log buffering has been disabled, for example by adding
	'-logbuflevel=-1' to the daemon startup options.
	"""
	pattern = re.compile(line_regex)
	found = 0
	log_file_path = os.path.join(self.impala_log_dir, daemon + "." + level)
	# Resolve symlinks to make finding the file easier.
	log_file_path = os.path.realpath(log_file_path)
	with open(log_file_path) as log_file:
	for line in log_file:
	if pattern.search(line):
	found += 1
	if expected_count == -1:
	assert found > 0, "Expected at least one line in file %s matching regex '%s'"\
	", but found none." % (log_file_path, line_regex)
	else:
	assert found == expected_count, "Expected %d lines in file %s matching regex '%s'"\
	", but found %d lines. Last line was: \n%s" %\
	(expected_count, log_file_path, line_regex, found, line)