Add pagination and session-SQL support from 0.0.10 build

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

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": {

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"
   },

src/ls/driver.ts

@@ -30,13 +30,15 @@ export default class BigQueryDriver extends AbstractDriver<DriverLib, DriverOpti
 
   queries = queries;
 
+  private _sessionId?: string;
+  private _paginationCache?: Map<string, { total: number; exact: boolean; resultId: string }>;
 
   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<DriverLib, DriverOpti
         const access_token = this.credentials.token;
         const oauth = new OAuth2Client();
         oauth.setCredentials({ access_token });
-        
+
         return {
           // is this a legit way to handle this typescript error
           authClient: oauth as JSONClient,
@@ -73,13 +75,36 @@ export default class BigQueryDriver extends AbstractDriver<DriverLib, DriverOpti
       }
     });
 
+    const initSql = this.credentials.connectionInitSql;
+    if (initSql && !this._sessionId) {
+      const bigquery = await this.connection;
+      try {
+        const [job] = await bigquery.createQueryJob({
+          query: initSql,
+          location: this.credentials.location,
+          createSession: true,
+        });
+        await job.getQueryResults();
+        const [metadata] = await job.getMetadata();
+        const sessionId = metadata?.statistics?.sessionInfo?.sessionId;
+        if (!sessionId) {
+          throw new Error('BigQuery did not return a session id for the init session');
+        }
+        this._sessionId = sessionId;
+      } catch (error) {
+        this.connection = null;
+        throw new Error('Connection init SQL failed: ' + (error && error.message || error));
+      }
+    }
   }
 
 
   public async close() {
     if (!this.connection) return Promise.resolve();
 
     this.connection = null;
+    this._sessionId = undefined;
+    this._paginationCache = undefined;
   }
 
   public async testConnection() {
@@ -94,41 +119,172 @@ export default class BigQueryDriver extends AbstractDriver<DriverLib, DriverOpti
     }
   }
 
-  public query: (typeof AbstractDriver)['prototype']['query'] = async (query, opt = {}) => {
+  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<NSDatabase.IResult> {
+    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;
   }
 

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)."
   }
 
 }