Improve multi-line and multi-statement evaluation

txase · txase · commit afb281087b90 · 2021-01-25T10:55:31.000-08:00
* Execute multiple statements, separated by ';', individually
* Parse comments out of statements
* Execute statements sequentially, waiting for each to complete before
  executing the next
diff --git a/README.md b/README.md
@@ -39,5 +39,25 @@ adasql will show you information up top to help you ensure you're connecting to
 
 Transactions are supported, though note the Data API doesn't support save points or nested transactions.
 
-### Canceling commands
-If you find yourself in the middle of a multi-line statement and wish to cancel it, enter `.clear`. This will reset the state of the REPL, though it will not affect a transaction if it is in progress.
+### Multiple statements and multi-line statements
+SQL statements are executed sequentially, and multiple statements can be executed when separated with the ';' delimiter. For example, the following query when executed will insert a new record and then return the same record as the first statement executes to completion before the second statement is executed:
+
+```SQL
+INSERT INTO people (id, name, age) VALUES (1, 'Alice', 39); SELECT * from mytable WHERE id = 1;
+```
+
+#### Canceling commands
+If you find yourself in the middle of a multi-line statement and wish to cancel it, enter `.clear`. This will reset the state of the REPL, though it will not affect a transaction if it is in progress.
+
+## FAQ
+### Can adasql be used for migrations, seeding data, or restoring backups?
+Maybe! But for most use cases adasql will not work. The Aurora Data API has three issues that make migrations and restoring backups difficult:
+
+* SQL connections to the DB used by the Data API are multiplexed, and there is no way to ensure affinity for statements that set per-connection variables. This is important when restoring backups where the backups need to execute commands like `SET FOREIGN_KEY_CHECKS=0` to disable foreign key checks when inserting records using the same connection. Multiple statements that need to be executed on the same connection with these connection-specific variables will likely fail part-way through execution.
+* SQL statements larger than 64 KB are not supported. If you want to insert a batch of records, make sure they are separated into chunks of statements smaller than 64 KB.
+* SQL statements that take longer than 45 seconds to complete will timeout. The adasql client does not tell the Data API to continue these statements if they do timeout for data integrity purposes. But, DDL statements that timeout may cause non-reversible changes when canceled.
+
+Your use case may not hit these limitations, in which case have at it! But many use cases will, especially the first and second limitations when attempting to restore a mysqldump backup. You may find you can work around the limitations by:
+
+1. Running mysqldump with the `--extended-insert=FALSE` argument to insert every record using a separate statement to keep statements under 64 KB in size
+1. If your tables have foreign key constraints, you may be able to re-order them with children tables and records first, then parent tables and records second, allowing you to create tables and insert records without violating the constraints
diff --git a/query.ts b/query.ts
@@ -30,16 +30,36 @@ import {
   Keywords
 } from './types';
 
+type REPLCallback = (err: Error | null, result?: string | Record<string, unknown>) => void
+
+interface Statement {
+  command: string,
+  cb: REPLCallback
+}
+
 import hydrateRecords from './hydrateRecords';
 import fetchKeywords from './fetchKeywords';
 
