Allow displaying multiple results

src/driver.ts

@@ -11,7 +11,7 @@ export interface Pool {
   end: () => void;
 }
 
-export type QueryResult = ResultTable | string;
+export type QueryResult = ResultTable[] | string;
 
 export type Row = { [key: string]: string | number | null };
 
@@ -36,6 +36,7 @@ interface BaseConfig {
 
 interface MySQLConfig extends BaseConfig {
   driver: 'mysql';
+  multipleStatements: boolean;
 }
 
 export async function getPool(c: PoolConfig): Promise<Pool> {
@@ -57,8 +58,16 @@ async function createMySQLPool({
   user,
   password,
   database,
+  multipleStatements,
 }: MySQLConfig): Promise<Pool> {
-  return mysqlPool(mysql.createPool({ host, port, user, password, database }));
+  return mysqlPool(mysql.createPool({
+    host,
+    port,
+    user,
+    password,
+    database,
+    multipleStatements,
+  }));
 }
 
 function mysqlPool(pool: mysql.Pool): Pool {
@@ -77,12 +86,16 @@ function mysqlConn(conn: mysql.PoolConnection): Conn {
     destroy() {
       conn.destroy();
     },
-    async query(q: string): Promise<ResultTable> {
+    async query(q: string): Promise<ResultTable[]> {
       const [result] = (await conn.query(q)) as any;
       if (!result.length) {
-        return [result] as ResultTable;
+        return [[result] as ResultTable];
       }
-      return result as ResultTable;
+      if (result.length > 0 && !!result[0].length) {
+        // Nested row data for multiple statements.
+        return result as ResultTable[];
+      }
+      return [result as ResultTable];
     },
     release() {
       conn.release();
@@ -125,9 +138,17 @@ function postgresPool(pool: pg.Pool): Pool {
 
 function postgresConn(conn: pg.PoolClient): Conn {
   return {
-    async query(q: string): Promise<ResultTable> {
+    async query(q: string): Promise<ResultTable[]> {
       const response = await conn.query(q);
-      return response.rows;
+
+      // Typings for pg unfortunately miss that `query` may return an array of
+      // results when the query strings contains multiple sql statements.
+      const maybeResponses = response as any as pg.QueryResult<any>[];
+      if (!!maybeResponses.length) {
+        return maybeResponses.map(r => r.rows);
+      }
+
+      return [response.rows];
     },
     destroy() {
       // TODO: verify
@@ -180,7 +201,7 @@ function mssqlConn(req: mssql.Request): Conn {
       if (res.recordsets.length < 1) {
         return `Rows affected: ${res.rowsAffected}`;
       }
-      return res.recordsets[0];
+      return [res.recordsets[0]];
     },
     release() {
       // TODO: verify correctness

src/extension.ts

@@ -172,7 +172,8 @@ class SQLNotebookController {
       writeSuccess(execution, result);
       return;
     }
-    writeSuccess(execution, resultToMarkdownTable(result), 'text/markdown');
+    const tables = result.map(r => resultToMarkdownTable(r));
+    writeSuccess(execution, tables, 'text/markdown');
   }
 }
 
@@ -220,14 +221,15 @@ const writeErr = (execution: vscode.NotebookCellExecution, err: string) => {
 
 const writeSuccess = (
   execution: vscode.NotebookCellExecution,
-  text: string,
+  text: string | string[],
   mimeType?: string
 ) => {
-  execution.replaceOutput([
-    new vscode.NotebookCellOutput([
-      vscode.NotebookCellOutputItem.text(text, mimeType),
-    ]),
-  ]);
+  const items = typeof text === 'string' ? [text] : text;
+  execution.replaceOutput(
+    items.map(item => new vscode.NotebookCellOutput([
+      vscode.NotebookCellOutputItem.text(item, mimeType),
+    ])),
+  );
   execution.end(true, Date.now());
 };
 

webview/Form.tsx

@@ -91,7 +91,13 @@ function useDropdownValue() {
 function showDriverConfig(driver: string) {
   switch (driver) {
     case 'mysql':
-      return <></>;
+      return (
+        <>
+          <vscode-checkbox name="multipleStatements" checked>
+            Multiple statements
+          </vscode-checkbox>
+        </>
+      );
     case 'postgres':
       return <></>;
     case 'mssql':