add query timeout configuration option (#27)

Allow user to specify query timeout in the vscode settings menu. Default to 30 seconds.

Author: cmoog

CHANGELOG.md

@@ -6,3 +6,4 @@
   - When running on a compatible arch/os, notebooks will now
     benefit from enhacned typed autocomplete and hover information
     when connected to a valid database connection.
+- New configuration option for query timeout in milliseconds. Defaults to 30000.

README.md

@@ -15,8 +15,8 @@ and view output interactively.
 - Open any `.sql` file as a Notebook.
 - Execute query blocks in the Notebook UI and view output.
 - Configure database connections in the SQL Notebook sidepanel.
-- Supports MySQL PostgreSQL, and MSSQL (Oracle support coming soon).
-- (coming soon) Built-in typed autocomplete.
+- Supports MySQL PostgreSQL, and MSSQL (OracleDB and SQLite support coming soon).
+- (unstable) Built-in typed autocomplete with an embeded language server.
 
 ## Usage
 

package.json

@@ -97,6 +97,11 @@
           "type": "boolean",
           "default": false,
           "description": "(Unstable) Use embeded language server for intelligent completion and hover information."
+        },
+        "SQLNotebook.queryTimeout": {
+          "type": "number",
+          "default": 30000,
+          "description": "Query timeout in milliseconds for cell query execution."
         }
       }
     },

src/commands.ts

@@ -69,7 +69,11 @@ export const connectToDatabase =
     }
 
     try {
-      globalConnPool.pool = await getPool({ ...match, password } as PoolConfig);
+      globalConnPool.pool = await getPool({
+        ...match,
+        password,
+        queryTimeout: getQueryTimeoutConfiguration(),
+      } as PoolConfig);
       const conn = await globalConnPool.pool.getConnection();
       await conn.query('SELECT 1'); // essentially a ping to see if the connection works
       connectionsSidepanel.setActive(match.name);
@@ -125,3 +129,11 @@ function shouldUseLanguageServer(): boolean {
     false
   );
 }
+
+function getQueryTimeoutConfiguration() {
+  const defaultTimeout = 30000; // make this the same as the package.json-level configuration default
+  return (
+    vscode.workspace.getConfiguration('SQLNotebook').get('queryTimeout') ??
+    defaultTimeout
+  );
+}

src/driver.ts

@@ -40,6 +40,8 @@ interface BaseConfig {
   user: string;
   password?: string;
   database?: string;
+
+  queryTimeout: number;
 }
 
 interface MySQLConfig extends BaseConfig {
@@ -67,6 +69,7 @@ async function createMySQLPool({
   password,
   database,
   multipleStatements,
+  queryTimeout,
 }: MySQLConfig): Promise<Pool> {
   return mysqlPool(
     mysql.createPool({
@@ -76,28 +79,32 @@ async function createMySQLPool({
       password,
       database,
       multipleStatements,
-    })
+    }),
+    queryTimeout
   );
 }
 
-function mysqlPool(pool: mysql.Pool): Pool {
+function mysqlPool(pool: mysql.Pool, queryTimeout: number): Pool {
   return {
     async getConnection(): Promise<Conn> {
-      return mysqlConn(await pool.getConnection());
+      return mysqlConn(await pool.getConnection(), queryTimeout);
     },
     end() {
       pool.end();
     },
   };
 }
 
-function mysqlConn(conn: mysql.PoolConnection): Conn {
+function mysqlConn(conn: mysql.PoolConnection, queryTimeout: number): Conn {
   return {
     destroy() {
       conn.destroy();
     },
     async query(q: string): Promise<ExecutionResult> {
-      const [result, ok] = (await conn.query(q)) as any;
+      const [result, ok] = (await conn.query({
+        sql: q,
+        timeout: queryTimeout,
+      })) as any;
       console.debug('mysql query result', { result, ok });
 
       if (!result.length) {
@@ -132,13 +139,15 @@ async function createPostgresPool({
   user,
   password,
   database,
+  queryTimeout,
 }: PostgresConfig): Promise<Pool> {
   const pool = new pg.Pool({
     host,
     port,
     password,
     database,
     user,
+    query_timeout: queryTimeout,
   });
   return postgresPool(pool);
 }
@@ -197,6 +206,7 @@ async function createMSSQLPool(config: MSSQLConfig): Promise<Pool> {
     user: config.user,
     password: config.password,
     database: config.database,
+    requestTimeout: config.queryTimeout,
     options: {
       encrypt: config.encrypt,
       trustServerCertificate: config.trustServerCertificate,