Merge pull request #7 from NYPL/update-db-clients

Update db clients

Author: aaronfriedman6

CHANGELOG.md

@@ -1,21 +1,25 @@
 # Changelog
 
-## v0.0.4 - 2/13
+## v0.0.5 - 2/22/23
+- Support write queries to PostgreSQL and MySQL databases
+- Support different return formats when querying PostgreSQL, MySQL, and Redshift databases
+
+## v0.0.4 - 2/13/23
 - In PostgreSQLClient, allow reconnecting after `close_connection` has been called
 - Updated README with deployment information
 
-## v0.0.3 - 2/10
+## v0.0.3 - 2/10/23
 
 - Added GitHub Actions workflow for deploying to production
 - Switched PostgreSQLClient to use connection pooling
 
-## v0.0.2 - 2/6
+## v0.0.2 - 2/6/23
 
 - Added CODEOWNERS
 - Added GitHub Actions workflows for running tests and deploying to QA
 - Added tests for helper functions
 - Updated Avro encoder to avoid dependency on pandas
 
-## v0.0.1 - 1/26
+## v0.0.1 - 1/26/23
 
 Initial version. Includes the `avro_encoder`, `kinesis_client`, `mysql_client`, `postgresql_client`, `redshift_client`, and `s3_client` classes as well as the `config_helper`, `kms_helper`, `log_helper`, and `obfuscation_helper` functions.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nypl_py_utils"
-version = "0.0.5"
+version = "0.0.6"
 authors = [
   { name="Aaron Friedman", email="aaronfriedman@nypl.org" },
 ]
@@ -22,8 +22,7 @@ dependencies = [
     "boto3>=1.26.5",
     "botocore>=1.29.5",
     "mysql-connector-python>=8.0.32",
-    "psycopg>=3.1.0",
-    "psycopg-pool>=3.1.6",
+    "psycopg[binary,pool]>=3.1.6",
     "PyYAML>=6.0",
     "redshift-connector>=2.0.909",
     "requests>=2.28.1"

src/nypl_py_utils/classes/mysql_client.py

@@ -33,19 +33,41 @@ def connect(self):
                 'Error connecting to {name} database: {error}'.format(
                     name=self.database, error=e)) from None
 
-    def execute_query(self, query):
+    def execute_query(self, query, is_write_query=False, query_params=None,
+                      dictionary=False):
         """
         Executes an arbitrary query against the given database connection.
 
-        Returns a sequence of tuples representing the rows returned by the
-        query.
+        Parameters
+        ----------
+        query: str
+            The query to execute
+        is_write_query: bool, optional
+            Whether or not the query is writing to the database, in which case
+            the transaction needs to be committed and None should be returned
+        query_params: sequence, optional
+            The values to be used in a parameterized query
+        dictionary: bool, optional
+            Whether the data will be returned as a dictionary. 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 either tuples or
+            dictionaries (based on the dictionary input) if is_write_query is
+            False.
         """
         self.logger.info('Querying {} database'.format(self.database))
         self.logger.debug('Executing query {}'.format(query))
         try:
-            cursor = self.conn.cursor()
-            cursor.execute(query)
-            return cursor.fetchall()
+            cursor = self.conn.cursor(dictionary=dictionary)
+            cursor.execute(query, query_params)
+            if is_write_query:
+                self.conn.commit()
+                return None
+            else:
+                return cursor.fetchall()
         except Exception as e:
             self.conn.rollback()
             self.logger.error(

src/nypl_py_utils/classes/postgresql_client.py

@@ -1,6 +1,7 @@
 import psycopg
 
 from nypl_py_utils.functions.log_helper import create_log
+from psycopg.rows import tuple_row
 from psycopg_pool import ConnectionPool
 
 
@@ -45,19 +46,43 @@ def connect(self):
                 'Error connecting to {name} database: {error}'.format(
                     name=self.db_name, error=e)) from None
 
