Merge pull request #3 from NYPL/postgres-pooling

Use postgres connection pooling

Author: aaronfriedman6

.github/workflows/deploy-production.yml

@@ -0,0 +1,57 @@
+name: Deploy to production if tagged as production release
+
+on:
+  release:
+    types: [ released ]
+
+jobs:
+  check_production_tag:
+    name: Check if the release is tagged as production
+    runs-on: ubuntu-latest
+    outputs:
+      has_production_tag: ${{ steps.check-production-tag.outputs.run_jobs }}
+    steps:
+      - name: check production tag ${{ github.ref }}
+        id: check-production-tag
+        run: |
+          if [[ ${{ github.ref }} =~ refs\/tags\/production ]]; then
+            echo "run_jobs=true" >> $GITHUB_OUTPUT
+          else
+            echo "run_jobs=false" >> $GITHUB_OUTPUT
+          fi
+  publish_production:
+    needs: [ check_production_tag ]
+    if: needs.check_production_tag.outputs.has_production_tag == 'true'
+    name: Publish package to pypi.org
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v3
+      
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+          cache: 'pip'
+          cache-dependency-path: 'pyproject.toml'
+
+      - name: Install pypa/build
+        run: >-
+          python -m
+          pip install
+          build
+          --user
+      
+      - name: Build a binary wheel and a source tarball
+        run: >-
+          python -m
+          build
+          --sdist
+          --wheel
+          --outdir dist/
+          .
+
+      - name: Publish distribution package to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## v0.0.3 - 2/10
+
+- Added GitHub Actions workflow for deploying to production
+- Switched PostgreSQLClient to use connection pooling
+
 ## v0.0.2 - 2/6
 
 - Added CODEOWNERS

README.md

@@ -5,14 +5,14 @@ This package contains common Python utility classes and functions.
 ## Classes
 * Pushing records to Kinesis
 * Setting and retrieving a resource in S3
+* 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 Redshift
 
 ## Functions
 * Reading a YAML config file and putting the contents in os.environ
-* Decrypting a value using KMS
 * Creating a logger in the appropriate format
 * Obfuscating a value using bcrypt
 

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "botocore>=1.29.5",
     "mysql-connector-python>=8.0.32",
     "psycopg>=3.1.0",
-    "psycopg-binary>=3.1.0",
+    "psycopg-pool>=3.1.6",
     "PyYAML>=6.0",
     "redshift-connector>=2.0.909",
     "requests>=2.28.1"

src/nypl_py_utils/classes/postgresql_client.py

@@ -1,32 +1,35 @@
 import psycopg
 
 from nypl_py_utils.functions.log_helper import create_log
+from psycopg_pool import ConnectionPool
 
 
 class PostgreSQLClient:
     """
     Client for managing connections to a PostgreSQL database (such as Sierra)
     """
 
-    def __init__(self, host, port, db_name, user, password):
+    def __init__(self, host, port, db_name, user, password, **kwargs):
         self.logger = create_log('postgresql_client')
-        self.conn = None
-        self.host = host
-        self.port = port
         self.db_name = db_name
-        self.user = user
-        self.password = password
+        self.timeout = kwargs.get('timeout', 300)
+
+        conn_info = 'postgresql://{user}:{password}@{host}:{port}/{db_name}'\
+            .format(user=user, password=password, host=host, port=port,
+                    db_name=db_name)
+        self.pool = ConnectionPool(
+            conn_info, open=False,
+            min_size=kwargs.get('min_size', 1),
+            max_size=kwargs.get('max_size', None))
 
     def connect(self):
-        """Connects to a PostgreSQL database using the given credentials"""
+        """
+        Opens the connection pool and connects to the given PostgreSQL database
+        min_size number of times.
+        """
         self.logger.info('Connecting to {} database'.format(self.db_name))
         try:
