Merge pull request #11 from NYPL/auto-close-pg-connections

Separate pooling and regular PostgreSQL clients

Author: aaronfriedman6

.github/workflows/run-unit-tests.yml

@@ -28,7 +28,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install .
-          pip install '.[tests]'
+          pip install '.[development]'
 
       - name: Run linter and test suite
         run: |

CHANGELOG.md

@@ -1,11 +1,11 @@
 # Changelog
 
-## v0.0.9 - 3/16/23
+## v1.0.0 - 3/22/23
 - Improve Oauth2ApiClient token refresh and method responses
-
-## v0.0.8 - 3/3/23
-- Pass in all kwargs from PostgreSQLClient to ConnectionPool so that all ConnectionPool settings
-can be set from the wrapper
+- Create separate PostgreSQLClient and PostgreSQLPoolClient classes
+- Update PostgreSQL and MySQL clients to accept write queries implicitly
+- Update RedshiftClient to ensure SSL is being used
+- Separate dependencies to slim down package installation
 
 ## v0.0.7 - 3/1/23
 - Added Oauth2ApiClient for oauth2 authenticated calls to our Platform API and Sierra

README.md

@@ -8,6 +8,7 @@ This package contains common Python utility classes and functions.
 * Decrypting values with KMS
 * Encoding and decoding records using a given Avro schema
 * Connecting to and querying a MySQL database
+* Connecting to and querying a PostgreSQL database
 * Connecting to and querying a PostgreSQL database using a connection pool
 * Connecting to and querying Redshift
 * Making requests to the Oauth2 authenticated APIs such as NYPL Platform API and Sierra
@@ -25,11 +26,16 @@ python3 -m venv testenv
 source testenv/bin/activate
 pip install --upgrade pip
 pip install .
