Merge pull request #20276 from github/tausbn/python-model-psycopg2-connection-pools

Python: Add support for Psycopg2 database connection pools

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.

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()
+    }
+  }
 }

python/ql/test/library-tests/frameworks/psycopg2/ConceptsTest.ql

@@ -0,0 +1,2 @@
+import python
+import experimental.meta.ConceptsTest

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()