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 000000000000..5a94d9829b48 --- /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 74016ebd639d..93011f222775 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() + } + } } 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 000000000000..e69de29bb2d1 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 000000000000..b557a0bccb69 --- /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 000000000000..507cdd59b822 --- /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()