SYNC: Limit query result rows #1098

packages/api/src/controllers/sessions.js

@@ -141,7 +141,7 @@ module.exports = {
   },
 
   executeQuery_meta: true,
-  async executeQuery({ sesid, sql, autoCommit }) {
+  async executeQuery({ sesid, sql, autoCommit, limitRows }) {
     const session = this.opened.find(x => x.sesid == sesid);
     if (!session) {
       throw new Error('Invalid session');
@@ -149,7 +149,7 @@ module.exports = {
 
     logger.info({ sesid, sql }, 'Processing query');
     this.dispatchMessage(sesid, 'Query execution started');
-    session.subprocess.send({ msgtype: 'executeQuery', sql, autoCommit });
+    session.subprocess.send({ msgtype: 'executeQuery', sql, autoCommit, limitRows });
 
     return { state: 'ok' };
   },

packages/api/src/proc/sessionProcess.js

@@ -117,7 +117,7 @@ async function handleExecuteControlCommand({ command }) {
   }
 }
 
-async function handleExecuteQuery({ sql, autoCommit }) {
+async function handleExecuteQuery({ sql, autoCommit, limitRows }) {
   lastActivity = new Date().getTime();
 
   await waitConnected();
@@ -146,7 +146,7 @@ async function handleExecuteQuery({ sql, autoCommit }) {
       ...driver.getQuerySplitterOptions('stream'),
       returnRichInfo: true,
     })) {
-      await handleQueryStream(dbhan, driver, queryStreamInfoHolder, sqlItem);
+      await handleQueryStream(dbhan, driver, queryStreamInfoHolder, sqlItem, undefined, limitRows);
       // const handler = new StreamHandler(resultIndex);
       // const stream = await driver.stream(systemConnection, sqlItem, handler);
       // handler.stream = stream;

packages/api/src/utility/handleQueryStream.js

@@ -82,20 +82,27 @@ class QueryStreamTableWriter {
   }
 
   close(afterClose) {
-    if (this.currentStream) {
-      this.currentStream.end(() => {
-        this.writeCurrentStats(true, true);
-        if (afterClose) afterClose();
-      });
-    }
+    return new Promise(resolve => {
+      if (this.currentStream) {
+        this.currentStream.end(() => {
+          this.writeCurrentStats(true, true);
+          if (afterClose) afterClose();
+          resolve();
+        });
+      } else {
+        resolve();
+      }
+    });
   }
 }
 
 class StreamHandler {
-  constructor(queryStreamInfoHolder, resolve, startLine, sesid = undefined) {
+  constructor(queryStreamInfoHolder, resolve, startLine, sesid = undefined, limitRows = undefined) {
     this.recordset = this.recordset.bind(this);
     this.startLine = startLine;
     this.sesid = sesid;
+    this.limitRows = limitRows;
+    this.rowsLimitOverflow = false;
     this.row = this.row.bind(this);
     // this.error = this.error.bind(this);
     this.done = this.done.bind(this);
@@ -107,6 +114,7 @@ class StreamHandler {
     this.plannedStats = false;
     this.queryStreamInfoHolder = queryStreamInfoHolder;
     this.resolve = resolve;
+    this.rowCounter = 0;
     // currentHandlers = [...currentHandlers, this];
   }
 
@@ -118,6 +126,9 @@ class StreamHandler {
   }
 
   recordset(columns) {
+    if (this.rowsLimitOverflow) {
+      return;
+    }
     this.closeCurrentWriter();
     this.currentWriter = new QueryStreamTableWriter(this.sesid);
     this.currentWriter.initializeFromQuery(
@@ -125,6 +136,7 @@ class StreamHandler {
       this.queryStreamInfoHolder.resultIndex
     );
     this.queryStreamInfoHolder.resultIndex += 1;
+    this.rowCounter = 0;
 
     // this.writeCurrentStats();
 
@@ -135,8 +147,36 @@ class StreamHandler {
     // }, 500);
   }
   row(row) {
-    if (this.currentWriter) this.currentWriter.row(row);
-    else if (row.message) process.send({ msgtype: 'info', info: { message: row.message }, sesid: this.sesid });
+    if (this.rowsLimitOverflow) {
+      return;
+    }
+
+    if (this.limitRows && this.rowCounter >= this.limitRows) {
+      process.send({
+        msgtype: 'info',
+        info: { message: `Rows limit overflow, loaded ${this.rowCounter} rows, canceling query`, severity: 'error' },
+        sesid: this.sesid,
+      });
+      this.rowsLimitOverflow = true;
+
+      this.queryStreamInfoHolder.canceled = true;
+      if (this.currentWriter) {
+        this.currentWriter.close().then(() => {
+          process.exit(0);
+        });
+      } else {
+        process.exit(0);
+      }
+
+      return;
+    }
+
+    if (this.currentWriter) {
+      this.currentWriter.row(row);
+      this.rowCounter += 1;
+    } else if (row.message) {
+      process.send({ msgtype: 'info', info: { message: row.message }, sesid: this.sesid });
+    }
     // this.onRow(this.jslid);
   }
   // error(error) {
@@ -161,10 +201,10 @@ class StreamHandler {
   }
 }
 
-function handleQueryStream(dbhan, driver, queryStreamInfoHolder, sqlItem, sesid = undefined) {
+function handleQueryStream(dbhan, driver, queryStreamInfoHolder, sqlItem, sesid = undefined, limitRows = undefined) {
   return new Promise((resolve, reject) => {
     const start = sqlItem.trimStart || sqlItem.start;
-    const handler = new StreamHandler(queryStreamInfoHolder, resolve, start && start.line, sesid);
+    const handler = new StreamHandler(queryStreamInfoHolder, resolve, start && start.line, sesid, limitRows);
     driver.stream(dbhan, sqlItem.text, handler);
   });
 }

packages/web/src/icons/FontIcon.svelte

@@ -222,6 +222,7 @@
 
     'icon premium': 'mdi mdi-star',
     'icon upload': 'mdi mdi-upload',
+    'icon limit': 'mdi mdi-car-speed-limiter',
 
     'img ok': 'mdi mdi-check-circle color-icon-green',
     'img ok-inv': 'mdi mdi-check-circle color-icon-inv-green',

packages/web/src/modals/RowsLimitModal.svelte

@@ -0,0 +1,41 @@
+<script lang="ts">
+  import FormStyledButton from '../buttons/FormStyledButton.svelte';
+
+  import FormProvider from '../forms/FormProvider.svelte';
+  import FormSubmit from '../forms/FormSubmit.svelte';
+  import FormTextField from '../forms/FormTextField.svelte';
+  import ModalBase from './ModalBase.svelte';
+  import { closeCurrentModal } from './modalTools';
+
+  export let value;
+  export let onConfirm;
+
+  const handleSubmit = async value => {
+    closeCurrentModal();
+    onConfirm(value);
+  };
+</script>
+
+<FormProvider initialValues={{ value }}>
+  <ModalBase {...$$restProps}>
+    <svelte:fragment slot="header">Rows limit</svelte:fragment>
+
+    <FormTextField
+      label="Return only N rows from query"
+      name="value"
+      focused
+      data-testid="RowsLimitModal_value"
+      placeholder="(No rows limit)"
+    />
+
+    <svelte:fragment slot="footer">
+      <FormSubmit
+        value="OK"
+        on:click={e => handleSubmit(parseInt(e.detail.value) || null)}
+        data-testid="RowsLimitModal_setLimit"
+      />
+      <FormStyledButton value="Set no limit" on:click={e => handleSubmit(null)} data-testid="RowsLimitModal_setNoLimit" />
+      <FormStyledButton type="button" value="Cancel" on:click={closeCurrentModal} data-testid="RowsLimitModal_cancel" />
+    </svelte:fragment>
+  </ModalBase>
+</FormProvider>

packages/web/src/settings/SettingsModal.svelte

@@ -227,6 +227,12 @@ ORDER BY
               </FormFieldTemplateLarge>
             </div>
           </div>
+
+          <FormTextField
+            name="sqlEditor.limitRows"
+            label="Return only N rows from query"
+            placeholder="(No rows limit)"
+          />
         </svelte:fragment>
         <svelte:fragment slot="2">
           <div class="heading">Connection</div>

packages/web/src/tabs/QueryTab.svelte

@@ -144,6 +144,9 @@
   import HorizontalSplitter from '../elements/HorizontalSplitter.svelte';
   import QueryAiAssistant from '../query/QueryAiAssistant.svelte';
   import uuidv1 from 'uuid/v1';
+  import ToolStripButton from '../buttons/ToolStripButton.svelte';
+  import { getIntSettingsValue } from '../settings/settingsTools';
+  import RowsLimitModal from '../modals/RowsLimitModal.svelte';
 
   export let tabid;
   export let conid;
@@ -197,6 +200,21 @@
   let isInTransaction = false;
   let isAutocommit = false;
 
+  const queryRowsLimitLocalStorageKey = `tabdata_limitRows_${tabid}`;
+  function getInitialRowsLimit() {
+    const storageValue = localStorage.getItem(queryRowsLimitLocalStorageKey);
+    if (storageValue == 'nolimit') {
+      return null;
+    }
+    if (storageValue) {
+      return parseInt(storageValue) ?? null;
+    }
+    return getIntSettingsValue('sqlEditor.limitRows', null, 1);
+  }
+
+  let queryRowsLimit = getInitialRowsLimit();
+  $: localStorage.setItem(queryRowsLimitLocalStorageKey, queryRowsLimit ? queryRowsLimit.toString() : 'nolimit');
+
   onMount(() => {
     intervalId = setInterval(() => {
       if (!driver?.singleConnectionOnly && sessionId) {
@@ -362,6 +380,7 @@
         sesid,
         sql,
         autoCommit: driver?.implicitTransactions && isAutocommit,
+        limitRows: queryRowsLimit ? queryRowsLimit : undefined,
       });
     }
     await apiCall('query-history/write', {
@@ -713,6 +732,20 @@
     <ToolStripCommandButton command="query.kill" data-testid="QueryTab_killButton" />
     <ToolStripSaveButton idPrefix="query" />
     <ToolStripCommandButton command="query.formatCode" />
+    {#if !driver?.singleConnectionOnly}
+      <ToolStripButton
+        icon="icon limit"
+        on:click={() =>
+          showModal(RowsLimitModal, {
+            value: queryRowsLimit,
+            onConfirm: value => {
+              queryRowsLimit = value;
+            },
+          })}
+      >
+        {queryRowsLimit ? `Limit ${queryRowsLimit} rows` : 'Unlimited rows'}</ToolStripButton
+      >
+    {/if}
     {#if resultCount == 1}
       <ToolStripExportButton command="jslTableGrid.export" {quickExportHandlerRef} label="Export result" />
     {/if}