JS: Improve @google-cloud/spanner model

Author: asgerf

javascript/ql/src/semmle/javascript/frameworks/SQL.qll

@@ -543,68 +543,90 @@ private module Spanner {
   API::Node database() {
     result =
       spanner().getReturn().getMember("instance").getReturn().getMember("database").getReturn()
+    or
+    result = API::Node::ofType("@google-cloud/spanner", "Database")
   }
 
   /**
    * Gets a node that refers to an instance of the `v1.SpannerClient` class.
    */
   API::Node v1SpannerClient() {
     result = spanner().getMember("v1").getMember("SpannerClient").getInstance()
+    or
+    result = API::Node::ofType("@google-cloud/spanner", "v1.SpannerClient")
   }
 
   /**
    * Gets a node that refers to a transaction object.
    */
   API::Node transaction() {
-    result = database().getMember("runTransaction").getParameter(0).getParameter(1)
+    result =
+      database()
+          .getMember(["runTransaction", "runTransactionAsync"])
+          .getParameter([0, 1])
+          .getParameter(1)
+    or
+    result = API::Node::ofType("@google-cloud/spanner", "Transaction")
+  }
+
+  /** Gets an API node referring to a `BatchTransaction` object. */
+  API::Node batchTransaction() {
+    result = database().getMember("batchTransaction").getReturn()
+    or
+    result = database().getMember("createBatchTransaction").getReturn().getPromised()
+    or
+    result = API::Node::ofType("@google-cloud/spanner", "BatchTransaction")
   }
 
   /**
    * A call to a Spanner method that executes a SQL query.
    */
-  abstract class SqlExecution extends DatabaseAccess, DataFlow::InvokeNode {
-    /**
-     * Gets the position of the query argument; default is zero, which can be overridden
-     * by subclasses.
-     */
-    int getQueryArgumentPosition() { result = 0 }
-
-    override DataFlow::Node getAQueryArgument() {
-      result = getArgument(getQueryArgumentPosition()) or
-      result = getOptionArgument(getQueryArgumentPosition(), "sql")
-    }
-  }
+  abstract class SqlExecution extends DatabaseAccess, DataFlow::InvokeNode { }
 
   /**
-   * A call to `Database.run`, `Database.runPartitionedUpdate` or `Database.runStream`.
+   * A SQL execution that takes the input directly in the first argument or in the `sql` option.
    */
-  class DatabaseRunCall extends SqlExecution {
-    DatabaseRunCall() {
+  class SqlExecutionDirect extends SqlExecution {
+    SqlExecutionDirect() {
       this = database().getMember(["run", "runPartitionedUpdate", "runStream"]).getACall()
+      or
+      this = transaction().getMember(["run", "runStream", "runUpdate"]).getACall()
+      or
+      this = batchTransaction().getMember("createQueryPartitions").getACall()
+    }
+
+    override DataFlow::Node getAQueryArgument() {
+      result = getArgument(0)
+      or
+      result = getOptionArgument(0, "sql")
     }
   }
 
   /**
-   * A call to `Transaction.run`, `Transaction.runStream` or `Transaction.runUpdate`.
+   * A SQL execution that takes an array of SQL strings or { sql: string } objects.
    */
-  class TransactionRunCall extends SqlExecution {
-    TransactionRunCall() {
-      this = transaction().getMember(["run", "runStream", "runUpdate"]).getACall()
+  class SqlExecutionBatch extends SqlExecution, API::CallNode {
+    SqlExecutionBatch() { this = transaction().getMember("batchUpdate").getACall() }
+
+    override DataFlow::Node getAQueryArgument() {
+      // just use the whole array as the query argument, as arrays becomes tainted if one of the elements
+      // are tainted
+      result = getArgument(0)
+      or
+      result = getParameter(0).getUnknownMember().getMember("sql").getARhs()
     }
   }
 
   /**
-   * A call to `v1.SpannerClient.executeSql` or `v1.SpannerClient.executeStreamingSql`.
+   * A SQL execution that only takes the input in the `sql` option, and do not accept query strings
+   * directly.
    */
-  class ExecuteSqlCall extends SqlExecution {
-    ExecuteSqlCall() {
+  class SqlExecutionWithOption extends SqlExecution {
+    SqlExecutionWithOption() {
       this = v1SpannerClient().getMember(["executeSql", "executeStreamingSql"]).getACall()
     }
 
-    override DataFlow::Node getAQueryArgument() {
-      // `executeSql` and `executeStreamingSql` do not accept query strings directly
-      result = getOptionArgument(0, "sql")
-    }
+    override DataFlow::Node getAQueryArgument() { result = getOptionArgument(0, "sql") }
   }
 
   /**

javascript/ql/test/library-tests/frameworks/SQL/SqlString.expected

@@ -41,6 +41,7 @@
 | sequelizeImport.js:3:17:3:118 | 'SELECT ... Y name' |
 | spanner2.js:5:26:5:35 | "SQL code" |
 | spanner2.js:7:35:7:44 | "SQL code" |
+| spanner-types.ts:4:12:4:23 | 'SELECT 123' |
 | spanner.js:6:8:6:17 | "SQL code" |
 | spanner.js:7:8:7:26 | { sql: "SQL code" } |
 | spanner.js:7:15:7:24 | "SQL code" |
@@ -59,6 +60,8 @@
 | spanner.js:18:16:18:25 | "SQL code" |
 | spanner.js:19:16:19:34 | { sql: "SQL code" } |
 | spanner.js:19:23:19:32 | "SQL code" |
+| spanner.js:23:12:23:23 | 'SELECT 123' |
+| spanner.js:26:12:26:38 | 'UPDATE ... = @baz' |
 | spannerImport.js:4:8:4:17 | "SQL code" |
 | sqlite-types.ts:4:12:4:49 | "UPDATE ... id = ?" |
 | sqlite.js:7:8:7:45 | "UPDATE ... id = ?" |

javascript/ql/test/library-tests/frameworks/SQL/spanner-types.ts

@@ -0,0 +1,5 @@
+import { Database } from "@google-cloud/spanner";
+
+export function doSomething(db: Database) {
+    db.run('SELECT 123');
+}

javascript/ql/test/library-tests/frameworks/SQL/spanner.js

@@ -17,6 +17,18 @@ db.runTransaction((err, tx) => {
   tx.runStream({ sql: "SQL code" });
   tx.runUpdate("SQL code");
   tx.runUpdate({ sql: "SQL code" });
+
+  const queries = [
+    {
+      sql: 'SELECT 123',
+    },
+    {
+      sql: 'UPDATE foo SET bar = @baz',
+      params: {key: 'baz', value: '123'}
+    }
+  ];
+  
+  tx.batchUpdate(queries, () => {});
 });
 
 exports.instance = instance;