-    def execute_query(self, query):
+    def execute_query(self, query, is_write_query=False, query_params=None,
+                      row_factory=tuple_row):
         """
         Requests a connection from the pool and uses it to execute an arbitrary
         query. After the query is complete, returns the connection to the pool.
 
-        Returns a sequence of tuples representing the rows returned by the
-        query.
+        Parameters
+        ----------
+        query: str
+            The query to execute
+        is_write_query: bool, optional
+            Whether or not the query is writing to the database, in which case
+            the transaction needs to be committed and None should be returned
+        query_params: sequence, optional
+            The values to be used in a parameterized query
+        row_factory: RowFactory, optional
+            A psycopg RowFactory that determines how the data will be returned.
+            Defaults to tuple_row, which returns the rows as a list of tuples.
+
+        Returns
+        -------
+        None or sequence
+            None if is_write_query is True. Some type of sequence based on
+            the row_factory input if is_write_query is False.
         """
         self.logger.info('Querying {} database'.format(self.db_name))
         self.logger.debug('Executing query {}'.format(query))
+        self.pool.check()
         with self.pool.connection() as conn:
             try:
-                return conn.execute(query).fetchall()
+                conn.row_factory = row_factory
+                cursor = conn.execute(query, query_params)
+                if is_write_query:
+                    conn.commit()
+                    return None
+                else:
+                    return cursor.fetchall()
             except Exception as e:
                 conn.rollback()
                 self.logger.error(

src/nypl_py_utils/classes/redshift_client.py

@@ -32,19 +32,34 @@ def connect(self):
                 'Error connecting to {name} database: {error}'.format(
                     name=self.database, error=e)) from None
 
-    def execute_query(self, query):
+    def execute_query(self, query, dataframe=False):
         """
         Executes an arbitrary query against the given database connection.
 
-        Returns a sequence of tuples representing the rows returned by the
-        query.
+        Parameters
+        ----------
+        query: str
+            The query to execute
+        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.
         """
         self.logger.info('Querying {} database'.format(self.database))
         self.logger.debug('Executing query {}'.format(query))
         try:
             cursor = self.conn.cursor()
             cursor.execute(query)
-            return cursor.fetchall()
+            if dataframe:
+                return cursor.fetch_dataframe()
+            else:
+                return cursor.fetchall()
         except Exception as e:
             self.conn.rollback()
             self.logger.error(

tests/test_mysql_client.py

@@ -22,7 +22,7 @@ def test_connect(self, mock_mysql_conn, test_instance):
                                                 user='test_user',
                                                 password='test_password')
 
-    def test_execute_query(self, mock_mysql_conn, test_instance, mocker):
+    def test_execute_read_query(self, mock_mysql_conn, test_instance, mocker):
         test_instance.connect()
 
         mock_cursor = mocker.MagicMock()
@@ -31,7 +31,56 @@ def test_execute_query(self, mock_mysql_conn, test_instance, mocker):
 
         assert test_instance.execute_query(
             'test query') == [(1, 2, 3), ('a', 'b', 'c')]
