Merge pull request #28 from NYPL/insert-many

Add executemany functionality to RedshiftClient

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
 

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="[email protected]" },
 ]

src/nypl_py_utils/classes/redshift_client.py

@@ -35,22 +35,21 @@ 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.
 
         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<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)
+            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:

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()