-pip install '.[tests]'
+pip install '.[development]'
 deactivate && source testenv/bin/activate
 ```
 
-Add any new dependencies required by code in the `nypl_py_utils` directory to the `dependencies` section of `pyproject.toml`. Add dependencies only required by code in the `tests` directory to the `[project.optional-dependencies]` section.
+## Managing dependencies
+In order to prevent dependency bloat, this package has no required dependencies. Instead, each class and helper file has its own optional dependency set. For instance, if an app needs to use the KMS client and the obfuscation helper, it should add `nypl-py-utils[kms-client, obfuscation-helper]` to the app's requirements. This way, only the required dependencies are installed.
+
+When a new client or helper file is created, a new optional dependency set should be added to `pyproject.toml`. The `development` dependency set, which includes all the dependencies required by all of the classes and tests, should also be updated.
+
+The optional dependency sets also give the developer the option to manually list out the dependencies of the clients rather than relying upon what the package thinks is required, which can be beneficial in certain circumstances. For instance, AWS lambda functions come with `boto3` and `botocore` pre-installed, so it's not necessary to include these (rather hefty) dependencies in the lambda deployment package.
 
 ### Troubleshooting
 If running `main.py` in this virtual environment produces the following error:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nypl_py_utils"
-version = "0.0.8"
+version = "1.0.0"
 authors = [
   { name="Aaron Friedman", email="aaronfriedman@nypl.org" },
 ]
@@ -16,30 +16,59 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
 ]
-dependencies = [
+dependencies = []
+
+[project.urls]
+"Homepage" = "https://github.com/NYPL/python-utils"
+"Bug Tracker" = "https://github.com/NYPL/python-utils/issues"
+
+[project.optional-dependencies]
+avro-encoder = [
     "avro>=1.11.1",
-    "bcrypt>=4.0.1",
+    "requests>=2.28.1"
+]
+kinesis-client = [
     "boto3>=1.26.5",
-    "botocore>=1.29.5",
-    "mysql-connector-python>=8.0.32",
+    "botocore>=1.29.5"
+]
+kms-client = [
+    "boto3>=1.26.5",
+    "botocore>=1.29.5"
+]
+mysql-client = [
+    "mysql-connector-python>=8.0.32"
+]
+oauth2-api-client = [
     "oauthlib>=3.2.2",
-    "psycopg[binary,pool]>=3.1.6",
-    "PyYAML>=6.0",
-    "redshift-connector>=2.0.909",
-    "requests>=2.28.1",
     "requests_oauthlib>=1.3.1"
 ]
-
-[project.optional-dependencies]
-tests = [
+postgresql-client = [
+    "psycopg[binary]>=3.1.6"
+]
+postgresql-pool-client = [
+    "psycopg[binary,pool]>=3.1.6"
+]
+redshift-client = [
+    "botocore>=1.29.5",
+    "redshift-connector>=2.0.909"
+]
+s3-client = [
+    "boto3>=1.26.5",
+    "botocore>=1.29.5"
+]
+config-helper = [
+    "nypl_py_utils[kms-client]",
+    "PyYAML>=6.0"
+]
+obfuscation-helper = [
+    "bcrypt>=4.0.1"
+]
+development = [
+    "nypl_py_utils[avro-encoder,kinesis-client,kms-client,mysql-client,oauth2-api-client,postgresql-client,postgresql-pool-client,redshift-client,s3-client,config-helper,obfuscation-helper]",
     "flake8>=6.0.0",
     "freezegun>=1.2.2",
     "mock>=4.0.3",
     "pytest>=7.2.0",
     "pytest-mock>=3.10.0",
     "requests-mock>=1.10.0"
 ]
-
-[project.urls]
-"Homepage" = "https://github.com/NYPL/python-utils"
-"Bug Tracker" = "https://github.com/NYPL/python-utils/issues"

src/nypl_py_utils/__init__.py

@@ -1,8 +0,0 @@
-from .classes.avro_encoder import AvroEncoder, AvroEncoderError  # noqa
-from .classes.kinesis_client import KinesisClient, KinesisClientError  # noqa
-from .classes.kms_client import KmsClient, KmsClientError # noqa
-from .classes.mysql_client import MySQLClient, MySQLClientError  # noqa
-from .classes.oauth2_api_client import Oauth2ApiClient, Oauth2ApiClientError  # noqa
-from .classes.postgresql_client import PostgreSQLClient, PostgreSQLClientError  # noqa
-from .classes.redshift_client import RedshiftClient, RedshiftClientError  # noqa
-from .classes.s3_client import S3Client, S3ClientError  # noqa

src/nypl_py_utils/classes/mysql_client.py

@@ -15,16 +15,28 @@ def __init__(self, host, port, database, user, password):
         self.user = user
         self.password = password
 
-    def connect(self):
-        """Connects to a MySQL database using the given credentials"""
+    def connect(self, **kwargs):
+        """
+        Connects to a MySQL database using the given credentials.
+
+        Keyword args can be passed into the connection to set certain options.
+        All possible arguments can be found here:
+        https://dev.mysql.com/doc/connector-python/en/connector-python-connectargs.html.
+
+        Common arguments include:
+            autocommit: bool
+                Whether to automatically commit each query rather than running
+                them as part of a transaction. By default False.
+        """
         self.logger.info('Connecting to {} database'.format(self.database))
         try:
             self.conn = mysql.connector.connect(
                 host=self.host,
                 port=self.port,
                 database=self.database,
                 user=self.user,
-                password=self.password)
+                password=self.password,
+                **kwargs)
         except mysql.connector.Error as e:
             self.logger.error(
                 'Error connecting to {name} database: {error}'.format(
@@ -33,37 +45,38 @@ def connect(self):
                 'Error connecting to {name} database: {error}'.format(
                     name=self.database, error=e)) from None
 
-    def execute_query(self, query, is_write_query=False, query_params=None,
-                      dictionary=False):
+    def execute_query(self, query, query_params=None, **kwargs):
         """
         Executes an arbitrary query against the given database connection.
 
         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.
+        kwargs:
+            All possible arguments can be found here:
+            https://dev.mysql.com/doc/connector-python/en/connector-python-api-mysqlconnection-cursor.html.
+
+            Common arguments include:
+                dictionary: bool
+                    Whether the data will be returned as a dictionary. Defaults
+                    to False, meaning 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.
+            None if the cursor has nothing to return. A list of either tuples
+            or dictionaries (based on the dictionary input) if there's
+            something to return (even if the result set is empty).
         """
         self.logger.info('Querying {} database'.format(self.database))
         self.logger.debug('Executing query {}'.format(query))
         try:
-            cursor = self.conn.cursor(dictionary=dictionary)
+            cursor = self.conn.cursor(**kwargs)
             cursor.execute(query, query_params)
-            if is_write_query:
+            if cursor.description is None:
                 self.conn.commit()
                 return None
             else:

src/nypl_py_utils/classes/postgresql_client.py

@@ -1,40 +1,41 @@
 import psycopg
 
 from nypl_py_utils.functions.log_helper import create_log
-from psycopg.rows import tuple_row
-from psycopg_pool import ConnectionPool
 
 
 class PostgreSQLClient:
-    """
-    Client for managing connections to a PostgreSQL database (such as Sierra)
-    """
+    """Client for managing individual connections to a PostgreSQL database"""
 
-    def __init__(self, host, port, db_name, user, password, **kwargs):
+    def __init__(self, host, port, db_name, user, password):
         self.logger = create_log('postgresql_client')
-        self.db_name = db_name
-        self.timeout = kwargs.get('timeout', 300)
-
+        self.conn = None
         self.conn_info = ('postgresql://{user}:{password}@{host}:{port}/'
                           '{db_name}').format(user=user, password=password,
                                               host=host, port=port,
                                               db_name=db_name)
-        self.kwargs = kwargs
-        self.kwargs['min_size'] = kwargs.get('min_size', 0)
-        self.kwargs['max_size'] = kwargs.get('max_size', 1)
-        self.pool = ConnectionPool(self.conn_info, open=False, **self.kwargs)
 
-    def connect(self):
+        self.db_name = db_name
+
+    def connect(self, **kwargs):
         """
-        Opens the connection pool and connects to the given PostgreSQL database
-        min_size number of times.
+        Connects to a PostgreSQL database using the given credentials.
+
+        Keyword args can be passed into the connection to set certain options.
+        All possible arguments can be found here:
+        https://www.psycopg.org/psycopg3/docs/api/connections.html#psycopg.Connection.connect.
+
+        Common arguments include:
+            autocommit: bool
+                Whether to automatically commit each query rather than running
+                them as part of a transaction. By default False.
+            row_factory: RowFactory
+                A psycopg RowFactory that determines how the data will be
+                returned. Defaults to tuple_row, which returns the rows as a
+                list of tuples.
         """
         self.logger.info('Connecting to {} database'.format(self.db_name))
         try:
-            if self.pool is None:
-                self.pool = ConnectionPool(
-                    self.conn_info, open=False, **self.kwargs)
-            self.pool.open(wait=True, timeout=self.timeout)
+            self.conn = psycopg.connect(self.conn_info, **kwargs)
         except psycopg.Error as e:
             self.logger.error(
                 'Error connecting to {name} database: {error}'.format(
@@ -43,59 +44,52 @@ def connect(self):
                 'Error connecting to {name} database: {error}'.format(
                     name=self.db_name, error=e)) from None
 
-    def execute_query(self, query, is_write_query=False, query_params=None,
-                      row_factory=tuple_row):
+    def execute_query(self, query, query_params=None, **kwargs):
         """
-        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.
+        Executes an arbitrary query against the given database connection.
 
         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.
+        kwargs:
+            All possible arguments can be found here:
+            https://www.psycopg.org/psycopg3/docs/api/cursors.html#psycopg.Cursor.execute
 
         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.
+            None if the cursor has nothing to return. Some type of sequence
+            based on the connection's row_factory if there's something to
+            return (even if the result set is empty).
         """
         self.logger.info('Querying {} database'.format(self.db_name))
         self.logger.debug('Executing query {}'.format(query))
-        with self.pool.connection() as conn:
-            try:
-                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(
-                    ('Error executing {name} database query \'{query}\': '
-                     '{error}').format(
-                        name=self.db_name, query=query, error=e))
-                raise PostgreSQLClientError(
-                    ('Error executing {name} database query \'{query}\': '
-                     '{error}').format(
-                        name=self.db_name, query=query, error=e)) from None
+        try:
+            cursor = self.conn.cursor()
+            cursor.execute(query, query_params, **kwargs)
+            self.conn.commit()
+            return None if cursor.description is None else cursor.fetchall()
+        except Exception as e:
+            self.conn.rollback()
+            self.logger.error(
+                ('Error executing {name} database query \'{query}\': '
+                    '{error}').format(
+                    name=self.db_name, query=query, error=e))
+            raise PostgreSQLClientError(
+                ('Error executing {name} database query \'{query}\': '
+                    '{error}').format(
+                    name=self.db_name, query=query, error=e)) from None
+        finally:
+            cursor.close()
 
     def close_connection(self):
-        """Closes the connection pool"""
+        """Closes the database connection"""
         self.logger.debug('Closing {} database connection'.format(
             self.db_name))
-        self.pool.close()
-        self.pool = None
+        self.conn.close()
 
 
 class PostgreSQLClientError(Exception):