Create pg_query_state extension in the case of TPC-DS test

Author: ololobus

tests/pg_qs_test_runner.py

@@ -83,7 +83,7 @@ def teardown(con):
 def main(config):
 	''' Main test function '''
 	conn_params = {
-			key:config.__dict__[key] for key in ('host', 'port', 'user', 'database', 'password')
+		key:config.__dict__[key] for key in ('host', 'port', 'user', 'database', 'password')
 	}
 
 	if config.use_tpcds:

tests/test_cases.py

@@ -526,14 +526,17 @@ def test_timing_buffers_conflicts(config):
 class DataLoadException(Exception): pass
 class StressTestException(Exception): pass
 
-def load_tpcds_data(config):
-	print('Loading TPC-DS data...')
+def setup_tpcds(config):
+	print('Setting up TPC-DS test...')
 	subprocess.call(['./tests/prepare_stress.sh'])
 
 	try:
 		conn = psycopg2.connect(**config)
 		cur = conn.cursor()
 
+		# Create pg_query_state extension
+		cur.execute('CREATE EXTENSION IF NOT EXISTS pg_query_state')
+
 		# Create tables
 		with open('tmp_stress/tpcds-kit/tools/tpcds.sql', 'r') as f:
 			cur.execute(f.read())
@@ -561,7 +564,7 @@ def test_tpc_ds(config):
 	global stress_in_progress
 
 	stress_in_progress = True
-	load_tpcds_data(config)
+	setup_tpcds(config)
 
 	print('Preparing TPC-DS queries...')
 	# # Execute query in separate thread