diff --git a/CHANGELOG.md b/CHANGELOG.md index 221e263..d799373 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## 0.0.10 +- Added `disablePagination` connection setting to turn off the new query-pagination behavior entirely and go back to fetching the full result set in one call +- Added `disablePaginationCount` connection setting to skip the exact `COUNT(1)` query and always show an estimated total instead +- Fixed a bug where moving from page 1 to page 2 of a paginated result opened a new results tab while the original tab was left stuck on a loading spinner (subsequent page changes were unaffected). Cause: a fresh random `resultId` was generated for every page; it's now generated once per query run and kept stable across all of its pages +- Console `SELECT`/`WITH` queries now support pagination (`LIMIT n+1 OFFSET m`), with an exact row total via a `COUNT(1)` wrapper subquery (falling back to a look-ahead estimate), cached per query/run to avoid re-counting on every page turn +- Added `previewLimit` connection setting used as the default page size +- Added `connectionInitSql` connection setting: SQL executed once in a BigQuery session (`createSession: true`) right after the connection opens; the resulting `session_id` is reused on subsequent queries; failure aborts the connection +- No-result DML/DDL statements (`INSERT`/`UPDATE`/`DELETE`/`MERGE`/`CREATE`/`DROP`/`ALTER`/...) now show a one-row `Statement`/`Result` grid instead of a blank "No rows returned" grid; empty `SELECT`s are unaffected +- Internal metadata/explorer/"Show Records" queries are now tagged and excluded from the new pagination logic so they aren't truncated to the page size + ## 0.0.9 - Adds support for ARRAY results - Upgrades BigQuery (@google-cloud/bigquery) to 7.9.0 diff --git a/connection.schema.json b/connection.schema.json index 39e8015..c8502a4 100644 --- a/connection.schema.json +++ b/connection.schema.json @@ -28,6 +28,30 @@ "type": "string", "minLength": 1, "default": "us" + }, + "previewLimit": { + "title": "Preview Limit", + "description": "Number of rows fetched per page when browsing query results or table data", + "type": "integer", + "minimum": 1, + "default": 50 + }, + "connectionInitSql": { + "title": "Connection Init SQL", + "description": "SQL script executed once in a BigQuery session when the connection is opened (e.g. DECLARE/SET session variables, temp tables). Requires BigQuery session support.", + "type": "string" + }, + "disablePagination": { + "title": "Disable Query Pagination", + "description": "Disable automatic LIMIT/OFFSET pagination for console SELECT/WITH queries and always fetch the full result set in one go", + "type": "boolean", + "default": false + }, + "disablePaginationCount": { + "title": "Disable Exact Row Count", + "description": "Skip the extra COUNT(1) query used to show an exact total number of rows when paginating; shows an estimate instead (\"at least N rows\")", + "type": "boolean", + "default": false } }, "dependencies": { diff --git a/package-lock.json b/package-lock.json index f04a7d2..074554b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "sqltools-bigquery-driver", - "version": "0.0.9", + "version": "0.0.11", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "sqltools-bigquery-driver", - "version": "0.0.9", + "version": "0.0.11", "license": "MIT", "devDependencies": { "@sqltools/base-driver": "latest", diff --git a/package.json b/package.json index 8a57510..b6ddf8f 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "sqltools-bigquery-driver", "displayName": "BigQuery Driver for SQLTools", "description": "Run Queries and Explore your BigQuery Database in VSCode", - "version": "0.0.9", + "version": "0.0.10", "engines": { "vscode": "^1.42.0" }, diff --git a/src/ls/driver.ts b/src/ls/driver.ts index 54ebf55..d991e25 100644 --- a/src/ls/driver.ts +++ b/src/ls/driver.ts @@ -30,13 +30,15 @@ export default class BigQueryDriver extends AbstractDriver; public async open() { const BigQuery = this.requireDep('@google-cloud/bigquery').BigQuery; const OAuth2Client = this.requireDep('google-auth-library').OAuth2Client; const getCredentials = () => { - + const authentication_method = this.credentials.authenticator if (authentication_method === 'CLI') { return { @@ -47,7 +49,7 @@ export default class BigQueryDriver extends AbstractDriver { + public singleQuery: (typeof AbstractDriver)['prototype']['singleQuery'] = ((query: any, opt: any) => { + return this.query(query, { ...opt, __internal: true }).then(([res]) => res); + }) as any; + + private _isPaginatableSelect(sql: string): boolean { + if (this.credentials.disablePagination) return false; + const withoutComments = sql.toString().replace(/^\s*(--[^\n]*\n)+/, '').trim(); + if (!/^(SELECT|WITH)\b/i.test(withoutComments)) return false; + const withoutTrailingSemi = withoutComments.replace(/;\s*$/, ''); + // reject multiple statements + return !/;\s*\S/.test(withoutTrailingSemi); + } + + private _stripTrailingSemicolon(sql: string): string { + return sql.toString().replace(/;\s*$/, ''); + } + + private _buildConnectionProperties() { + if (this._sessionId) { + return [{ key: 'session_id', value: this._sessionId }]; + } + return undefined; + } + + private _buildDmlOutcomeMessage(statementType: string, metadata: any): string { + const queryStats = metadata && metadata.statistics && metadata.statistics.query; + let affected: number | undefined; + const dmlStats = queryStats && queryStats.dmlStats; + if (dmlStats) { + affected = Number(dmlStats.insertedRowCount || 0) + Number(dmlStats.updatedRowCount || 0) + Number(dmlStats.deletedRowCount || 0); + } else if (queryStats && queryStats.numDmlAffectedRows !== undefined) { + affected = Number(queryStats.numDmlAffectedRows); + } + const label = statementType || 'Statement'; + return `${label} executed successfully.${affected !== undefined ? ` ${affected} rows were affected.` : ''}`; + } + + private async _getPaginationState(bigquery: any, baseSql: string, baseOptions: any, requestId: string, page: number, offset: number, rowsLen: number, hasMore: boolean) { + this._paginationCache = this._paginationCache || new Map(); + const key = `${requestId} ${baseSql}`; + const cached = this._paginationCache.get(key); + const resultId = (cached && cached.resultId) || generateId(); + + if (cached && cached.exact) { + return { total: cached.total, exact: true, resultId }; + } + + const estimate = offset + rowsLen + (hasMore ? 1 : 0); + let total = estimate; + let exact = false; + + const skipCount = !!this.credentials.disablePaginationCount; + if (page === 0 && !skipCount) { + try { + const [job] = await bigquery.createQueryJob({ ...baseOptions, query: `SELECT COUNT(1) AS total FROM (${baseSql})` }); + const [countRows] = await job.getQueryResults(); + total = Number(countRows[0].total); + exact = true; + } catch (error) { + total = estimate; + exact = false; + } + } + + if (this._paginationCache.size >= 100 && !this._paginationCache.has(key)) { + const firstKey = this._paginationCache.keys().next().value; + this._paginationCache.delete(firstKey); + } + this._paginationCache.set(key, { total, exact, resultId }); + return { total, exact, resultId }; + } + + private async _execPaginatedSelect(bigquery: any, rawSql: string, baseOptions: any, opt: any): Promise { + const page = opt.page || 0; + const pageSize = opt.pageSize || this.credentials.previewLimit || 50; + const offset = page * pageSize; + const base = this._stripTrailingSemicolon(rawSql); + const limitedSql = `${base} LIMIT ${pageSize + 1} OFFSET ${offset}`; + + const [job] = await bigquery.createQueryJob({ ...baseOptions, query: limitedSql }); + const [rows] = await job.getQueryResults(); + const hasMore = rows.length > pageSize; + const pageRows = hasMore ? rows.slice(0, pageSize) : rows; + const standardizedRows = await standardizeResult(pageRows); + + const { total, exact, resultId } = await this._getPaginationState(bigquery, base, baseOptions, opt.requestId, page, offset, pageRows.length, hasMore); + const message = exact + ? `Showing page ${page + 1} of ${Math.max(1, Math.ceil(total / pageSize))} (${total} rows).` + : `Showing page ${page + 1} (at least ${total} rows).`; + + return { + cols: standardizedRows && standardizedRows.length ? Object.keys(standardizedRows[0]) : ['No rows returned'], + connId: this.getId(), + messages: [{ date: new Date(), message }], + results: standardizedRows, + query: rawSql, + requestId: opt.requestId, + resultId, + page, + pageSize, + total, + queryType: 'executeQuery', + queryParams: base, + } as unknown as NSDatabase.IResult; + } + + public query: (typeof AbstractDriver)['prototype']['query'] = async (query, opt: any = {}) => { await this.open(); const bigquery = await this.connection; - const options = { - // typescript complains if this is not an array - query: [query], + const rawSql = String(query); + const baseOptions: any = { location: this.credentials.location, }; + const connectionProperties = this._buildConnectionProperties(); + if (connectionProperties) { + baseOptions.connectionProperties = connectionProperties; + } + const resultsAgg: NSDatabase.IResult[] = []; - const [rows] = await bigquery.query(options); + if (!opt.__internal && this._isPaginatableSelect(rawSql)) { + resultsAgg.push(await this._execPaginatedSelect(bigquery, rawSql, baseOptions, opt)); + return resultsAgg; + } + + const [job] = await bigquery.createQueryJob({ ...baseOptions, query: rawSql }); + const [rows] = await job.getQueryResults(); + const [metadata] = await job.getMetadata(); + const statementType = metadata?.statistics?.query?.statementType; + const isSelectLike = !statementType || statementType === 'SELECT' || statementType === 'SCRIPT'; const standardizedRows = await standardizeResult(rows); + if (!Array.isArray(rows) || !rows.length) { + if (isSelectLike) { + resultsAgg.push({ + cols: ['No rows returned'], + connId: this.getId(), + messages: [{ date: new Date(), message: `Query executed successfully but no data was returned` }], + results: [], + query: rawSql, + requestId: opt.requestId, + resultId: generateId(), + }); + } else { + const outcome = this._buildDmlOutcomeMessage(statementType, metadata); + resultsAgg.push({ + cols: ['Statement', 'Result'], + connId: this.getId(), + messages: [{ date: new Date(), message: outcome }], + results: [{ Statement: rawSql, Result: outcome }], + query: rawSql, + requestId: opt.requestId, + resultId: generateId(), + }); + } + } else { resultsAgg.push({ - cols: ['No rows returned'], + cols: standardizedRows && standardizedRows.length && Object.keys(standardizedRows[0]), connId: this.getId(), - messages: [{ date: new Date(), message: `Query executed successfully but no data was returned` }], - results: [], - // back to string - query: query[0], + messages: [{ date: new Date(), message: `Query executed successfully` }], + results: standardizedRows, + query: rawSql, requestId: opt.requestId, resultId: generateId(), - }) - } else { - resultsAgg.push({ - cols: standardizedRows && standardizedRows.length && Object.keys(standardizedRows[0]), - connId: this.getId(), - messages: [{ date: new Date(), message: `Query executed successfully` }], - results: standardizedRows, - // back to string - query: query[0], - requestId: opt.requestId, - resultId: generateId(), - }); - } + }); + } return resultsAgg; } diff --git a/ui.schema.json b/ui.schema.json index 6155165..1978749 100644 --- a/ui.schema.json +++ b/ui.schema.json @@ -3,13 +3,18 @@ "authenticator", "keyfile", "projectId", - "token" + "token", + "location", + "previewLimit", + "disablePagination", + "disablePaginationCount", + "connectionInitSql" ], "authenticator": { "ui:autofocus": true, "ui:help": "See Authenticator Guide: https://docs.evidence.dev/core-concepts/data-sources/#bigquery" }, - "keyfile": { + "keyfile": { "ui:widget": "file" }, "projectId": { @@ -19,6 +24,16 @@ }, "location": { "ui:help": "Default `us`. Each connection can access one location. Full location list: https://cloud.google.com/bigquery/docs/locations" + }, + "connectionInitSql": { + "ui:widget": "textarea", + "ui:help": "Runs once in a BigQuery session right after the connection opens. If it fails, the connection is aborted." + }, + "disablePagination": { + "ui:help": "Turn off if you prefer the old behavior (full result set, no page controls) or if pagination causes issues with certain queries." + }, + "disablePaginationCount": { + "ui:help": "Turn on to skip the extra COUNT(1) query on the first page of results (faster, but the row total shown will be an estimate)." } } \ No newline at end of file