From d5e029899919db069537f29c81d8fedbfa99dee9 Mon Sep 17 00:00:00 2001 From: Taus Date: Mon, 25 Aug 2025 12:35:57 +0000 Subject: [PATCH 1/2] Python: Add support for Psycopg2 database connection pools Our current modelling only treated `psycopg2` insofar as it implemented PEP 249 (which does not define any notion of connection pool), which meant we were missing database connections that arose from such pools. With these changes, we add support for the three classes relating to database pools that are defined in `psycopg2`. (Note that `getAnInstance` automatically looks at subclasses, which means this should also handle cases where the user has defined a new subclass that inherits from one of these three classes.) --- ...2025-08-25-psycopg2-connection-pool-modelling.md | 5 +++++ python/ql/lib/semmle/python/frameworks/Psycopg2.qll | 13 +++++++++++++ 2 files changed, 18 insertions(+) create mode 100644 python/ql/lib/change-notes/2025-08-25-psycopg2-connection-pool-modelling.md diff --git a/python/ql/lib/change-notes/2025-08-25-psycopg2-connection-pool-modelling.md b/python/ql/lib/change-notes/2025-08-25-psycopg2-connection-pool-modelling.md new file mode 100644 index 00000000000..5a94d9829b4 --- /dev/null +++ b/python/ql/lib/change-notes/2025-08-25-psycopg2-connection-pool-modelling.md @@ -0,0 +1,5 @@ +--- +category: minorAnalysis +--- + +- The modelling of Psycopg2 now supports the use of `psycopg2.pool` connection pools for handling database connections. diff --git a/python/ql/lib/semmle/python/frameworks/Psycopg2.qll b/python/ql/lib/semmle/python/frameworks/Psycopg2.qll index 74016ebd639..93011f22277 100644 --- a/python/ql/lib/semmle/python/frameworks/Psycopg2.qll +++ b/python/ql/lib/semmle/python/frameworks/Psycopg2.qll @@ -29,4 +29,17 @@ private module Psycopg2 { class Psycopg2 extends PEP249::PEP249ModuleApiNode { Psycopg2() { this = API::moduleImport("psycopg2") } } + + /** A database connection obtained from a psycopg2 connection pool. */ + class Psycopg2ConnectionPoolMember extends PEP249::DatabaseConnection { + Psycopg2ConnectionPoolMember() { + this = + any(Psycopg2 p) + .getMember("pool") + .getMember(["SimpleConnectionPool", "ThreadedConnectionPool", "AbstractConnectionPool"]) + .getAnInstance() + .getMember("getconn") + .getReturn() + } + } } From 1008ca974430eb5ccb72a323c101377ce207e58a Mon Sep 17 00:00:00 2001 From: Taus Date: Mon, 25 Aug 2025 14:14:16 +0000 Subject: [PATCH 2/2] Python: Add `psycopg2.pool` tests --- .../frameworks/psycopg2/ConceptsTest.expected | 0 .../frameworks/psycopg2/ConceptsTest.ql | 2 + .../frameworks/psycopg2/connectionpool.py | 46 +++++++++++++++++++ 3 files changed, 48 insertions(+) create mode 100644 python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.expected create mode 100644 python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.ql create mode 100644 python/ql/test/library-tests/frameworks/psycopg2/connectionpool.py diff --git a/python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.expected b/python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.expected new file mode 100644 index 00000000000..e69de29bb2d diff --git a/python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.ql b/python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.ql new file mode 100644 index 00000000000..b557a0bccb6 --- /dev/null +++ b/python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.ql @@ -0,0 +1,2 @@ +import python +import experimental.meta.ConceptsTest diff --git a/python/ql/test/library-tests/frameworks/psycopg2/connectionpool.py b/python/ql/test/library-tests/frameworks/psycopg2/connectionpool.py new file mode 100644 index 00000000000..507cdd59b82 --- /dev/null +++ b/python/ql/test/library-tests/frameworks/psycopg2/connectionpool.py @@ -0,0 +1,46 @@ +# Examples using psycopg2 connection pools. + +import psycopg2 +from psycopg2.pool import SimpleConnectionPool, AbstractConnectionPool + + +DSN = "dbname=test user=test password=test host=localhost port=5432" + + +def run_simple_pool_query(): + pool = SimpleConnectionPool(1, 4, dsn=DSN) + try: + conn = pool.getconn() + try: + cur = conn.cursor() + try: + # Simple, parameterless query + cur.execute("SELECT 1") # $ getSql="SELECT 1" + _ = cur.fetchall() if hasattr(cur, "fetchall") else None # $ threatModelSource[database]=cur.fetchall() + finally: + cur.close() + finally: + pool.putconn(conn) + finally: + pool.closeall() + + +class LocalPool(AbstractConnectionPool): + pass + + +def run_custom_pool_query(): + pool = LocalPool(1, 3, dsn=DSN) + try: + conn = pool.getconn() + try: + cur = conn.cursor() + try: + cur.execute("SELECT 2") # $ getSql="SELECT 2" + _ = cur.fetchone() if hasattr(cur, "fetchone") else None # $ threatModelSource[database]=cur.fetchone() + finally: + cur.close() + finally: + pool.putconn(conn) + finally: + pool.closeall()