diff --git a/packages/api/src/controllers/sessions.js b/packages/api/src/controllers/sessions.js index 08e3e11e2..4f2751483 100644 --- a/packages/api/src/controllers/sessions.js +++ b/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' }; }, diff --git a/packages/api/src/proc/sessionProcess.js b/packages/api/src/proc/sessionProcess.js index 7560f30f9..8dd193db5 100644 --- a/packages/api/src/proc/sessionProcess.js +++ b/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; diff --git a/packages/api/src/utility/handleQueryStream.js b/packages/api/src/utility/handleQueryStream.js index 76e573a57..f57cb46fd 100644 --- a/packages/api/src/utility/handleQueryStream.js +++ b/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); }); } diff --git a/packages/web/src/icons/FontIcon.svelte b/packages/web/src/icons/FontIcon.svelte index 8190919fc..bd9d0df51 100644 --- a/packages/web/src/icons/FontIcon.svelte +++ b/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', diff --git a/packages/web/src/modals/RowsLimitModal.svelte b/packages/web/src/modals/RowsLimitModal.svelte new file mode 100644 index 000000000..12f62606f --- /dev/null +++ b/packages/web/src/modals/RowsLimitModal.svelte @@ -0,0 +1,41 @@ + + + + + Rows limit + + + + + handleSubmit(parseInt(e.detail.value) || null)} + data-testid="RowsLimitModal_setLimit" + /> + handleSubmit(null)} data-testid="RowsLimitModal_setNoLimit" /> + + + + diff --git a/packages/web/src/settings/SettingsModal.svelte b/packages/web/src/settings/SettingsModal.svelte index 0cec20629..6e66932ab 100644 --- a/packages/web/src/settings/SettingsModal.svelte +++ b/packages/web/src/settings/SettingsModal.svelte @@ -227,6 +227,12 @@ ORDER BY + +
Connection
diff --git a/packages/web/src/tabs/QueryTab.svelte b/packages/web/src/tabs/QueryTab.svelte index 41ae484d2..da772adb9 100644 --- a/packages/web/src/tabs/QueryTab.svelte +++ b/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 @@ + {#if !driver?.singleConnectionOnly} + + showModal(RowsLimitModal, { + value: queryRowsLimit, + onConfirm: value => { + queryRowsLimit = value; + }, + })} + > + {queryRowsLimit ? `Limit ${queryRowsLimit} rows` : 'Unlimited rows'} + {/if} {#if resultCount == 1} {/if}