Python: model aiopg

python/ql/lib/semmle/python/Frameworks.qll

@@ -6,6 +6,7 @@
 // `docs/codeql/support/reusables/frameworks.rst`
 private import semmle.python.frameworks.Aioch
 private import semmle.python.frameworks.Aiohttp
+private import semmle.python.frameworks.Aiopg
 private import semmle.python.frameworks.Asyncpg
 private import semmle.python.frameworks.ClickhouseDriver
 private import semmle.python.frameworks.Cryptodome

python/ql/lib/semmle/python/frameworks/Aiopg.qll

@@ -0,0 +1,74 @@
+/**
+ * Provides classes modeling security-relevant aspects of the `aiopg` PyPI package.
+ * See
+ * - https://aiopg.readthedocs.io/en/stable/index.html
+ * - https://pypi.org/project/aiopg/
+ */
+
+private import python
+private import semmle.python.dataflow.new.DataFlow
+private import semmle.python.Concepts
+private import semmle.python.ApiGraphs
+
+/** Provides models for the `aiopg` PyPI package. */
+private module Aiopg {
+  private import semmle.python.internal.Awaited
+
+  /** A `ConectionPool` is created when the result of `aiopg.create_pool()` is awaited. */
+  API::Node connectionPool() {
+    result = API::moduleImport("aiopg").getMember("create_pool").getReturn().getAwaited()
+  }
+
+  /**
+   * A `Connection` is created when
+   * - the result of `aiopg.connect()` is awaited.
+   * - the result of calling `aquire` on a `ConnectionPool` is awaited.
+   */
+  API::Node connection() {
+    result = API::moduleImport("aiopg").getMember("connect").getReturn().getAwaited()
+    or
+    result = connectionPool().getMember("acquire").getReturn().getAwaited()
+  }
+
+  /**
+   * A `Cursor` is created when
+   * - the result of calling `cursor` on a `ConnectionPool` is awaited.
+   * - the result of calling `cursor` on a `Connection` is awaited.
+   */
+  API::Node cursor() {
+    result = connectionPool().getMember("cursor").getReturn().getAwaited()
+    or
+    result = connection().getMember("cursor").getReturn().getAwaited()
+  }
+
+  class CursorExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
+    CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
+
+    override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("operation")] }
+  }
+
+  /**
+   * This is only needed to connect the argument to the execute call with the subsequnt awaiting.
+   * It should be obsolete once we have `API::CallNode` available.
+   */
+  private DataFlow::TypeTrackingNode cursorExecuteCall(DataFlow::TypeTracker t, DataFlow::Node sql) {
+    // cursor created from connection
+    t.start() and
+    sql = result.(CursorExecuteCall).getSql()
+    or
+    exists(DataFlow::TypeTracker t2 | result = cursorExecuteCall(t2, sql).track(t2, t))
+  }
+
+  DataFlow::Node cursorExecuteCall(DataFlow::Node sql) {
+    cursorExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
+  }
+
+  /** Awaiting the result of calling `execute` executes the query. */
+  class AwaitedCursorExecuteCall extends SqlExecution::Range {
+    DataFlow::Node sql;
+
+    AwaitedCursorExecuteCall() { this = awaited(cursorExecuteCall(sql)) }
+
+    override DataFlow::Node getSql() { result = sql }
+  }
+}

python/ql/test/library-tests/frameworks/aiopg/test.py

@@ -5,15 +5,15 @@ async def test_cursor():
     # Create connection directly
     conn = await aiopg.connect()
     cur = await conn.cursor()
-    await cur.execute("sql")  # $ MISSING: getSql="sql"
+    await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
     # Create connection via pool
     async with aiopg.create_pool() as pool:
         # Create Cursor via Connection
         async with pool.acquire() as conn:
             cur = await conn.cursor()
-            await cur.execute("sql")  # $ MISSING: getSql="sql"
+            await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"
 
         # Create Cursor directly
         async with pool.cursor() as cur:
-            await cur.execute("sql")  # $ MISSING: getSql="sql"
+            await cur.execute("sql")  # $ getSql="sql" constructedSql="sql"