From cad20cbb90794fc13a8e50352b60b4e4f466b135 Mon Sep 17 00:00:00 2001 From: aaronfriedman Date: Thu, 6 Jun 2024 16:09:55 -0400 Subject: [PATCH] Add executemany functionality to RedshiftClient In addition to unit tests, tested locally that this works as expected --- CHANGELOG.md | 3 +++ pyproject.toml | 2 +- src/nypl_py_utils/classes/redshift_client.py | 24 ++++++++++++++------ tests/test_redshift_client.py | 22 ++++++++++++++++++ 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a565708..2f1f08d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,7 @@ # Changelog +## v1.1.5 6/6/24 +- Use executemany instead of execute when appropriate in RedshiftClient.execute_transaction + ## v1.1.4 3/14/24 - Fix bug with oauth2 requests after token refresh diff --git a/pyproject.toml b/pyproject.toml index 676394f..6939787 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "nypl_py_utils" -version = "1.1.4" +version = "1.1.5" authors = [ { name="Aaron Friedman", email="aaronfriedman@nypl.org" }, ] diff --git a/src/nypl_py_utils/classes/redshift_client.py b/src/nypl_py_utils/classes/redshift_client.py index 61b1275..17c4558 100644 --- a/src/nypl_py_utils/classes/redshift_client.py +++ b/src/nypl_py_utils/classes/redshift_client.py @@ -35,12 +35,12 @@ def connect(self): def execute_query(self, query, dataframe=False): """ - Executes an arbitrary query against the given database connection. + Executes an arbitrary read query against the given database connection. Parameters ---------- query: str - The query to execute + The query to execute, assumed to be a read query dataframe: bool, optional Whether the data will be returned as a pandas DataFrame. Defaults to False, which means the data is returned as a list of tuples. @@ -48,9 +48,8 @@ def execute_query(self, query, dataframe=False): Returns ------- None or sequence - None if is_write_query is True. A list of tuples or a pandas - DataFrame (based on the dataframe input) if is_write_query is - False. + A list of tuples or a pandas DataFrame (based on the `dataframe` + input) """ self.logger.info('Querying {} database'.format(self.database)) self.logger.debug('Executing query {}'.format(query)) @@ -82,7 +81,12 @@ def execute_transaction(self, queries): ---------- queries: list A list of tuples containing a query and the values to be used if - the query is parameterized (or None if it's not) + the query is parameterized (or None if it's not). The values can + be for a single insert query -- e.g. execute_transaction( + "INSERT INTO x VALUES (%s, %s)", (1, "a")) + or for multiple -- e.g execute_transaction( + "INSERT INTO x VALUES (%s, %s)", [(1, "a"), (2, "b")]) + """ self.logger.info('Executing transaction against {} database'.format( self.database)) @@ -91,7 +95,13 @@ def execute_transaction(self, queries): cursor.execute('BEGIN TRANSACTION;') for query in queries: self.logger.debug('Executing query {}'.format(query)) - cursor.execute(query[0], query[1]) + if query[1] is not None and all( + isinstance(el, tuple) or isinstance(el, list) + for el in query[1] + ): + cursor.executemany(query[0], query[1]) + else: + cursor.execute(query[0], query[1]) cursor.execute('END TRANSACTION;') self.conn.commit() except Exception as e: diff --git a/tests/test_redshift_client.py b/tests/test_redshift_client.py index a086b92..7d6219d 100644 --- a/tests/test_redshift_client.py +++ b/tests/test_redshift_client.py @@ -76,6 +76,28 @@ def test_execute_transaction(self, mock_redshift_conn, test_instance, mocker.call('query 1', None), mocker.call('query 2 %s %s', ('a', 1)), mocker.call('END TRANSACTION;')]) + mock_cursor.executemany.assert_not_called() + test_instance.conn.commit.assert_called_once() + mock_cursor.close.assert_called_once() + + def test_execute_transaction_with_many(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', (None, 1)), + ('query 3 %s %s', [(None, 10), ('b', 20)]), ('query 4', None)]) + mock_cursor.execute.assert_has_calls([ + mocker.call('BEGIN TRANSACTION;'), + mocker.call('query 1', None), + mocker.call('query 2 %s %s', (None, 1)), + mocker.call('query 4', None), + mocker.call('END TRANSACTION;')]) + mock_cursor.executemany.assert_called_once_with( + 'query 3 %s %s', [(None, 10), ('b', 20)]) test_instance.conn.commit.assert_called_once() mock_cursor.close.assert_called_once()