-            self.conn = psycopg.connect(
-                host=self.host,
-                port=self.port,
-                dbname=self.db_name,
-                user=self.user,
-                password=self.password)
+            self.pool.open(wait=True, timeout=self.timeout)
         except psycopg.Error as e:
             self.logger.error(
                 'Error connecting to {name} database: {error}'.format(
@@ -37,32 +40,33 @@ def connect(self):
 
     def execute_query(self, query):
         """
-        Executes an arbitrary query against the given database connection.
+        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.
         """
         self.logger.info('Querying {} database'.format(self.db_name))
         self.logger.debug('Executing query {}'.format(query))
-        try:
-            cursor = self.conn.cursor()
-            return cursor.execute(query).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()
+        with self.pool.connection() as conn:
+            try:
+                return conn.execute(query).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
 
     def close_connection(self):
-        """Closes the database connection"""
+        """Closes the connection pool"""
         self.logger.debug('Closing {} database connection'.format(
             self.db_name))
-        self.conn.close()
+        self.pool.close()
 
 
 class PostgreSQLClientError(Exception):

tests/test_postgresql_client.py

@@ -6,50 +6,59 @@
 class TestPostgreSQLClient:
 
     @pytest.fixture
-    def mock_pg_conn(self, mocker):
-        return mocker.patch('psycopg.connect')
-
-    @pytest.fixture
-    def test_instance(self):
+    def test_instance(self, mocker):
+        mocker.patch('psycopg_pool.ConnectionPool.open')
+        mocker.patch('psycopg_pool.ConnectionPool.close')
         return PostgreSQLClient('test_host', 'test_port', 'test_db_name',
                                 'test_user', 'test_password')
 
-    def test_connect(self, mock_pg_conn, test_instance):
-        test_instance.connect()
-        mock_pg_conn.assert_called_once_with(host='test_host',
-                                             port='test_port',
-                                             dbname='test_db_name',
-                                             user='test_user',
-                                             password='test_password')
+    def test_init(self, test_instance):
+        assert test_instance.pool.conninfo == (
+            'postgresql://test_user:test_password@test_host:test_port/' +
+            'test_db_name')
+        assert test_instance.pool._opened is False
+        assert test_instance.pool.min_size == 1
+        assert test_instance.pool.max_size == 1
 
-    def test_execute_query(self, mock_pg_conn, test_instance, mocker):
+    def test_init_with_kwargs(self):
+        test_instance = PostgreSQLClient(
+            'test_host', 'test_port', 'test_db_name', 'test_user',
+            'test_password', min_size=5, max_size=10)
+        assert test_instance.pool.conninfo == (
+            'postgresql://test_user:test_password@test_host:test_port/' +
+            'test_db_name')
+        assert test_instance.pool._opened is False
+        assert test_instance.pool.min_size == 5
+        assert test_instance.pool.max_size == 10
+
+    def test_connect(self, test_instance):
         test_instance.connect()
+        test_instance.pool.open.assert_called_once()
 
-        mock_cursor = mocker.MagicMock()
-        mock_cursor.execute.return_value = mock_cursor
-        mock_cursor.fetchall.return_value = [(1, 2, 3), ('a', 'b', 'c')]
-        test_instance.conn.cursor.return_value = mock_cursor
+    def test_connect_with_exception(self):
+        test_instance = PostgreSQLClient(
+            'test_host', 'test_port', 'test_db_name', 'test_user',
+            'test_password', timeout=1.0)
 
-        assert test_instance.execute_query(
-            'test query') == [(1, 2, 3), ('a', 'b', 'c')]
-        mock_cursor.execute.assert_called_once_with('test query')
-        mock_cursor.close.assert_called_once()
+        with pytest.raises(PostgreSQLClientError):
+            test_instance.connect()
 
-    def test_execute_query_with_exception(
-            self, mock_pg_conn, test_instance, mocker):
+    def test_execute_query(self, test_instance, mocker):
         test_instance.connect()
 
-        mock_cursor = mocker.MagicMock()
-        mock_cursor.execute.side_effect = Exception()
-        test_instance.conn.cursor.return_value = mock_cursor
+        mock_conn = mocker.MagicMock()
+        mock_conn.execute.side_effect = Exception()
+        mock_conn_context = mocker.MagicMock()
+        mock_conn_context.__enter__.return_value = mock_conn
+        mocker.patch('psycopg_pool.ConnectionPool.connection',
+                     return_value=mock_conn_context)
 
         with pytest.raises(PostgreSQLClientError):
             test_instance.execute_query('test query')
 
-        test_instance.conn.rollback.assert_called_once()
-        mock_cursor.close.assert_called_once()
+        mock_conn.rollback.assert_called_once()
 
-    def test_close_connection(self, mock_pg_conn, test_instance):
+    def test_close_connection(self, test_instance):
         test_instance.connect()
         test_instance.close_connection()
-        test_instance.conn.close.assert_called_once()
+        test_instance.pool.close.assert_called_once()