@malloydata/db-bigquery 0.0.336 → 0.0.338

package/dist/bigquery_connection.d.ts

@@ -1,7 +1,7 @@
 import type { RowMetadata } from '@google-cloud/bigquery';
 import type bigquery from '@google-cloud/bigquery/build/src/types';
 import type { ResourceStream } from '@google-cloud/paginator';
-import type { Connection, ConnectionConfig, MalloyQueryData, PersistSQLResults, QueryData, QueryDataRow, QueryOptionsReader, QueryRunStats, RunSQLOptions, StreamingConnection, TableSourceDef, SQLSourceDef, SQLSourceRequest } from '@malloydata/malloy';
+import type { Connection, ConnectionConfig, MalloyQueryData, PersistSQLResults, QueryData, QueryRecord, QueryOptionsReader, QueryRunStats, RunSQLOptions, StreamingConnection, TableSourceDef, SQLSourceDef, SQLSourceRequest } from '@malloydata/malloy';
 import type { TableMetadata } from '@malloydata/malloy/connection';
 import { BaseConnection } from '@malloydata/malloy/connection';
 export interface BigQueryManagerOptions {
@@ -26,6 +26,7 @@ interface BigQueryConnectionConfiguration {
     timeoutMs?: string;
     billingProjectId?: string;
     credentials?: CredentialBody;
+    setupSQL?: string;
 }
 interface BigQueryConnectionOptions extends ConnectionConfig {
     /** This ID is used for Bigquery Table Normalization */
@@ -37,6 +38,7 @@ interface BigQueryConnectionOptions extends ConnectionConfig {
     billingProjectId?: string;
     client_email?: string;
     private_key?: string;
+    setupSQL?: string;
 }
 interface SchemaInfo {
     schema: bigquery.ITableFieldSchema;
@@ -58,9 +60,11 @@ export declare class BigQueryConnection extends BaseConnection implements Connec
     private queryOptions?;
     private config;
     private location?;
+    private setupSQL;
     constructor(option: BigQueryConnectionOptions, queryOptions?: QueryOptionsReader);
     constructor(name: string, queryOptions?: QueryOptionsReader, config?: BigQueryConnectionConfiguration);
     get dialectName(): string;
+    private prependSetupSQL;
     private readQueryOptions;
     canPersist(): this is PersistSQLResults;
     canStream(): this is StreamingConnection;
@@ -88,7 +92,7 @@ export declare class BigQueryConnection extends BaseConnection implements Connec
     private createBigQueryJobAndGetResults;
     private createBigQueryJob;
     initiateJobAndGetLinkToConsole(sqlCommand: string, dryRun?: boolean): Promise<string>;
-    runSQLStream(sqlCommand: string, { rowLimit, abortSignal }?: RunSQLOptions): AsyncIterableIterator<QueryDataRow>;
+    runSQLStream(sqlCommand: string, { rowLimit, abortSignal }?: RunSQLOptions): AsyncIterableIterator<QueryRecord>;
     fetchTableMetadata(tablePath: string): Promise<TableMetadata>;
     close(): Promise<void>;
 }

package/dist/bigquery_connection.js

@@ -129,10 +129,19 @@ class BigQueryConnection extends connection_1.BaseConnection {
         this.queryOptions = queryOptions;
         this.config = config;
         this.location = config.location;
+        this.setupSQL = config.setupSQL;
     }
     get dialectName() {
         return 'standardsql';
     }
+    prependSetupSQL(sql) {
+        if (!this.setupSQL)
+            return sql;
+        const setup = this.setupSQL.trimEnd().endsWith(';')
+            ? this.setupSQL
+            : this.setupSQL + ';';
+        return setup + '\n' + sql;
+    }
     readQueryOptions() {
         const options = BigQueryConnection.DEFAULT_QUERY_OPTIONS;
         if (this.queryOptions) {
@@ -154,9 +163,8 @@ class BigQueryConnection extends connection_1.BaseConnection {
         return true;
     }
     getDigest() {
-        var _a;
-        const data = `bigquery:${this.projectId}:${(_a = this.location) !== null && _a !== void 0 ? _a : 'US'}`;
-        return (0, malloy_1.makeDigest)(data);
+        var _a, _b;
+        return (0, malloy_1.makeDigest)('bigquery', this.billingProjectId, this.projectId, (_a = this.location) !== null && _a !== void 0 ? _a : 'US', (_b = this.setupSQL) !== null && _b !== void 0 ? _b : '');
     }
     get supportsNesting() {
         return true;
@@ -214,7 +222,7 @@ class BigQueryConnection extends connection_1.BaseConnection {
         try {
             const [result] = await this.bigQuery.createQueryJob({
                 location: this.location,
-                query: sqlCommand,
+                query: this.prependSetupSQL(sqlCommand),
                 dryRun: true,
             });
             return result;
@@ -354,7 +362,7 @@ class BigQueryConnection extends connection_1.BaseConnection {
             throw new Error(`Table ${tableName} already exists`);
         }
         const [job] = await this.bigQuery.createQueryJob({
-            query: sqlCommand,
+            query: this.prependSetupSQL(sqlCommand),
             location: this.location,
             destination: table,
         });
@@ -468,7 +476,7 @@ class BigQueryConnection extends connection_1.BaseConnection {
             try {
                 const [job] = await this.bigQuery.createQueryJob({
                     location: this.location,
-                    query: sqlRef.selectStr,
+                    query: this.prependSetupSQL(sqlRef.selectStr),
                     dryRun: true,
                 });
                 return job.metadata.statistics.query.schema;
@@ -529,11 +537,15 @@ class BigQueryConnection extends connection_1.BaseConnection {
         }
     }
     async createBigQueryJob(createQueryJobOptions) {
+        const options = { ...createQueryJobOptions };
+        if (options.query) {
+            options.query = this.prependSetupSQL(options.query);
+        }
         const [job] = await this.bigQuery.createQueryJob({
             location: this.location,
             maximumBytesBilled: this.config.maximumBytesBilled || MAXIMUM_BYTES_BILLED,
             jobTimeoutMs: Number(this.config.timeoutMs) || TIMEOUT_MS,
-            ...createQueryJobOptions,
+            ...options,
         });
         return job;
     }
@@ -557,7 +569,7 @@ class BigQueryConnection extends connection_1.BaseConnection {
                 }
             }
             this.bigQuery
-                .createQueryStream(sqlCommand)
+                .createQueryStream(this.prependSetupSQL(sqlCommand))
                 .on('error', onError)
                 .on('data', handleData)
                 .on('end', onEnd);

package/dist/index.js

@@ -25,4 +25,57 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.BigQueryConnection = void 0;
 var bigquery_connection_1 = require("./bigquery_connection");
 Object.defineProperty(exports, "BigQueryConnection", { enumerable: true, get: function () { return bigquery_connection_1.BigQueryConnection; } });
+const malloy_1 = require("@malloydata/malloy");
+const bigquery_connection_2 = require("./bigquery_connection");
+(0, malloy_1.registerConnectionType)('bigquery', {
+    factory: (config) => {
+        return new bigquery_connection_2.BigQueryConnection(config);
+    },
+    properties: [
+        {
+            name: 'projectId',
+            displayName: 'Project ID',
+            type: 'string',
+            optional: true,
+        },
+        {
+            name: 'serviceAccountKeyPath',
+            displayName: 'Service Account Key',
+            type: 'file',
+            optional: true,
+            fileFilters: { JSON: ['json'] },
+        },
+        {
+            name: 'location',
+            displayName: 'Location',
+            type: 'string',
+            optional: true,
+        },
+        {
+            name: 'maximumBytesBilled',
+            displayName: 'Maximum Bytes Billed',
+            type: 'string',
+            optional: true,
+        },
+        {
+            name: 'timeoutMs',
+            displayName: 'Timeout (ms)',
+            type: 'string',
+            optional: true,
+        },
+        {
+            name: 'billingProjectId',
+            displayName: 'Billing Project ID',
+            type: 'string',
+            optional: true,
+        },
+        {
+            name: 'setupSQL',
+            displayName: 'Setup SQL',
+            type: 'text',
+            optional: true,
+            description: 'SQL statements to run when the connection is established',
+        },
+    ],
+});
 //# sourceMappingURL=index.js.map

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@malloydata/db-bigquery",
-  "version": "0.0.336",
+  "version": "0.0.338",
   "license": "MIT",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -25,7 +25,7 @@
     "@google-cloud/bigquery": "^7.3.0",
     "@google-cloud/common": "^5.0.1",
     "@google-cloud/paginator": "^5.0.0",
-    "@malloydata/malloy": "0.0.336",
+    "@malloydata/malloy": "0.0.338",
     "gaxios": "^4.2.0"
   }
 }