def handler(event, context):
global logger
try:
# Empty list means no keys are required
args = common.get_parameters(event, [])
logger = common.initialise_logger(args)
processes = []
for query in copy_queries:
p = multiprocessing.Process(
target=execute_query_multiple_statements,
args=(
query,
args,
),
)
processes.append(p)
p.start()
for process in processes:
process.join()
connection = database.get_connection(args)
database.execute_multiple_statements(drop_query, connection)
database.execute_statement(rename_query, connection)
return 200
except Exception as e:
logger.error("Failed to execute one or more SQL statement(s)")
logger.error(e)
return 500
def handler(event, context):
global logger
args = common.get_parameters(event, ["table-name"])
logger = common.initialise_logger(args)
connection = database.get_connection(args)
script_dir = os.path.dirname(__file__)
rel_path = "../resources"
abs_file_path = os.path.join(script_dir, rel_path)
logger.info("Creating table if it does not exist")
database.execute_statement(
open(os.path.join(abs_file_path, "create_table.sql")).read().format(
table_name=common.get_table_name(args)),
connection,
)
logger.info("Creating user if not exists and grant access")
database.execute_multiple_statements(
open(os.path.join(abs_file_path, "grant_user.sql")).read().format(
table_name=common.get_table_name(args)),
connection,
)
logger.info("Create table alteration stored procedures")
database.execute_multiple_statements(
open(os.path.join(abs_file_path, "alter_table.sql")).read(),
connection,
)
logger.info("Execute table alteration stored procedures")
partition_count = args[
"partition-count"] if "partition-count" in args else 1
database.call_procedure(
connection,
"alter_reconciliation_table",
[common.get_table_name(args), partition_count],
)
logger.info("Validate table and users exist and the structure is correct")
table_valid = validate_table(args["rds_database_name"],
common.get_table_name(args), connection)
connection.close()
if not table_valid:
raise RuntimeError(
f"Schema is invalid in table: {common.get_table_name(args)}")
def test_drop_query(self):
execute_mock = MagicMock()
execute_mock.with_rows = True
cursor_mock = MagicMock()
cursor_mock.execute.return_value = [execute_mock]
connection_mock = MagicMock()
connection_mock.cursor.return_value = cursor_mock
database.execute_multiple_statements(drop_query, connection_mock)
cursor_mock.execute.assert_called_once_with(drop_query, multi=True)
def test_copy_queries(self):
execute_mock = MagicMock()
execute_mock.with_rows = True
cursor_mock = MagicMock()
cursor_mock.description = [("id"), ("data")]
cursor_mock.execute.return_value = [execute_mock]
connection_mock = MagicMock()
connection_mock.cursor.return_value = cursor_mock
for query in copy_queries:
database.execute_multiple_statements(query, connection_mock)
expected_calls = [call(i, multi=True) for i in copy_queries]
cursor_mock.execute.assert_has_calls(expected_calls)
self.assertEqual(3, cursor_mock.execute.call_count)
def execute_query_multiple_statements(query, args):
query_connection = database.get_connection(args)
database.execute_multiple_statements(query, query_connection)