-        mock_cursor.execute.assert_called_once_with('test query')
+        test_instance.conn.cursor.called_once_with(dictionary=False)
+        mock_cursor.execute.assert_called_once_with('test query', None)
+        mock_cursor.close.assert_called_once()
+
+    def test_execute_dictionary_read_query(self, mock_mysql_conn,
+                                           test_instance, mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        mock_cursor.fetchall.return_value = [
+            {'col1': 1, 'col2': 'a'},
+            {'col1': 2, 'col2': 'b'},
+            {'col1': 3, 'col2': 'c'}]
+        test_instance.conn.cursor.return_value = mock_cursor
+
+        assert test_instance.execute_query(
+            'test query', dictionary=True) == [{'col1': 1, 'col2': 'a'},
+                                               {'col1': 2, 'col2': 'b'},
+                                               {'col1': 3, 'col2': 'c'}]
+        test_instance.conn.cursor.called_once_with(dictionary=True)
+        mock_cursor.execute.assert_called_once_with('test query', None)
+        mock_cursor.close.assert_called_once()
+
+    def test_execute_write_query(self, mock_mysql_conn, test_instance, mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        test_instance.conn.cursor.return_value = mock_cursor
+
+        assert test_instance.execute_query(
+            'test query', is_write_query=True) is None
+        test_instance.conn.cursor.called_once_with(dictionary=False)
+        mock_cursor.execute.assert_called_once_with('test query', None)
+        test_instance.conn.commit.called_once()
+        mock_cursor.close.assert_called_once()
+
+    def test_execute_write_query_with_params(self, mock_mysql_conn,
+                                             test_instance, mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        test_instance.conn.cursor.return_value = mock_cursor
+
+        assert test_instance.execute_query(
+            'test query %s %s', is_write_query=True,
+            query_params=('a', 1)) is None
+        test_instance.conn.cursor.called_once_with(dictionary=False)
+        mock_cursor.execute.assert_called_once_with('test query %s %s',
+                                                    ('a', 1))
+        test_instance.conn.commit.called_once()
         mock_cursor.close.assert_called_once()
 
     def test_execute_query_with_exception(

tests/test_postgresql_client.py

@@ -43,7 +43,54 @@ def test_connect_with_exception(self):
         with pytest.raises(PostgreSQLClientError):
             test_instance.connect()
 
-    def test_execute_query(self, test_instance, mocker):
+    def test_execute_read_query(self, test_instance, mocker):
+        test_instance.connect()
+
+        mock_cursor = mocker.MagicMock()
+        mock_cursor.fetchall.return_value = [(1, 2, 3), ('a', 'b', 'c')]
+        mock_conn = mocker.MagicMock()
+        mock_conn.execute.return_value = mock_cursor
+        mock_conn_context = mocker.MagicMock()
+        mock_conn_context.__enter__.return_value = mock_conn
+        mocker.patch('psycopg_pool.ConnectionPool.connection',
+                     return_value=mock_conn_context)
+
+        assert test_instance.execute_query(
+            'test query') == [(1, 2, 3), ('a', 'b', 'c')]
+        mock_conn.execute.assert_called_once_with('test query', None)
+        mock_cursor.fetchall.assert_called_once()
+
+    def test_execute_write_query(self, test_instance, mocker):
+        test_instance.connect()
+
+        mock_conn = mocker.MagicMock()
+        mock_conn_context = mocker.MagicMock()
+        mock_conn_context.__enter__.return_value = mock_conn
+        mocker.patch('psycopg_pool.ConnectionPool.connection',
+                     return_value=mock_conn_context)
+
+        assert test_instance.execute_query(
+            'test query', is_write_query=True) is None
+        mock_conn.execute.assert_called_once_with('test query', None)
+        mock_conn.commit.assert_called_once()
+
+    def test_execute_write_query_with_params(self, test_instance, mocker):
+        test_instance.connect()
+
+        mock_conn = mocker.MagicMock()
+        mock_conn_context = mocker.MagicMock()
+        mock_conn_context.__enter__.return_value = mock_conn
+        mocker.patch('psycopg_pool.ConnectionPool.connection',
+                     return_value=mock_conn_context)
+
+        assert test_instance.execute_query(
+            'test query %s %s', is_write_query=True,
+            query_params=('a', 1)) is None
+        mock_conn.execute.assert_called_once_with('test query %s %s',
+                                                  ('a', 1))
+        mock_conn.commit.assert_called_once()
+
+    def test_execute_query_with_exception(self, test_instance, mocker):
         test_instance.connect()
 
         mock_conn = mocker.MagicMock()

tests/test_redshift_client.py

@@ -31,6 +31,19 @@ def test_execute_query(self, mock_redshift_conn, test_instance, mocker):
         assert test_instance.execute_query(
             'test query') == [[1, 2, 3], ['a', 'b', 'c']]
         mock_cursor.execute.assert_called_once_with('test query')
+        mock_cursor.fetchall.assert_called_once()
+        mock_cursor.close.assert_called_once()
+
+    def test_execute_dataframe_query(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_query('test query', dataframe=True)
+        mock_cursor.execute.assert_called_once_with('test query')
+        mock_cursor.fetch_dataframe.assert_called_once()
         mock_cursor.close.assert_called_once()
 
     def test_execute_query_with_exception(