-const DB_SELECT_COMMAND_RE = /^(?:use|\\c(?:onnect)?)\s+([^;\s]+);?$/i;
-const TX_BEGIN_COMMAND_RE = /^begin;$/i
-const TX_ROLLBACK_COMMAND_RE = /^rollback;$/i
-const TX_COMMIT_COMMAND_RE = /^commit;$/i
+const DB_SELECT_COMMAND_RE = /^\s*(?:use|\\c(?:onnect)?)\s+(`?)([^;\s]+)\1;?\s*$/i;
+const TX_BEGIN_COMMAND_RE = /^\s*begin\s*;$/i
+const TX_ROLLBACK_COMMAND_RE = /^\s*rollback\s*;$/i
+const TX_COMMIT_COMMAND_RE = /^\s*commit\s*;$/i
+const COMMENT_BEGIN_RE = /\/\*(?!!)/;
+const COMMENT_END_RE = /\*\//;
+const IN_COMMENT_END_RE = /(.|\n|\r)*?\*\//m;
+const START_COMMENT_RE = /\/\*(?!!)(.|\n|\r)*/m;
+const FULL_COMMENT_RE = /\/\*(?!!)(.|\n|\r)*?\*\//mg;
+
+// Comments that start like like /*! or /*!57384 contain statements that *should* be executed by MySQL clients. See https://dev.mysql.com/doc/refman/en/comments.html.
+const FULL_NON_MYSQL_COMMENT_RE = /\/\*!\d*((?:.|\n|\r)*?)\*\//mg;
 
 let keywords: Promise<Keywords>;
 
+let currentStatementPromiseResolver: (value: Statement) => void;
+const statementPromises: Promise<Statement>[] = [
+  new Promise(resolve => currentStatementPromiseResolver = resolve)
+];
+
 export default async function query(awsConfig: AWSConfig, awsInfo: AWSInfo, clusterId: string, secretName: string, database?: string): Promise<void> {
   const resourceArn = `arn:${awsInfo.partition}:rds:${awsInfo.region}:${awsInfo.accountId}:cluster:${clusterId}`;
   const secretArn = `arn:${awsInfo.partition}:secretsmanager:${awsInfo.region}:${awsInfo.accountId}:secret:${secretName}`;
@@ -61,121 +81,49 @@ export default async function query(awsConfig: AWSConfig, awsInfo: AWSInfo, clus
 
   keywords = fetchKeywords(rdsDataClient, resourceArn, secretArn, database);
 
-  let transactionId: string | undefined;
+  // Start function to process statements sequentially
+  processStatements(rdsDataClient, resourceArn, secretArn, database);
 
-  replStart({
-    eval: async function (this: REPLServer, command: string, context: Context, file: string, cb: (err: Error | null, result: string | Record<string, unknown>) => void) {
-      command = command.trim();
+  let partial = '';
+  let partialInComment = false;
 
-      const dbMatch = command.match(DB_SELECT_COMMAND_RE);
-      if (dbMatch) {
-        database = dbMatch[1];
-        keywords = fetchKeywords(rdsDataClient, resourceArn, secretArn, database);
-        cb(null, `Now using database ${database}`);
-        return;
-      }
-
-      if (TX_BEGIN_COMMAND_RE.test(command)) {
-        if (transactionId) {
-          cb(null, `Error: Transaction '${transactionId}' currently in progress, cannot create a new one`);
-          return;
-        }
+  const replServer = replStart({
+    eval: function (this: REPLServer, command: string, context: Context, file: string, cb: REPLCallback) {
+      partial += command;
 
-        try {
-          ({ transactionId } = await rdsDataClient.send(
-            new BeginTransactionCommand({
-              resourceArn,
-              secretArn,
-              database
-            })
-          ));
-          cb(null, `Transaction '${transactionId}' begun`);
-        } catch (err) {
-          cb(null, `Failed to begin transaction: ${err.message} (${err.code})`);
-        }
-
-        return;
+      if (partialInComment && COMMENT_END_RE.test(partial)) {
+        partial.replace(IN_COMMENT_END_RE, '');
+        partialInComment = false;
       }
 
-      if (TX_ROLLBACK_COMMAND_RE.test(command)) {
-        if (!transactionId) {
-          cb(null, `Error: No transaction currently in progress`);
-          return;
-        }
-
-        try {
-          await rdsDataClient.send(
-            new RollbackTransactionCommand({
-              resourceArn,
-              secretArn,
-              transactionId
-            })
-          );
-          cb(null, `Transaction '${transactionId}' rolled back`);
-          transactionId = undefined;
-        } catch (err) {
-          cb(null, `Failed to rollback transaction '${transactionId}': ${err.message} (${err.code})`);
-        }
+      partial = partial
+        .replace(FULL_COMMENT_RE, '')
+        .replace(FULL_NON_MYSQL_COMMENT_RE, '$1');
 
-        return;
+      if (!partialInComment && COMMENT_BEGIN_RE.test(partial)) {
+        partial.replace(START_COMMENT_RE, '');
+        partialInComment = true;
       }
 
-      if (TX_COMMIT_COMMAND_RE.test(command)) {
-        if (!transactionId) {
-          cb(null, `Error: No transaction currently in progress`);
-          return;
-        }
+      const statements = partial.split(';');
 
-        try {
-          await rdsDataClient.send(
-            new CommitTransactionCommand({
-              resourceArn,
-              secretArn,
-              transactionId
-            })
-          );
-          cb(null, `Transaction '${transactionId}' committed`);
-          transactionId = undefined;
-        } catch (err) {
-          cb(null, `Failed to commit transaction '${transactionId}': ${err.message} (${err.code})`);
-        }
+      partial = statements.pop() as string;
 
-        return;
+      // `use <database>` is a command and doesn't require a semicolon terminator, check for it specially
+      if (DB_SELECT_COMMAND_RE.test(partial)) {
+        statements.push(partial);
+        partial = '';
       }
 
-      let records: Field[][] | undefined;
-      let columnMetadata: ColumnMetadata[] | undefined;
-      let numberOfRecordsUpdated;
-      try {
-        ({ records, columnMetadata, numberOfRecordsUpdated } = await rdsDataClient.send(
-          new ExecuteStatementCommand({
-            resourceArn,
-            secretArn,
-            database,
-            sql: command,
-            includeResultMetadata: true,
-            transactionId
-          })
-        ));
-      } catch (err) {
-        cb(null, `Failed to execute statement: ${err.message} (${err.code})`);
-        return;
-      }
-
-      const output: {
-        Records?: Row[]
-        'Record Count'?: number
-        'Number of Affected Records'?: number
-      } = {};
-
-      if (records) {
-        output.Records = hydrateRecords(records, columnMetadata as ColumnMetadata[]);
-        output['Record Count'] = output.Records.length
-      } else if (typeof numberOfRecordsUpdated === 'number') {
-        output['Number of Affected Records'] = numberOfRecordsUpdated
+      for (const statement of statements) {
+        // Enqueue command to run statements sequentially
+        const resolver = currentStatementPromiseResolver;
+        statementPromises.push(new Promise(resolve => currentStatementPromiseResolver = resolve));
+        resolver({
+          command: statement,
+          cb
+        });
       }
-
-      cb(null, output);
     },
 
     // Loosely based on the mysql CLI completion algorithm
@@ -224,6 +172,145 @@ export default async function query(awsConfig: AWSConfig, awsInfo: AWSInfo, clus
         callback(err);
       }
     },
-    writer: prettyoutput
+    writer: prettyoutput,
+    preview: true
   });
+
+  replServer.on('reset', () => {
+    partial = '';
+    partialInComment = false;
+  });
+}
+
+function handleError (cb: REPLCallback, message: string) {
+  if (process.stdin.isTTY) {
+    cb(null, message);
+  } else {
+    console.error(message);
+    process.exit(1);
+  }
+}
+
+async function processStatements (rdsDataClient: RDSDataClient, resourceArn: string, secretArn: string, database?: string) {
+  let transactionId: string | undefined;
+
+  for await (const { command, cb } of statementPromises) {
+    if (/^\s*$/.test(command)) {
+      cb(null);
+      continue;
+    }
+
+    const dbMatch = command.match(DB_SELECT_COMMAND_RE);
+    if (dbMatch) {
+      database = dbMatch[2];
+      keywords = fetchKeywords(rdsDataClient, resourceArn, secretArn, database);
+      cb(null, `Now using database ${database}`);
+      continue;
+    }
+
+    if (TX_BEGIN_COMMAND_RE.test(command)) {
+      if (transactionId) {
+        handleError(cb, `Error: Transaction '${transactionId}' currently in progress, cannot create a new one`);
+        continue;
+      }
+
+      try {
+        ({ transactionId } = await rdsDataClient.send(
+          new BeginTransactionCommand({
+            resourceArn,
+            secretArn,
+            database
+          })
+        ));
+        cb(null, `Transaction '${transactionId}' begun`);
+      } catch (err) {
+        handleError(cb, `Failed to begin transaction: ${err.message} (${err.code})`);
+      }
+
+      continue;
+    }
+
+    if (TX_ROLLBACK_COMMAND_RE.test(command)) {
+      if (!transactionId) {
+        cb(null, `Error: No transaction currently in progress`);
+        continue;
+      }
+
+      try {
+        await rdsDataClient.send(
+          new RollbackTransactionCommand({
+            resourceArn,
+            secretArn,
+            transactionId
+          })
+        );
+        cb(null, `Transaction '${transactionId}' rolled back`);
+        transactionId = undefined;
+      } catch (err) {
+        handleError(cb, `Failed to rollback transaction '${transactionId}': ${err.message} (${err.code})`);
+      }
+
+      continue;
+    }
+
+    if (TX_COMMIT_COMMAND_RE.test(command)) {
+      if (!transactionId) {
+        handleError(cb, `Error: No transaction currently in progress`);
+        continue;
+      }
+
+      try {
+        await rdsDataClient.send(
+          new CommitTransactionCommand({
+            resourceArn,
+            secretArn,
+            transactionId
+          })
+        );
+        cb(null, `Transaction '${transactionId}' committed`);
+        transactionId = undefined;
+      } catch (err) {
+        handleError(cb, `Failed to commit transaction '${transactionId}': ${err.message} (${err.code})`);
+      }
+
+      continue;
+    }
+
+    let records: Field[][] | undefined;
+    let columnMetadata: ColumnMetadata[] | undefined;
+    let numberOfRecordsUpdated;
+    try {
+      ({ records, columnMetadata, numberOfRecordsUpdated } = await rdsDataClient.send(
+        new ExecuteStatementCommand({
+          resourceArn,
+          secretArn,
+          database,
+          sql: command,
+          includeResultMetadata: true,
+          transactionId
+        })
+      ));
+    } catch (err) {
+      handleError(cb, `Failed to execute statement: ${err.message} (${err.code})`);
+      continue;
+    }
+
+    const output: {
+      Records?: Row[]
+      'Record Count'?: number
+      'Statement': string
+      'Number of Affected Records'?: number
+    } = {
+      Statement: command
+    };
+
+    if (records) {
+      output.Records = hydrateRecords(records, columnMetadata as ColumnMetadata[]);
+      output['Record Count'] = output.Records.length
+    } else if (typeof numberOfRecordsUpdated === 'number') {
+      output['Number of Affected Records'] = numberOfRecordsUpdated
+    }
+
+    cb(null, output);
+  }
 }
