Merge branch 'main' into qa

Author: aaronfriedman6

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.0.1 - 4/3/23
+- Add transaction support to RedshiftClient
+
 ## v1.0.0 - 3/22/23
 - Improve Oauth2ApiClient token refresh and method responses
 - Create separate PostgreSQLClient and PostgreSQLPoolClient classes

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nypl_py_utils"
-version = "1.0.0"
+version = "1.0.1"
 authors = [
   { name="Aaron Friedman", email="aaronfriedman@nypl.org" },
 ]

src/nypl_py_utils/classes/redshift_client.py

@@ -72,6 +72,39 @@ def execute_query(self, query, dataframe=False):
         finally:
             cursor.close()
 
+    def execute_transaction(self, queries):
+        """
+        Executes a series of queries within a single transaction against the
+        given database connection. Assumes each of these queries is a write
+        query and so does not return anything.
+
+        Parameters
+        ----------
+        queries: list<tuple>
+            A list of tuples containing a query and the values to be used if
+            the query is parameterized (or None if it's not)
+        """
+        self.logger.info('Executing transaction against {} database'.format(
+            self.database))
+        try:
+            cursor = self.conn.cursor()
+            cursor.execute('BEGIN TRANSACTION;')
+            for query in queries:
+                self.logger.debug('Executing query {}'.format(query))
+                cursor.execute(query[0], query[1])
+            cursor.execute('END TRANSACTION;')
+            self.conn.commit()
+        except Exception as e:
+            self.conn.rollback()
+            self.logger.error(
+                ('Error executing {name} database transaction: {error}')
+                .format(name=self.database, error=e))
+            raise RedshiftClientError(
+                ('Error executing {name} database transaction: {error}')
+                .format(name=self.database, error=e)) from None
+        finally:
+            cursor.close()
+
     def close_connection(self):
         """Closes the database connection"""
         self.logger.debug('Closing {} database connection'.format(

tests/test_redshift_client.py

@@ -62,6 +62,43 @@ def test_execute_query_with_exception(
         test_instance.conn.rollback.assert_called_once()
         mock_cursor.close.assert_called_once()
 
+    def test_execute_transaction(self, mock_redshift_conn, test_instance,
+                                 mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        test_instance.conn.cursor.return_value = mock_cursor
+
+        test_instance.execute_transaction([('query 1', None),
+                                           ('query 2 %s %s', ('a', 1))])
+        mock_cursor.execute.assert_has_calls([
+            mocker.call('BEGIN TRANSACTION;'),
+            mocker.call('query 1', None),
+            mocker.call('query 2 %s %s', ('a', 1)),
+            mocker.call('END TRANSACTION;')])
+        test_instance.conn.commit.assert_called_once()
+        mock_cursor.close.assert_called_once()
+
+    def test_execute_transaction_with_exception(
+            self, mock_redshift_conn, test_instance, mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        mock_cursor.execute.side_effect = [None, None, Exception()]
+        test_instance.conn.cursor.return_value = mock_cursor
+
+        with pytest.raises(RedshiftClientError):
+            test_instance.execute_transaction(
+                [('query 1', None), ('query 2', None)])
+
+        mock_cursor.execute.assert_has_calls([
+            mocker.call('BEGIN TRANSACTION;'),
+            mocker.call('query 1', None),
+            mocker.call('query 2', None)])
+        test_instance.conn.commit.assert_not_called()
+        test_instance.conn.rollback.assert_called_once()
+        mock_cursor.close.assert_called_once()
+
     def test_close_connection(self, mock_redshift_conn, test_instance):
         test_instance.connect()
         test_instance.close_connection()