Add transaction support to RedshiftClient

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,38 @@ 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<str>
+            A list of the queries to execute in order
+        """
+        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)
+            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,41 @@ 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', 'query 2'])
+        mock_cursor.execute.assert_has_calls([
+            mocker.call('BEGIN TRANSACTION;'),
+            mocker.call('query 1'),
+            mocker.call('query 2'),
+            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', 'query 2'])
+
+        mock_cursor.execute.assert_has_calls([
+            mocker.call('BEGIN TRANSACTION;'),
+            mocker.call('query 1'),
+            mocker.call('query 2')])
+        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()