def test_context_manager_multiple_2(self): self.assertEqual(TRN._contexts_entered, 0) def tester(): self.assertEqual(TRN._contexts_entered, 1) with TRN: self.assertEqual(TRN._contexts_entered, 2) sql = """SELECT EXISTS( SELECT * FROM labcontrol.test_table WHERE int_column=%s) """ TRN.add(sql, [2]) self.assertTrue(TRN.execute_fetchlast()) self.assertEqual(TRN._contexts_entered, 1) with TRN: self.assertEqual(TRN._contexts_entered, 1) sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) tester() self.assertEqual(TRN._contexts_entered, 1) self._assert_sql_equal([]) self.assertEqual(TRN._contexts_entered, 0) self._assert_sql_equal([('insert1', True, 1), ('insert2', True, 2), ('insert3', True, 3)]) self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)
def test_context_manager_multiple(self): self.assertEqual(TRN._contexts_entered, 0) with TRN: self.assertEqual(TRN._contexts_entered, 1) TRN.add("SELECT 42") with TRN: self.assertEqual(TRN._contexts_entered, 2) sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) # We exited the second context, nothing should have been executed self.assertEqual(TRN._contexts_entered, 1) self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE) self._assert_sql_equal([]) # We have exited the first context, everything should have been # executed and committed self.assertEqual(TRN._contexts_entered, 0) self._assert_sql_equal([('insert1', True, 1), ('insert2', True, 2), ('insert3', True, 3)]) self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)
def test_post_commit_funcs_error(self): def func(): raise ValueError() with self.assertRaises(RuntimeError): with TRN: TRN.add("SELECT 42") TRN.add_post_commit_func(func)
def tester(): self.assertEqual(TRN._contexts_entered, 1) with TRN: self.assertEqual(TRN._contexts_entered, 2) sql = """SELECT EXISTS( SELECT * FROM labcontrol.test_table WHERE int_column=%s)""" TRN.add(sql, [2]) self.assertTrue(TRN.execute_fetchlast()) self.assertEqual(TRN._contexts_entered, 1)
def test_add_many(self): with TRN: self.assertEqual(TRN._queries, []) sql = "INSERT INTO labcontrol.test_table (int_column) VALUES (%s)" args = [[1], [2], [3]] TRN.add(sql, args, many=True) exp = [(sql, [1]), (sql, [2]), (sql, [3])] self.assertEqual(TRN._queries, exp)
def test_execute_return(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" TRN.add(sql, ['test_insert', 2]) sql = """UPDATE labcontrol.test_table SET bool_column = %s WHERE str_column = %s RETURNING int_column""" TRN.add(sql, [False, 'test_insert']) obs = TRN.execute() self.assertEqual(obs, [[['test_insert', 2]], [[2]]])
def test_context_manager_execute(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) self._assert_sql_equal([]) self._assert_sql_equal([('insert1', True, 1), ('insert2', True, 2), ('insert3', True, 3)]) self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)
def test_execute_fetchindex(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) self.assertEqual(TRN.execute_fetchindex(), [['insert3', 3]]) sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert4', 4], ['insert5', 5], ['insert6', 6]] TRN.add(sql, args, many=True) self.assertEqual(TRN.execute_fetchindex(3), [['insert4', 4]])
def test_execute(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s)""" TRN.add(sql, ["test_insert", 2]) sql = """UPDATE labcontrol.test_table SET int_column = %s, bool_column = %s WHERE str_column = %s""" TRN.add(sql, [20, False, "test_insert"]) obs = TRN.execute() self.assertEqual(obs, [None, None]) self._assert_sql_equal([]) self._assert_sql_equal([("test_insert", False, 20)])
def test_post_commit_funcs(self): fd, fp = mkstemp() close(fd) self._files_to_remove.append(fp) def func(fp): with open(fp, 'w') as f: f.write('\n') with TRN: TRN.add("SELECT 42") TRN.add_post_commit_func(func, fp) self.assertTrue(exists(fp))
def test_context_manager_rollback(self): try: with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) TRN.execute() raise ValueError("Force exiting the context manager") except ValueError: pass self._assert_sql_equal([]) self.assertEqual(TRN._connection.get_transaction_status(), TRANSACTION_STATUS_IDLE)
def test_execute_commit_false_rollback(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) obs = TRN.execute() exp = [[['insert1', 1]], [['insert2', 2]], [['insert3', 3]]] self.assertEqual(obs, exp) self._assert_sql_equal([]) TRN.rollback() self._assert_sql_equal([])
def test_execute_many(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s)""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) sql = """UPDATE labcontrol.test_table SET int_column = %s, bool_column = %s WHERE str_column = %s""" TRN.add(sql, [20, False, 'insert2']) obs = TRN.execute() self.assertEqual(obs, [None, None, None, None]) self._assert_sql_equal([]) self._assert_sql_equal([('insert1', True, 1), ('insert3', True, 3), ('insert2', False, 20)])
def test_context_manager_checker(self): with self.assertRaises(RuntimeError): TRN.add("SELECT 42") with self.assertRaises(RuntimeError): TRN.execute() with self.assertRaises(RuntimeError): TRN.commit() with self.assertRaises(RuntimeError): TRN.rollback() with TRN: TRN.add("SELECT 42") with self.assertRaises(RuntimeError): TRN.execute()
def test_execute_commit_false_wipe_queries(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) obs = TRN.execute() exp = [[['insert1', 1]], [['insert2', 2]], [['insert3', 3]]] self.assertEqual(obs, exp) self._assert_sql_equal([]) sql = """UPDATE labcontrol.test_table SET bool_column = %s WHERE str_column = %s""" args = [False, 'insert2'] TRN.add(sql, args) self.assertEqual(TRN._queries, [(sql, args)]) TRN.execute() self._assert_sql_equal([]) self._assert_sql_equal([('insert1', True, 1), ('insert3', True, 3), ('insert2', False, 2)])
def test_add_error(self): with TRN: with self.assertRaises(TypeError): TRN.add("SELECT 42", 1) with self.assertRaises(TypeError): TRN.add("SELECT 42", {'foo': 'bar'}, many=True) with self.assertRaises(TypeError): TRN.add("SELECT 42", [1, 1], many=True)
def test_execute_fetchflatten(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s)""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) sql = "SELECT str_column, int_column FROM labcontrol.test_table" TRN.add(sql) sql = "SELECT int_column FROM labcontrol.test_table" TRN.add(sql) obs = TRN.execute_fetchflatten() self.assertEqual(obs, [1, 2, 3]) sql = "SELECT 42" TRN.add(sql) obs = TRN.execute_fetchflatten(idx=3) self.assertEqual(obs, ['insert1', 1, 'insert2', 2, 'insert3', 3])
def test_execute_fetchlast(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) sql = """SELECT EXISTS( SELECT * FROM labcontrol.test_table WHERE int_column=%s)""" TRN.add(sql, [2]) self.assertTrue(TRN.execute_fetchlast()) sql = """SELECT str_column FROM labcontrol.test_table WHERE int_column = %s""" TRN.add(sql, [2]) self.assertEqual(TRN.execute_fetchlast(), 'insert2') TRN.add(sql, [4]) self.assertIsNone(TRN.execute_fetchlast())
def test_add(self): with TRN: self.assertEqual(TRN._queries, []) sql1 = "INSERT INTO labcontrol.test_table (bool_column) VALUES (%s)" args1 = [True] TRN.add(sql1, args1) sql2 = "INSERT INTO labcontrol.test_table (int_column) VALUES (1)" TRN.add(sql2) args3 = (False, ) TRN.add(sql1, args3) sql3 = """INSERT INTO labcontrol.test_table (int_column) VALEUS (%(foo)s)""" args4 = {'foo': 1} TRN.add(sql3, args4) exp = [(sql1, args1), (sql2, None), (sql1, args3), (sql3, args4)] self.assertEqual(TRN._queries, exp) # Remove queries so __exit__ doesn't try to execute it TRN._queries = []
def test_index(self): with TRN: self.assertEqual(TRN.index, 0) TRN.add("SELECT 42") self.assertEqual(TRN.index, 1) sql = "INSERT INTO labcontrol.test_table (int_column) VALUES (%s)" args = [[1], [2], [3]] TRN.add(sql, args, many=True) self.assertEqual(TRN.index, 4) TRN.execute() self.assertEqual(TRN.index, 4) TRN.add(sql, args, many=True) self.assertEqual(TRN.index, 7) self.assertEqual(TRN.index, 0)
def test_execute_huge_transaction(self): with TRN: # Add a lot of inserts to the transaction sql = "INSERT INTO labcontrol.test_table (int_column) VALUES (%s)" for i in range(1000): TRN.add(sql, [i]) # Add some updates to the transaction sql = """UPDATE labcontrol.test_table SET bool_column = %s WHERE int_column = %s""" for i in range(500): TRN.add(sql, [False, i]) # Make the transaction fail with the last insert sql = """INSERT INTO labcontrol.table_to_make (the_trans_to_fail) VALUES (1)""" TRN.add(sql) with self.assertRaises(ValueError): TRN.execute() # make sure rollback correctly self._assert_sql_equal([])
def test_execute_return_many(self): with TRN: sql = """INSERT INTO labcontrol.test_table (str_column, int_column) VALUES (%s, %s) RETURNING str_column, int_column""" args = [['insert1', 1], ['insert2', 2], ['insert3', 3]] TRN.add(sql, args, many=True) sql = """UPDATE labcontrol.test_table SET bool_column = %s WHERE str_column = %s""" TRN.add(sql, [False, 'insert2']) sql = "SELECT * FROM labcontrol.test_table" TRN.add(sql) obs = TRN.execute() exp = [ [['insert1', 1]], # First query of the many query [['insert2', 2]], # Second query of the many query [['insert3', 3]], # Third query of the many query None, # Update query [ ['insert1', True, 1], # First result select ['insert3', True, 3], # Second result select ['insert2', False, 2] ] ] # Third result select self.assertEqual(obs, exp)
def get_samples(): with TRN: TRN.add("SELECT sample_id FROM qiita.study_sample") return TRN.execute_fetchflatten()
from labcontrol.db.sql_connection import TRN import logging with TRN: # although TRN has its own queue, use a local list to queue statements # until prerequisite queries are completed. statements = [] # # rename kappa_hyper_plus_kit_id in library_prep_shotgun_process # sql = """SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'library_prep_shotgun_process' AND column_name = 'kappa_hyper_plus_kit_id';""" logging.debug(sql) TRN.add(sql) # because we are COUNT()ing, execute_fetchflatten() can be expected to # return a single scalar value within a list. The [0] in this case is to # access the first (and only) scalar value within the results. result = TRN.execute_fetchflatten()[0] # for now, assume result will be either 0 or 1. # if 1, then the column 'kappa_hyper_plus_kit_id' needs # to be renamed to 'kapa_hyperplus_kit_id', otherwise # it has already been renamed. if result == 1: msg = ('labcontrol.library_prep_shotgun_process contains the ' 'column name kappa_hyper_plus_kit_id. Will rename to ' 'kapa_hyperplus_kit_id...') logging.debug(msg)