Python: Share implementation of awaited

python/ql/lib/semmle/python/ApiGraphs.qll

@@ -302,6 +302,8 @@ module API {
      * API graph node for the prefix `foo`), in accordance with the usual semantics of Python.
      */
 
+    private import semmle.python.internal.Awaited
+
     cached
     newtype TApiNode =
       /** The root of the API graph. */
@@ -485,43 +487,6 @@ module API {
       )
     }
 
-    /**
-     * Holds if `result` is the result of awaiting `awaitedValue`.
-     */
-    cached
-    DataFlow::Node awaited(DataFlow::Node awaitedValue) {
-      // `await` x
-      // - `awaitedValue` is `x`
-      // - `result` is `await x`
-      exists(Await await |
-        await.getValue() = awaitedValue.asExpr() and
-        result.asExpr() = await
-      )
-      or
-      // `async for x in l`
-      // - `awaitedValue` is `l`
-      // - `result` is `l` (`x` is behind a read step)
-      exists(AsyncFor asyncFor |
-        // To consider `x` the result of awaiting, we would use asyncFor.getTarget() = awaitedValue.asExpr(),
-        // but that is behind a read step rather than a flow step.
-        asyncFor.getIter() = awaitedValue.asExpr() and
-        result.asExpr() = asyncFor.getIter()
-      )
-      or
-      // `async with x as y`
-      // - `awaitedValue` is `x`
-      // - `result` is `x` and `y` if it exists
-      exists(AsyncWith asyncWith |
-        awaitedValue.asExpr() = asyncWith.getContextExpr() and
-        result.asExpr() in [
-            // `x`
-            asyncWith.getContextExpr(),
-            // `y`, if it exists
-            asyncWith.getOptionalVars()
-          ]
-      )
-    }
-
     /**
      * Holds if `ref` is a use of a node that should have an incoming edge from `base` labeled
      * `lbl` in the API graph.

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

@@ -10,6 +10,8 @@ private import semmle.python.ApiGraphs
 
 /** Provides models for the `asyncpg` PyPI package. */
 private module Asyncpg {
+  private import semmle.python.internal.Awaited
+
   /** A `ConectionPool` is created when the result of `asyncpg.create_pool()` is awaited. */
   API::Node connectionPool() {
     result = API::moduleImport("asyncpg").getMember("create_pool").getReturn().getAwaited()
@@ -63,41 +65,6 @@ private module Asyncpg {
     }
   }
 
-  /**
-   * Holds if `result` is the result of awaiting `awaitedValue`.
-   *
-   * Internal helper predicate to achieve the same as `.awaited()` does for API graphs,
-   * but sutiable for use with type-tracking.
-   */
-  pragma[inline]
-  DataFlow::Node awaited(DataFlow::Node awaitedValue) {
-    // `await` x
-    // - `awaitedValue` is `x`
-    // - `result` is `await x`
-    exists(Await await |
-      await.getValue() = awaitedValue.asExpr() and
-      result.asExpr() = await
-    )
-    or
-    // `async for x in l`
-    // - `awaitedValue` is local source of `l`
-    // - `result` is `l`
-    exists(AsyncFor asyncFor, DataFlow::Node awaited |
-      asyncFor.getIter() = awaited.asExpr() and
-      awaited.getALocalSource() = awaitedValue and
-      result.asExpr() = asyncFor.getIter()
-    )
-    or
-    // `async with x as y`
-    // - `awaitedValue` is local source of `x`
-    // - `result` is `x` and `y`
-    exists(AsyncWith asyncWith, DataFlow::Node awaited |
-      awaited.asExpr() = asyncWith.getContextExpr() and
-      awaited.getALocalSource() = awaitedValue and
-      result.asExpr() in [asyncWith.getContextExpr(), asyncWith.getOptionalVars()]
-    )
-  }
-
   /**
    * Provides models of the `PreparedStatement` class in `asyncpg`.
    * `PreparedStatement`s are created when the result of calling `prepare(query)` on a connection is awaited.

python/ql/lib/semmle/python/internal/Awaited.qll

@@ -0,0 +1,45 @@
+/**
+ * INTERNAL: Do not use.
+ *
+ * Provides helper class for defining additional API graph edges.
+ */
+
+private import python
+private import semmle.python.dataflow.new.DataFlow
+
+/**
+ * Holds if `result` is the result of awaiting `awaitedValue`.
+ */
+cached
+DataFlow::Node awaited(DataFlow::Node awaitedValue) {
+  // `await` x
+  // - `awaitedValue` is `x`
+  // - `result` is `await x`
+  exists(Await await |
+    await.getValue() = awaitedValue.asExpr() and
+    result.asExpr() = await
+  )
+  or
+  // `async for x in l`
+  // - `awaitedValue` is `l`
+  // - `result` is `l` (`x` is behind a read step)
+  exists(AsyncFor asyncFor |
+    // To consider `x` the result of awaiting, we would use asyncFor.getTarget() = awaitedValue.asExpr(),
+    // but that is behind a read step rather than a flow step.
+    asyncFor.getIter() = awaitedValue.asExpr() and
+    result.asExpr() = asyncFor.getIter()
+  )
+  or
+  // `async with x as y`
+  // - `awaitedValue` is `x`
+  // - `result` is `x` and `y` if it exists
+  exists(AsyncWith asyncWith |
+    awaitedValue.asExpr() = asyncWith.getContextExpr() and
+    result.asExpr() in [
+        // `x`
+        asyncWith.getContextExpr(),
+        // `y`, if it exists
+        asyncWith.getOptionalVars()
+      ]
+  )
+}

python/ql/test/experimental/dataflow/ApiGraphs/awaited.ql

@@ -9,18 +9,18 @@ class AwaitedTest extends InlineExpectationsTest {
   override string getARelevantTag() { result = "awaited" }
 
   override predicate hasActualResult(Location location, string element, string tag, string value) {
-    exists(API::Node a, DataFlow::Node n, API::Node pred |
-      a = pred.getAwaited() and
-      n = a.getAUse() and
-      location = n.getLocation() and
+    exists(API::Node awaited, DataFlow::Node use, API::Node pred |
+      awaited = pred.getAwaited() and
+      use = awaited.getAUse() and
+      location = use.getLocation() and
       // Module variable nodes have no suitable location, so it's best to simply exclude them entirely
       // from the inline tests.
-      not n instanceof DataFlow::ModuleVariableNode and
+      not use instanceof DataFlow::ModuleVariableNode and
       exists(location.getFile().getRelativePath())
     |
       tag = "awaited" and
       value = pred.getPath() and
-      element = n.toString()
+      element = use.toString()
     )
   }
 }