def generate_random_queries(impala, random_db):
"""Generator function to produce random queries. 'impala' is the Impala service
object. random_db is the name of the database that queries should be
generated for."""
with impala.cursor(db_name=random_db) as cursor:
tables = [cursor.describe_table(t) for t in cursor.list_table_names()]
query_generator = QueryGenerator(DefaultProfile())
model_translator = SqlWriter.create()
while True:
query_model = query_generator.generate_statement(tables)
sql = model_translator.write_query(query_model)
query = Query()
query.sql = sql
query.db_name = random_db
yield query
def generate_random_queries(impala, random_db):
"""Generator function to produce random queries. 'impala' is the Impala service
object. random_db is the name of the database that queries should be
generated for."""
with impala.cursor(db_name=random_db) as cursor:
tables = [cursor.describe_table(t) for t in cursor.list_table_names()]
query_generator = QueryGenerator(DefaultProfile())
model_translator = SqlWriter.create()
while True:
query_model = query_generator.generate_statement(tables)
sql = model_translator.write_query(query_model)
query = Query()
query.sql = sql
query.db_name = random_db
yield query
def test_use_nested_width_subquery():
"""
Tests that setting DefaultProfile.use_nested_with to False works properly. Setting this
method to return False should prevent a WITH clause from being used inside a sub-query.
"""
class MockQueryProfile(DefaultProfile):
"""
A mock QueryProfile that sets use_nested_with to False and forces the
QueryGenerator to created nested queries.
"""
def __init__(self):
super(MockQueryProfile, self).__init__()
# Force the QueryGenerator to create nested queries
self._bounds['MAX_NESTED_QUERY_COUNT'] = (4, 4)
# Force the QueryGenerator to use WITH clauses whenever possible
self._probabilities['OPTIONAL_QUERY_CLAUSES']['WITH'] = 1
# Force the QueryGenerator to create inline views whenever possible
self._probabilities['MISC']['INLINE_VIEW'] = 1
def use_nested_with(self):
return False
mock_query_gen = QueryGenerator(MockQueryProfile())
# Create two tables
table_expr_list = TableExprList()
right_table = Table("right_table")
right_table.add_col(Column("right_table", "right_col", Int))
table_expr_list.append(right_table)
left_table = Table("left_table")
left_table.add_col(Column("left_table", "left_col", Int))
table_expr_list.append(left_table)
# Check that each nested_query doesn't have a with clause
for nested_query in mock_query_gen.generate_statement(table_expr_list).nested_queries:
assert nested_query.with_clause is None
class InsertStatementGenerator(object):
def __init__(self, profile):
# QueryProfile-like object
self.profile = profile
# used to generate SELECT queries for INSERT ... SELECT statements;
# to ensure state is completely reset, this is created anew with each call to
# generate_statement()
self.select_stmt_generator = None
def generate_statement(self, tables, dml_table):
"""
Return a randomly generated INSERT statement.
tables should be a list of Table objects. A typical source of such a list comes from
db_connection.DbCursor.describe_common_tables(). This list describes the possible
"sources" of the INSERT's WITH and FROM/WHERE clauses.
dml_table is a required Table object. The INSERT will be into this table.
"""
if not (isinstance(tables, list) and len(tables) > 0 and all(
(isinstance(t, Table) for t in tables))):
raise Exception('tables must be a not-empty list of Table objects')
if not isinstance(dml_table, Table):
raise Exception('dml_table must be a Table')
self.select_stmt_generator = QueryGenerator(self.profile)
if dml_table.primary_keys:
insert_statement = InsertStatement(
conflict_action=InsertStatement.CONFLICT_ACTION_IGNORE)
else:
insert_statement = InsertStatement(
conflict_action=InsertStatement.CONFLICT_ACTION_DEFAULT)
insert_statement.execution = StatementExecutionMode.DML_TEST
# Choose whether this is a
# INSERT INTO table SELECT/VALUES
# or
# INSERT INTO table (col1, col2, ...) SELECT/VALUES
# If the method returns None, it's the former.
insert_column_list = self.profile.choose_insert_column_list(dml_table)
insert_statement.insert_clause = InsertClause(
dml_table, column_list=insert_column_list)
# We still need to internally track the columns we're inserting. Keep in mind None
# means "all" without an explicit column list. Since we've already created the
# InsertClause object though, we can fill this in for ourselves.
if insert_column_list is None:
insert_column_list = dml_table.cols
insert_item_data_types = [col.type for col in insert_column_list]
# Decide whether this is INSERT VALUES or INSERT SELECT
insert_source_clause = self.profile.choose_insert_source_clause()
if issubclass(insert_source_clause, Query):
# Use QueryGenerator()'s public interface to generate the SELECT.
select_query = self.select_stmt_generator.generate_statement(
tables, select_item_data_types=insert_item_data_types)
# To avoid many loss-of-precision errors, explicitly cast the SelectItems. The
# generator's type system is not near sophisticated enough to know how random
# expressions will be implicitly casted in the databases. This requires less work
# to implement. IMPALA-4693 considers alternative approaches.
self._cast_select_items(select_query, insert_column_list)
insert_statement.with_clause = deepcopy(select_query.with_clause)
select_query.with_clause = None
insert_statement.select_query = select_query
elif issubclass(insert_source_clause, ValuesClause):
insert_statement.values_clause = self._generate_values_clause(
insert_column_list)
else:
raise Exception('unsupported INSERT source clause: {0}'.format(
insert_source_clause))
return insert_statement
def _generate_values_clause(self, columns):
"""
Return a VALUES clause containing a variable number of rows.
The values corresponding to primary keys will be non-null constants. Any other
columns could be null, constants, or function trees that may or may not evaluate to
null.
"""
values_rows = []
for _ in xrange(self.profile.choose_insert_values_row_count()):
values_row = []
for col in columns:
if col.is_primary_key:
val = self.profile.choose_constant(
return_type=col.exact_type, allow_null=False)
elif 'constant' == self.profile.choose_values_item_expr():
val = self.profile.choose_constant(
return_type=col.exact_type, allow_null=True)
else:
func_tree = self.select_stmt_generator.create_func_tree(
col.type, allow_subquery=False)
val = self.select_stmt_generator.populate_func_with_vals(
func_tree)
# Only the generic type, not the exact type, of the value will be known. To
# avoid a lot of failed queries due to precision errors, we cast the val to
# the exact type of the column. This will still not prevent "out of range"
# conditions, as we don't try to evaluate the random expressions.
val = CastFunc(val, col.exact_type)
values_row.append(val)
values_rows.append(ValuesRow(values_row))
return ValuesClause(values_rows)
def _cast_select_items(self, select_query, column_list):
"""
For a given Query select_query and a column_list (list of Columns), cast each select
item in select_query to the exact type of the column.
A Query may have a UNION, recursively do this down the line.
"""
for col_idx, select_item in enumerate(
select_query.select_clause.items):
cast_val_expr = CastFunc(select_item.val_expr,
column_list[col_idx].exact_type)
select_item.val_expr = cast_val_expr
if select_query.union_clause:
self._cast_select_items(select_query.union_clause.query,
column_list)
class InsertStatementGenerator(object):
def __init__(self, profile):
# QueryProfile-like object
self.profile = profile
# used to generate SELECT queries for INSERT/UPSERT ... SELECT statements;
# to ensure state is completely reset, this is created anew with each call to
# generate_statement()
self.select_stmt_generator = None
def generate_statement(self, tables, dml_table):
"""
Return a randomly generated INSERT or UPSERT statement. Note that UPSERTs are very
similar to INSERTs, which is why this generator handles both.
tables should be a list of Table objects. A typical source of such a list comes from
db_connection.DbCursor.describe_common_tables(). This list describes the possible
"sources" of the INSERT/UPSERT's WITH and FROM/WHERE clauses.
dml_table is a required Table object. The INSERT/UPSERT will be into this table.
"""
if not (isinstance(tables, list) and len(tables) > 0 and
all((isinstance(t, Table) for t in tables))):
raise Exception('tables must be a not-empty list of Table objects')
if not isinstance(dml_table, Table):
raise Exception('dml_table must be a Table')
self.select_stmt_generator = QueryGenerator(self.profile)
insert_statement = InsertStatement(execution=StatementExecutionMode.DML_TEST)
# Choose whether this is a
# INSERT/UPSERT INTO table SELECT/VALUES
# or
# INSERT/UPSERT INTO table (col1, col2, ...) SELECT/VALUES
# If the method returns None, it's the former.
insert_column_list = self.profile.choose_insert_column_list(dml_table)
if dml_table.primary_keys:
# Having primary keys implies the table is a Kudu table, which makes it subject to
# both INSERTs (with automatic ignoring of primary key duplicates) and UPSERTs.
conflict_action = self.profile.choose_insert_vs_upsert()
else:
conflict_action = InsertClause.CONFLICT_ACTION_DEFAULT
insert_statement.insert_clause = InsertClause(
dml_table, column_list=insert_column_list, conflict_action=conflict_action)
# We still need to internally track the columns we're inserting. Keep in mind None
# means "all" without an explicit column list. Since we've already created the
# InsertClause object though, we can fill this in for ourselves.
if insert_column_list is None:
insert_column_list = dml_table.cols
insert_item_data_types = [col.type for col in insert_column_list]
# Decide whether this is INSERT/UPSERT VALUES or INSERT/UPSERT SELECT
insert_source_clause = self.profile.choose_insert_source_clause()
if issubclass(insert_source_clause, Query):
# Use QueryGenerator()'s public interface to generate the SELECT.
select_query = self.select_stmt_generator.generate_statement(
tables, select_item_data_types=insert_item_data_types)
# To avoid many loss-of-precision errors, explicitly cast the SelectItems. The
# generator's type system is not near sophisticated enough to know how random
# expressions will be implicitly casted in the databases. This requires less work
# to implement. IMPALA-4693 considers alternative approaches.
self._cast_select_items(select_query, insert_column_list)
insert_statement.with_clause = deepcopy(select_query.with_clause)
select_query.with_clause = None
insert_statement.select_query = select_query
elif issubclass(insert_source_clause, ValuesClause):
insert_statement.values_clause = self._generate_values_clause(insert_column_list)
else:
raise Exception('unsupported INSERT/UPSERT source clause: {0}'.format(
insert_source_clause))
return insert_statement
def _generate_values_clause(self, columns):
"""
Return a VALUES clause containing a variable number of rows.
The values corresponding to primary keys will be non-null constants. Any other
columns could be null, constants, or function trees that may or may not evaluate to
null.
"""
values_rows = []
for _ in xrange(self.profile.choose_insert_values_row_count()):
values_row = []
for col in columns:
if col.is_primary_key:
val = self.profile.choose_constant(return_type=col.exact_type, allow_null=False)
elif 'constant' == self.profile.choose_values_item_expr():
val = self.profile.choose_constant(return_type=col.exact_type, allow_null=True)
else:
func_tree = self.select_stmt_generator.create_func_tree(
col.type, allow_subquery=False)
val = self.select_stmt_generator.populate_func_with_vals(func_tree)
# Only the generic type, not the exact type, of the value will be known. To
# avoid a lot of failed queries due to precision errors, we cast the val to
# the exact type of the column. This will still not prevent "out of range"
# conditions, as we don't try to evaluate the random expressions.
val = CastFunc(val, col.exact_type)
values_row.append(val)
values_rows.append(ValuesRow(values_row))
return ValuesClause(values_rows)
def _cast_select_items(self, select_query, column_list):
"""
For a given Query select_query and a column_list (list of Columns), cast each select
item in select_query to the exact type of the column.
A Query may have a UNION, recursively do this down the line.
"""
for col_idx, select_item in enumerate(select_query.select_clause.items):
cast_val_expr = CastFunc(select_item.val_expr, column_list[col_idx].exact_type)
select_item.val_expr = cast_val_expr
if select_query.union_clause:
self._cast_select_items(select_query.union_clause.query, column_list)
class Job(object):
'''Represents a Query Generator Job. One ImpalaDockerEnv is associated with it. Able to
execute queries by either generaing them based on a provided query profile or by
extracting queries from an existing report. A report is generated when it finishes
running.
'''
def __init__(self,
query_profile,
job_id,
run_name = 'default',
time_limit_sec = 24 * 3600,
git_command = None,
parent_job = None):
self.git_hash = ''
self.job_id = job_id
self.job_name = run_name
self.parent_job = parent_job
self.query_profile = query_profile or (
ImpalaNestedTypesProfile() if NESTED_TYPES_MODE else DefaultProfile())
self.ref_connection = None
self.result_list = []
self.start_time = time()
self.stop_time = None
self.target_stop_time = time() + time_limit_sec
self.test_connection = None
self.num_queries_executed = 0
self.num_queries_returned_correct_data = 0
self.flatten_dialect = 'POSTGRESQL' if NESTED_TYPES_MODE else None
self.impala_env = ImpalaDockerEnv(git_command)
def __getstate__(self):
'''For pickling'''
result = {}
result['job_id'] = self.job_id
result['job_name'] = self.job_name
result['parent_job'] = self.parent_job
result['result_list'] = self.result_list
result['git_hash'] = self.git_hash
result['start_time'] = self.start_time
result['stop_time'] = self.stop_time
result['num_queries_executed'] = self.num_queries_executed
result['num_queries_returned_correct_data'] = self.num_queries_returned_correct_data
return result
def prepare(self):
'''Prepares the environment and connects to Postgres and Impala running inside the
Docker container.
'''
LOG.info('Starting Job Preparation')
self.impala_env.prepare()
LOG.info('Job Preparation Complete')
self.ref_connection = PostgresqlConnection(
user_name=POSTGRES_USER_NAME,
password=None,
host_name=self.impala_env.host,
port=self.impala_env.postgres_port,
db_name=POSTGRES_DATABASE_NAME)
LOG.info('Created ref_connection')
self.start_impala()
self.git_hash = self.impala_env.get_git_hash()
def get_stack(self):
stack_trace = self.impala_env.get_stack()
LOG.info('Stack Trace: {0}'.format(stack_trace))
return stack_trace
def start_impala(self):
'''Starts impala and creates a connection to it. '''
self.impala_env.start_impala()
self.test_connection = ImpalaConnection(
host_name=self.impala_env.host,
port=self.impala_env.impala_port,
user_name=None,
db_name=DATABASE_NAME)
self.test_connection.reconnect()
self.query_result_comparator = QueryResultComparator(
self.query_profile,
self.ref_connection,
self.test_connection,
query_timeout_seconds=4*60,
flatten_dialect='POSTGRESQL')
LOG.info('Created query result comparator')
LOG.info(str(self.query_result_comparator.__dict__))
def is_impala_running(self):
return self.impala_env.is_impala_running()
def save_pickle(self):
'''Saves self as pickle. This is normally done when the job finishes running. '''
with open(join_path(PATH_TO_FINISHED_JOBS, self.job_id), 'w') as f:
pickle.dump(self, f)
LOG.info('Saved Completed Job Pickle')
def queries_to_be_executed(self):
'''Generator that outputs query models. They are either generated based on the query
profile, or they are extracted from an existing report.
'''
if self.parent_job:
# If parent job is specified, get the queries from the parent job report
with open(join_path(PATH_TO_REPORTS, self.parent_job), 'r') as f:
parent_report = pickle.load(f)
for error_type in ['stack', 'row_counts', 'mismatch']:
for query in parent_report.grouped_results[error_type]:
yield query['model']
else:
# If parent job is not specified, generate queries with QueryGenerator
num_unexpected_errors = 0
while num_unexpected_errors < NUM_UNEXPECTED_ERRORS_THRESHOLD:
query = None
try:
# TODO: Support DML statements. Possibly this be part of IMPALA-4600.
self.query_generator = QueryGenerator(self.query_profile)
query = self.query_generator.generate_statement(self.common_tables)
except IndexError as e:
# This is a query generator bug that happens extremely rarely
LOG.info('Query Generator Choice Problem, {0}'.format(e))
continue
except Exception as e:
LOG.info('Unexpected error in queries_to_be_executed, {0}'.format(e))
num_unexpected_errors += 1
if num_unexpected_errors > NUM_UNEXPECTED_ERRORS_THRESHOLD:
LOG.error('Num Unexpected Errors above threshold')
raise
else:
continue
yield query
def generate_report(self):
'''Generate report and save it into the reports directory. '''
from report import Report
rep = Report(self.job_id)
rep.save_pickle()
def start(self):
try:
self.prepare()
self.query_generator = QueryGenerator(self.query_profile)
if NESTED_TYPES_MODE:
self.common_tables = DbCursor.describe_common_tables(
[self.test_connection.cursor()])
else:
self.common_tables = DbCursor.describe_common_tables(
[self.test_connection.cursor(), self.ref_connection.cursor()])
for query_model in self.queries_to_be_executed():
LOG.info('About to execute query.')
result_dict = self.run_query(query_model)
LOG.info('Query Executed successfully.')
self.num_queries_executed += 1
if result_dict:
self.result_list.append(result_dict)
LOG.info('Time Left: {0}'.format(self.target_stop_time - time()))
if time() > self.target_stop_time:
break
self.stop_time = time()
self.save_pickle()
self.generate_report()
LOG.info('Generated Report')
except:
LOG.exception('Unexpected Exception in start')
raise
finally:
self.impala_env.stop_docker()
LOG.info('Docker Stopped')
try:
os.remove(join_path(PATH_TO_SCHEDULE, self.job_id))
LOG.info('Schedule file removed')
except OSError:
LOG.info('Unable to remove schedule file.')
def reproduce_crash(self, query_model):
'''Check if the given query_model causes a crash. Returns the number of times the
query had to be run to cause a crash.
'''
NUM_TRIES = 5
self.start_impala()
for try_num in range(1, NUM_TRIES + 1):
self.query_result_comparator.compare_query_results(query_model)
if not self.is_impala_running():
return try_num
def run_query(self, query_model):
'''Runs a single query. '''
if not self.is_impala_running():
LOG.info('Impala is not running, starting Impala.')
self.start_impala()
def run_query_internal():
self.comparison_result = self.query_result_comparator.compare_query_results(
query_model)
self.comparison_result = None
internal_thread = Thread(
target=run_query_internal,
name='run_query_internal_{0}'.format(self.job_id))
internal_thread.daemon = True
internal_thread.start()
internal_thread.join(timeout=600)
if internal_thread.is_alive():
LOG.info('run_query_internal is alive, restarting Impala Environment')
self.impala_env.stop_docker()
self.prepare()
return None
else:
LOG.info('run_query_internal is dead as expected')
comparison_result = self.comparison_result
if comparison_result.query_timed_out:
LOG.info('Query Timeout Exception')
restart_impala = True
else:
restart_impala = False
result_dict = {}
if self.is_impala_running():
if comparison_result.error:
result_dict = self.comparison_result_analysis(comparison_result)
result_dict['model'] = query_model
elif comparison_result.query_resulted_in_data:
self.num_queries_returned_correct_data += 1
else:
LOG.info('CRASH OCCURED')
result_dict = self.comparison_result_analysis(comparison_result)
result_dict['model'] = query_model
result_dict['stack'] = self.get_stack()
result_dict['num_tries_to_reproduce'] = self.reproduce_crash(query_model)
if restart_impala:
self.start_impala()
return result_dict
def comparison_result_analysis(self, comparison_result):
'''Get useful information from the comparison_result. '''
result_dict = {}
result_dict['error'] = comparison_result.error
result_dict['mismatch_col'] = comparison_result.mismatch_at_col_number
result_dict['mismatch_ref_row'] = comparison_result.ref_row
result_dict['mismatch_test_row'] = comparison_result.test_row
result_dict['ref_row_count'] = comparison_result.ref_row_count
result_dict['ref_sql'] = comparison_result.ref_sql
result_dict['test_row_count'] = comparison_result.test_row_count
result_dict['test_sql'] = comparison_result.test_sql
return result_dict
