Python: model aiopg.sa

Author: yoff

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

@@ -41,6 +41,7 @@ private module Aiopg {
     result = connection().getMember("cursor").getReturn().getAwaited()
   }
 
+  /** Calling `execute` on a `Cursor` constructs a query. */
   class CursorExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
     CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
 
@@ -71,4 +72,49 @@ private module Aiopg {
 
     override DataFlow::Node getSql() { result = sql }
   }
+
+  /** An `Engine` is created when the result of calling `aiopg.sa.create_engine` is awaited. */
+  API::Node engine() {
+    result =
+      API::moduleImport("aiopg").getMember("sa").getMember("create_engine").getReturn().getAwaited()
+  }
+
+  /**
+   * A `SAConnection` is created when the result of calling `aquire` on an `Engine` is awaited.
+   */
+  API::Node saConnection() { result = engine().getMember("acquire").getReturn().getAwaited() }
+
+  /** Calling `execute` on a `SAConnection` constructs a query. */
+  class SAConnectionExecuteCall extends SqlConstruction::Range, DataFlow::CallCfgNode {
+    SAConnectionExecuteCall() { this = saConnection().getMember("execute").getACall() }
+
+    override DataFlow::Node getSql() { result in [this.getArg(0), this.getArgByName("query")] }
+  }
+
+  /**
+   * 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 saConnectionExecuteCall(
+    DataFlow::TypeTracker t, DataFlow::Node sql
+  ) {
+    // saConnection created from engine
+    t.start() and
+    sql = result.(SAConnectionExecuteCall).getSql()
+    or
+    exists(DataFlow::TypeTracker t2 | result = saConnectionExecuteCall(t2, sql).track(t2, t))
+  }
+
+  DataFlow::Node saConnectionExecuteCall(DataFlow::Node sql) {
+    saConnectionExecuteCall(DataFlow::TypeTracker::end(), sql).flowsTo(result)
+  }
+
+  /** Awaiting the result of calling `execute` executes the query. */
+  class AwaitedSAConnectionExecuteCall extends SqlExecution::Range {
+    DataFlow::Node sql;
+
+    AwaitedSAConnectionExecuteCall() { this = awaited(saConnectionExecuteCall(sql)) }
+
+    override DataFlow::Node getSql() { result = sql }
+  }
 }

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

@@ -23,4 +23,4 @@ async def test_cursor():
 async def test_engine():
     engine = await create_engine()
     conn = await engine.acquire()
-    await conn.execute("sql")  # $ MISSING: getSql="sql" constructedSql="sql"
+    await conn.execute("sql")  # $ getSql="sql" constructedSql="sql"