Python: Make asyncio version of PEP249 modeling library

so it's also easy to modeling asyncio libraries

Also ports aiomysql/aiopg to use this new modeling

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

@@ -9,11 +9,10 @@ private import python
 private import semmle.python.dataflow.new.DataFlow
 private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
+private import semmle.python.frameworks.PEP249
 
 /** Provides models for the `aiomysql` PyPI package. */
 private module Aiomysql {
-  private import semmle.python.internal.Awaited
-
   /**
    * Gets a `ConnectionPool` that is created when the result of `aiomysql.create_pool()` is awaited.
    * See https://aiomysql.readthedocs.io/en/stable/pool.html
@@ -23,49 +22,29 @@ private module Aiomysql {
   }
 
   /**
-   * Gets a `Connection` that is created when
+   * A Connection that is created when
    * - the result of `aiomysql.connect()` is awaited.
    * - the result of calling `acquire` on a `ConnectionPool` is awaited.
-   * See https://aiomysql.readthedocs.io/en/stable/connection.html#connection
+   * See
+   * - https://aiomysql.readthedocs.io/en/stable/connection.html#connection
+   * - https://aiomysql.readthedocs.io/en/stable/pool.html#Pool.acquire
    */
-  API::Node connection() {
-    result = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
-    or
-    result = connectionPool().getMember("acquire").getReturn().getAwaited()
+  class AiomysqlConnection extends PEP249::AsyncDatabaseConnection {
+    AiomysqlConnection() {
+      this = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
+      or
+      this = connectionPool().getMember("acquire").getReturn().getAwaited()
+    }
   }
 
   /**
-   * Gets a `Cursor` that is created when
+   * An additional cursor, that is created when
    * - the result of calling `cursor` on a `ConnectionPool` is awaited.
-   * - the result of calling `cursor` on a `Connection` is awaited.
-   * See https://aiomysql.readthedocs.io/en/stable/cursors.html
+   * See
+   * - https://aiomysql.readthedocs.io/en/stable/pool.html##Pool.cursor
    */
-  API::Node cursor() {
-    result = connectionPool().getMember("cursor").getReturn().getAwaited()
-    or
-    result = connection().getMember("cursor").getReturn().getAwaited()
-  }
-
-  /**
-   * A query. Calling `execute` on a `Cursor` constructs a query.
-   * See https://aiomysql.readthedocs.io/en/stable/cursors.html#Cursor.execute
-   */
-  class CursorExecuteCall extends SqlConstruction::Range, API::CallNode {
-    CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
-
-    override DataFlow::Node getSql() { result = this.getParameter(0, "operation").asSink() }
-  }
-
-  /**
-   * An awaited query. Awaiting the result of calling `execute` executes the query.
-   * See https://aiomysql.readthedocs.io/en/stable/cursors.html#Cursor.execute
-   */
-  class AwaitedCursorExecuteCall extends SqlExecution::Range {
-    CursorExecuteCall executeCall;
-
-    AwaitedCursorExecuteCall() { this = executeCall.getReturn().getAwaited().asSource() }
-
-    override DataFlow::Node getSql() { result = executeCall.getSql() }
+  class AiomysqlCursor extends PEP249::AsyncDatabaseCursor {
+    AiomysqlCursor() { this = connectionPool().getMember("cursor").getReturn().getAwaited() }
   }
 
   /**

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

@@ -9,11 +9,10 @@ private import python
 private import semmle.python.dataflow.new.DataFlow
 private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
+private import semmle.python.frameworks.PEP249
 
 /** Provides models for the `aiopg` PyPI package. */
 private module Aiopg {
-  private import semmle.python.internal.Awaited
-
   /**
    * Gets a `ConnectionPool` that is created when the result of `aiopg.create_pool()` is awaited.
    * See https://aiopg.readthedocs.io/en/stable/core.html#pool
@@ -23,49 +22,29 @@ private module Aiopg {
   }
 
   /**
-   * Gets a `Connection` that is created when
+   * A Connection that is created when
    * - the result of `aiopg.connect()` is awaited.
    * - the result of calling `acquire` on a `ConnectionPool` is awaited.
-   * See https://aiopg.readthedocs.io/en/stable/core.html#connection
+   * See
+   * - https://aiopg.readthedocs.io/en/stable/core.html#connection
+   * - https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Pool.acquire
    */
-  API::Node connection() {
-    result = API::moduleImport("aiopg").getMember("connect").getReturn().getAwaited()
-    or
-    result = connectionPool().getMember("acquire").getReturn().getAwaited()
+  class AiopgConnection extends PEP249::AsyncDatabaseConnection {
+    AiopgConnection() {
+      this = API::moduleImport("aiopg").getMember("connect").getReturn().getAwaited()
+      or
+      this = connectionPool().getMember("acquire").getReturn().getAwaited()
+    }
   }
 
   /**
-   * Gets a `Cursor` that is created when
+   * An additional cursor, that is created when
    * - the result of calling `cursor` on a `ConnectionPool` is awaited.
-   * - the result of calling `cursor` on a `Connection` is awaited.
-   * See https://aiopg.readthedocs.io/en/stable/core.html#cursor
+   * See
+   * - https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Pool.cursor
    */
-  API::Node cursor() {
-    result = connectionPool().getMember("cursor").getReturn().getAwaited()
-    or
-    result = connection().getMember("cursor").getReturn().getAwaited()
-  }
-
-  /**
-   * A query. Calling `execute` on a `Cursor` constructs a query.
-   * See https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Cursor.execute
-   */
-  class CursorExecuteCall extends SqlConstruction::Range, API::CallNode {
-    CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
-
-    override DataFlow::Node getSql() { result = this.getParameter(0, "operation").asSink() }
-  }
-
-  /**
-   * An awaited query. Awaiting the result of calling `execute` executes the query.
-   * See https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Cursor.execute
-   */
-  class AwaitedCursorExecuteCall extends SqlExecution::Range {
-    CursorExecuteCall execute;
-
-    AwaitedCursorExecuteCall() { this = execute.getReturn().getAwaited().asSource() }
-
-    override DataFlow::Node getSql() { result = execute.getSql() }
+  class AiopgCursor extends PEP249::AsyncDatabaseCursor {
+    AiopgCursor() { this = connectionPool().getMember("cursor").getReturn().getAwaited() }
   }
 
   /**

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

@@ -52,6 +52,8 @@ module PEP249 {
     result in ["sql", "statement", "operation", "query", "query_string", "sql_script"]
   }
 
+  private string getExecuteMethodName() { result in ["execute", "executemany", "executescript"] }
+
   /**
    * A call to an execute method on a database cursor or a connection, such as `execute`
    * or `executemany`.
@@ -68,7 +70,7 @@ module PEP249 {
       exists(API::Node start |
         start instanceof DatabaseCursor or start instanceof DatabaseConnection
       |
-        this = start.getMember(["execute", "executemany", "executescript"]).getACall()
+        this = start.getMember(getExecuteMethodName()).getACall()
       )
     }
 
@@ -77,6 +79,82 @@ module PEP249 {
     }
   }
 
+  // ---------------------------------------------------------------------------
+  // asyncio implementations
+  // ---------------------------------------------------------------------------
+  //
+  // we differentiate between normal and asyncio implementations, since we model the
+  // `execute` call differently -- as a SqlExecution vs SqlConstruction, since the SQL
+  // is only executed in asyncio after being awaited (which might happen in something
+  // like `asyncio.gather`)
+  /**
+   * An API graph node representing a module that implements PEP 249 using asyncio.
+   */
+  abstract class AsyncPEP249ModuleApiNode extends API::Node {
+    /** Gets a string representation of this element. */
+    override string toString() { result = this.(API::Node).toString() }
+  }
+
+  /**
+   * An API graph node representing a asyncio database connection (after being awaited).
+   */
+  abstract class AsyncDatabaseConnection extends API::Node {
+    /** Gets a string representation of this element. */
+    override string toString() { result = this.(API::Node).toString() }
+  }
+
+  private class DefaultAsyncDatabaseConnection extends AsyncDatabaseConnection {
+    DefaultAsyncDatabaseConnection() {
+      this = any(AsyncPEP249ModuleApiNode mod).getMember("connect").getReturn().getAwaited()
+    }
+  }
+
+  /**
+   * An API graph node representing a asyncio database cursor (after being awaited).
+   */
+  abstract class AsyncDatabaseCursor extends API::Node {
+    /** Gets a string representation of this element. */
+    override string toString() { result = this.(API::Node).toString() }
+  }
+
+  private class DefaultAsyncDatabaseCursor extends AsyncDatabaseCursor {
+    DefaultAsyncDatabaseCursor() {
+      this = any(AsyncDatabaseConnection conn).getMember("cursor").getReturn().getAwaited()
+    }
+  }
+
+  /**
+   * A call to an execute method on an asyncio database cursor or an asyncio connection,
+   * such as `execute` or `executemany`.
+   *
+   * (This is not an SqlExecution, since that only happens when the coroutine is
+   * awaited)
+   *
+   * See ExecuteMethodCall for more details.
+   */
+  private class AsyncExecuteMethodCall extends SqlConstruction::Range, API::CallNode {
+    AsyncExecuteMethodCall() {
+      exists(API::Node start |
+        start instanceof AsyncDatabaseCursor or start instanceof AsyncDatabaseConnection
+      |
+        this = start.getMember(getExecuteMethodName()).getACall()
+      )
+    }
+
+    override DataFlow::Node getSql() {
+      result in [this.getArg(0), this.getArgByName(getSqlKwargName()),]
+    }
+  }
+
+  /** Actual execution of the AsyncExecuteMethodCall coroutine. */
+  private class AwaitedAsyncExecuteMethodCall extends SqlExecution::Range {
+    AsyncExecuteMethodCall execute;
+
+    AwaitedAsyncExecuteMethodCall() { this = execute.getReturn().getAwaited().asSource() }
+
+    override DataFlow::Node getSql() { result = execute.getSql() }
+  }
+
   // ---------------------------------------------------------------------------
   // old impl
   // ---------------------------------------------------------------------------