Introduce DBUtils PooledDB db connection pooling (#3)

* refactored database.db to app.db to simplify project structure

* implement db connection pooling with DBUtils PooledDB (pooled_db)

* minor fix in Handler: only close db connection if it was initialized

* minor fix for previous commit

* start explicit transaction with db.begin() in Cleanup, as required by DBUtils PooledDB

* moved logger and db cleanup code into destructor of Handler
made DBUtils PooledDB params configurable via .env vars

* added db connection pooling configuration notes to README
getting ready for v0.6.0 release

Author: onlime

.env.example

@@ -7,6 +7,13 @@ DB_USER=policyd-rate-guard
 DB_PASSWORD=Example1234
 DB_DATABASE=policyd-rate-guard
 
+## DB connection pool configuration
+## See https://webwareforpython.github.io/DBUtils/main.html#pooleddb-pooled-db
+# DB_POOL_MINCACHED=0
+# DB_POOL_MAXCACHED=10
+# DB_POOL_MAXSHARED=10
+# DB_POOL_MAXUSAGE=10000
+
 ## PolicydRateGuard configuration
 # QUOTA=1000 # The default quota for a user (default: 1000)
 # ACTION_TEXT_BLOCKED="Rate limit reached, retry later." # Here you can put a custom message to be shown to a user who is over the ratelimit.

CHANGELOG.md

@@ -1,5 +1,21 @@
 # CHANGELOG
 
+## [v0.6.0](https://github.com/onlime/policyd-rate-guard/releases/tag/v0.6.0) (2023-09-01)
+
+**Improved:**
+
+- Improved performance and stability by introducing database connection pooling, using [DBUtils PooledDB (pooled_db)](https://webwareforpython.github.io/DBUtils/main.html#pooleddb-pooled-db)
+- Moved logger and db cleanup code into destructor of `Handler`.
+- Refactored `database.db` to `app.db` to simplify project structure.
+
+**Added:**
+
+- Added environment variables `DB_POOL_MINCACHED`, `DB_POOL_MAXCACHED`, `DB_POOL_MAXSHARED`, `DB_POOL_MAXUSAGE` for db connection pooling fine-tuning.
+
+**Fixed:**
+
+- Fix `Lost connection to MySQL server during query ` and ``AttributeError: 'NoneType' object has no attribute 'read'` (on db cursor) connectivity issues by introducing connection pooling.
+
 ## [v0.5.1](https://github.com/onlime/policyd-rate-guard/releases/tag/v0.5.1) (2023-08-30)
 
 **Added:**
@@ -9,7 +25,7 @@
 
 **Fixed:**
 
-- Fix `AttributeError: 'NoneType' object has no attribute 'read' in Ratelimit.find()` edge case where `sasl_username` was set in Postfix DATA but empty. We now bail out early if `sasl_username` either does not exist or is empty.
+- Fix edge case where `sasl_username` was set in Postfix DATA but empty. We now bail out early if `sasl_username` either does not exist or is empty.
 
 ## [v0.5.0](https://github.com/onlime/policyd-rate-guard/releases/tag/v0.5.0) (2023-08-29)
 

README.md

@@ -17,10 +17,11 @@ Actually, PolicydRateGuard is just a super simple Postfix policy daemon with onl
 But let me name some features that make it stand out from other solutions:
 
 - **Super easy Postfix integration** using `check_policy_service` in `smtpd_data_restrictions`
-- Set **individual sender (SASL username) quotas**
+- **Tuned for high performance**, using network or unix sockets, threading, and db connection pooling.
+- Set **individual sender (SASL username) quotas**, which can be both persistent or only for the current time period.
 - Limit senders to **number of recipients** per time period
 - Automatically fills `ratelimits` table with new senders (SASL username) upon first email sent
-- Set your own time period (usually 24hrs) by resetting the counters via Systemd cleanup timer (or cronjob)
+- Set your own **time period (usually 24hrs)** by resetting the counters via Systemd cleanup timer (or cronjob)
 - Continues to raise counters (`msg_counter`, `rcpt_counter`) even in over quota state, so you know if a sender keeps retrying/spamming.
 - Keeps totals of all messages/recipients sent for each sender (SASL username)
 - Stores both **message and recipient counters** in database (`ratelimits` table)
@@ -29,9 +30,10 @@ But let me name some features that make it stand out from other solutions:
 - **Maximum failure safety:** On any unexpected exception, the daemon still replies with a `DUNNO` action, so that the mail is not getting rejected by Postfix. This is done both on Postfix integration side and application exception handling side.
 - **Block action message** `"Rate limit reached, retry later."` can be configured.
 - Lots of configuration params via a simple `.env` 
-- **Tuned for high performance**, using network or unix sockets, and threading.
 - **Secure setup**, nothing running under `root`, only on `postfix` user.
-- A super slick minimal codebase with **only a few dependencies** ([PyMySQL](https://pypi.org/project/pymysql/), [python-dotenv](https://pypi.org/project/python-dotenv/), [yoyo-migrations](https://pypi.org/project/yoyo-migrations/)), using Python virtual environment for easy `pip` install. PyMySQL is a pure-Python MySQL client library, so you won't have any trouble on any future major system upgrades.
+- A multi-threaded app that uses [DBUtils PooledDB (pooled_db)](https://github.com/WebwareForPython/DBUtils) for **robust and efficient DB connection handling**.
+- Can be used with any [DB-API 2 (PEP 249)](https://peps.python.org/pep-0249/) conformant database adapter (currently supported: PyMySQL, sqlite3)
+- A super slick minimal codebase with **only a few dependencies** ([PyMySQL](https://pypi.org/project/pymysql/), [DBUtils](https://webwareforpython.github.io/DBUtils/), [python-dotenv](https://pypi.org/project/python-dotenv/), [yoyo-migrations](https://pypi.org/project/yoyo-migrations/)), using Python virtual environment for easy `pip` install. PyMySQL is a pure-Python MySQL client library, so you won't have any trouble on any future major system upgrades.
 - Provides an Ansible Galaxy role [`onlime.policyd_rate_guard`](https://galaxy.ansible.com/onlime/policyd_rate_guard) for easy installation on a Debian mailserver.
 - A **well maintained** project, as it is in active use at [Onlime GmbH](https://www.onlime.ch/), a Swiss webhoster with a rock-solid mailserver architecture.
 
@@ -145,11 +147,12 @@ smtpd_data_restrictions =
         permit 
 ```
 
-> **IMPORTANT:** We strongly recommend the advanced policy client configuration (supported since Postfix 3.0), using above syntax with **default action `DUNNO`**, instead of just using `check_policy_service inet:127.0.0.1:10033`.
->
-> It ensures that if PolicydRateGuard becomes unavailable for any reason, Postfix will ignore it and keep accepting mail as if the rule was not there. PolicydRateGuard should be considered a "non-critical" policy service and you should use some monitoring solution to ensure it is always running as expected.
+**IMPORTANT:** We strongly recommend the advanced policy client configuration (supported since Postfix 3.0), using above syntax with **default action `DUNNO`**, instead of just using `check_policy_service inet:127.0.0.1:10033`.
+
+It ensures that if PolicydRateGuard becomes unavailable for any reason, Postfix will ignore it and keep accepting mail as if the rule was not there. PolicydRateGuard should be considered a "non-critical" policy service and you should use some monitoring solution to ensure it is always running as expected.
 
-> **NOTE:** You may use `unix:rateguard/policyd` instead of `inet:127.0.0.1:10033` if you have configured PolicydRateGuard to use a unix socket (`SOCKET="/var/spool/postfix/rateguard/policyd"` environment variable).
+> [!NOTE]
+> You may use `unix:rateguard/policyd` instead of `inet:127.0.0.1:10033` if you have configured PolicydRateGuard to use a unix socket (`SOCKET="/var/spool/postfix/rateguard/policyd"` environment variable).
 
 Make sure to reload Postfix after this change:
 
@@ -225,13 +228,21 @@ PolicydRateGuard can be fully configured through environment variables in `.env`
 - `SENTRY_ENVIRONMENT`
   Sentry environment. Suggested values: `dev` or `prod`, but can be any custom string. Defaults to `dev`.
 
+You may also tune the database connection pooling by modifying the following environment variables (defaults are fine for most environments, and you'll find e detailed description in the [DBUtils PooledDB](https://webwareforpython.github.io/DBUtils/main.html#pooleddb-pooled-db-1) usage docs):
+
+- `DB_POOL_MINCACHED` (default: `0`)
+- `DB_POOL_MAXCACHED` (default: `10`
+- `DB_POOL_MAXSHARED` (default: `10`)
+- `DB_POOL_MAXUSAGE`  (default: `10000`)
+
 For production, we recommend to start by copying `.env.example` and then fine-tune your `.env`:
 
 ```bash
 $ cp .env.example .env
 ```
 
-> **NOTE:** Minimally, you should set `DB_PASSWORD`, and maybe enable `SYSLOG` logging. For all the other config params it's usually fine to stick with the defaults.
+> [!NOTE]
+> Minimally, you should set `DB_PASSWORD`, and maybe enable `SYSLOG` logging. For all the other config params it's usually fine to stick with the defaults.
 
 ## Development 👩‍💻
 
@@ -337,6 +348,7 @@ $ . venv/bin/activate
 (venv)$ ./tests.sh
 ```
 
+> [!IMPORTANT]
 > Make sure to always run the tests inside your venv!
 
 ### Configure Sentry SDK
@@ -389,8 +401,8 @@ Planned features (coming soon):
 - [x] **Sentry** integration for exception reporting
 - [x] **Ansible role** for easy production deployment
 - [x] **Github workflow** for CI/testing
-- [ ] **Publish package** to [PyPI](https://pypi.org/)
 - [ ] Implement a **configurable webhook API** call for notification to sender on reaching quota limit (on first block) to external service.
+- [ ] **Publish package** to [PyPI](https://pypi.org/) (Might need some restructuring. Any help greatly appreciated!)
 
 ## Credits 🙏
 

app/db.py

@@ -0,0 +1,41 @@
+from importlib import import_module
+from dbutils.pooled_db import PooledDB
+
+class DbConnectionPool:
+    def __init__(self, conf: object):
+        self.driver = conf.get('DB_DRIVER', 'pymysql').lower()
+        self.backend = import_module(self.driver)
+        db_config = getattr(self, 'get_dbconfig_' + self.driver)(conf)
+
+        self.pool = PooledDB(
+            creator=self.backend, # pymysql or sqlite3
+            mincached=int(conf.get('DB_POOL_MINCACHED', 0)),
+            maxcached=int(conf.get('DB_POOL_MAXCACHED', 10)),
+            maxshared=int(conf.get('DB_POOL_MAXSHARED', 10)),
+            maxusage=int(conf.get('DB_POOL_MAXUSAGE', 10000)),
+            # maxconnections=int(conf.get('DB_POOL_MAXCONNECTIONS', 20)),
+            **db_config
+        )
+    
+    def connection(self):
+        connection = self.pool.connection()
+        if self.driver == 'sqlite3':
+            # https://docs.python.org/3/library/sqlite3.html#sqlite3.Row
+            connection.row_factory = self.backend.Row
+        return connection
+
+
+    def get_dbconfig_pymysql(self, conf: object) -> dict:
+        return {
+            'host': conf.get('DB_HOST', 'localhost'),
+            'user': conf.get('DB_USER', 'policyd-rate-guard'),
+            'password': conf.get('DB_PASSWORD', ''),
+            'database': conf.get('DB_DATABASE', 'policyd-rate-guard'),
+            'port': int(conf.get('DB_PORT', 3306)),
+            'cursorclass': self.backend.cursors.DictCursor
+        }
+
+    def get_dbconfig_sqlite3(self, conf: object) -> dict:
+        return {
+            'database': conf.get('DB_DATABASE', ':memory:'),
+        }

app/handler.py

@@ -3,21 +3,21 @@
 class Handler:
     """Handle request"""
 
+    db = None
     data = ''
 
-    def __init__(self, conn: object, addr: str, conf: object, logger: object, db: object):
+    def __init__(self, conn: object, addr: str, conf: object, logger: object, db_pool: object):
         self.conn = conn
         self.addr = addr
         self.conf = conf
         self.logger = logger
-        self.db = db
+        self.db_pool = db_pool
         try:
             self.handle()
         except Exception as e: # pragma: no cover
             self.logger.exception('Unhandled Exception: %s', e)
             self.logger.warning('Received DATA: %s', self.data)
             self.send_response('DUNNO') # use DUNNO as accept action, just to distinguish between OK and unhandled exception
-            self.conn.close()
 
     def handle(self):
         """Handle request"""
@@ -46,16 +46,14 @@ def handle(self):
         if not request.get('sasl_username'):
             self.logger.debug('sasl_username is empty, accepting message and reply with DUNNO')
             self.send_response('DUNNO')
-            self.conn.close()
             return
 
         # Temp debugging of message data without recipient (e.g. on multiple To: addresses)
         # if not request.get('recipient'):
         #     self.logger.warning('Received DATA with no recipient: %s', self.data)
 
-        # PyMySQL: Ensure the db connection is alive
-        if self.conf.get('DB_DRIVER', 'pymysql').lower() == 'pymysql':
-            self.db.ping(reconnect=True)
+        # Get database connection from DB pool
+        self.db = self.db_pool.connection()
 
         # Handle message
         message = Message(
@@ -98,21 +96,28 @@ def handle(self):
         # Create response
         if blocked:
             self.logger.warning('Message BLOCKED: %s', message.get_props_description())
-            self.send_response('DEFER_IF_PERMIT ' + self.conf.get('ACTION_TEXT_BLOCKED', 'Rate limit reached, retry later'))
             if not was_blocked: # TODO: Implement webhook API call for notification to sender on quota limit reached (only on first block)
                 self.logger.debug('Quota limit reached for %s, notifying sender via webhook!', message.sender)
             self.logger.warning(log_message)
+            self.send_response('DEFER_IF_PERMIT ' + self.conf.get('ACTION_TEXT_BLOCKED', 'Rate limit reached, retry later'))
         else:
             self.logger.debug('Message ACCEPTED: %s', message.get_props_description())
-            self.send_response('OK')
             self.logger.info(log_message)
-
-        self.logger.msgid = None # Reset msgid in logger
-        self.conn.close()
+            self.send_response('OK')
 
     def send_response(self, message: str = 'OK'):
         """Send response"""
-        # actions return to postfix, see http://www.postfix.org/access.5.html for a list of actions.
+        # actions return to Postfix, see http://www.postfix.org/access.5.html for a list of actions.
         data = 'action={}\n\n'.format(message)
         self.logger.debug('Sending data: %s', data)
         self.conn.send(data.encode('utf-8'))
+        self.conn.close()
+
+    def __del__(self):
+        """Destructor"""
+        self.logger.msgid = None # Reset msgid in logger
+        # TODO: Do we need to close the cursor as well? (prior to closing the connection)
+        if self.db is not None:
+            self.db.close() # Close database connection
+            self.db = None
+        self.logger.debug('Handler destroyed')

app/ratelimit.py

@@ -138,12 +138,18 @@ def find(sender: str, db: object, logger: object, conf: object) -> object:
         )
 
     @staticmethod
-    def reset_all_counters(db: object, logger: object):
+    def reset_all_counters(db_pool: object, logger: object):
         """Reset all ratelimit counters"""
         logger.debug('Reset all counters')
-        cursor = db.cursor()
-        # reset all counters, but don't change updated_at timestamp
-        cursor.execute('UPDATE `ratelimits` SET `msg_counter` = 0, `rcpt_counter` = 0, `updated_at` = `updated_at`')
-        # only reset quota if it is not locked
-        cursor.execute('UPDATE `ratelimits` SET `quota` = `quota_reset`, `updated_at` = `updated_at` WHERE `quota_locked` = 0')
-        db.commit()
+        db = db_pool.connection()
+        try:
+            # With DBUtils PooledDB, we need to explicitly start a transaction
+            db.begin()
+            cursor = db.cursor()
+            # reset all counters, but don't change updated_at timestamp
+            cursor.execute('UPDATE `ratelimits` SET `msg_counter` = 0, `rcpt_counter` = 0, `updated_at` = `updated_at`')
+            # only reset quota if it is not locked
+            cursor.execute('UPDATE `ratelimits` SET `quota` = `quota_reset`, `updated_at` = `updated_at` WHERE `quota_locked` = 0')
+            db.commit()
+        finally:
+            db.close()

cleanup.py

@@ -1,20 +1,20 @@
 from app.conf import Config
 from app.logging import get_logger
-from database.db import connect_database
+from app.db import DbConnectionPool
 from app.ratelimit import Ratelimit
 
 class Cleaner:
 
     def __init__(self, conf: object = None) -> None:
         self.conf = conf or Config()
         self.logger = get_logger(self.conf)
-        self.db = connect_database(self.conf)
+        self.db_pool = DbConnectionPool(self.conf)
         self.cleanup()
 
     def cleanup(self) -> None:
         """Cleanup database"""
         self.logger.debug('Cleaning up database')
-        Ratelimit.reset_all_counters(self.db, self.logger)
+        Ratelimit.reset_all_counters(self.db_pool, self.logger)
 
 
 if __name__ == '__main__': # pragma: no cover

database/db.py

@@ -1,4 +1,5 @@
 cryptography==41.0.*
+DBUtils==3.0.*
 PyMySQL==1.1.*
 python-dotenv==1.0.*
 sentry-sdk==1.29.*

requirements.txt

@@ -18,14 +18,14 @@
 from app.conf import Config
 from app.handler import Handler
 from app.logging import get_logger
-from database.db import connect_database
+from app.db import DbConnectionPool
 
 class Daemon:
 
     def __init__(self) -> None:
         self.conf = Config()
         self.logger = get_logger(self.conf)
-        self.db = connect_database(self.conf)
+        self.db_pool = DbConnectionPool(self.conf)
         # TODO: Improve socket configuration parsing
         self.socket_conf = self.conf.get_array('SOCKET', '127.0.0.1,10033')
         self.init_sentry()
@@ -49,7 +49,7 @@ def run(self) -> None:
                 conn, addr = self.socket.accept()
                 threading.Thread(
                     target=Handler,
-                    args=(conn, addr, self.conf, self.logger, self.db)
+                    args=(conn, addr, self.conf, self.logger, self.db_pool)
                 ).start()
             except KeyboardInterrupt:
                 break