Making changes to required files to add in maxBillableBytes.

Author: duncan771

pkg/bigquery/driver/conn.go

@@ -124,6 +124,11 @@ func (c *Conn) execContext(ctx context.Context, query string, args []driver.Valu
 	q := c.client.Query(query)
 
 	q.QueryConfig.Labels = c.headersAsLabels()
+
+	if c.cfg.MaxBillableBytes > 0 {
+		q.QueryConfig.MaxBytesBilled = c.cfg.MaxBillableBytes
+	}
+
 	// q.DefaultProjectID = c.cfg.Project // allows omitting project in table reference
 	// q.DefaultDatasetID = c.cfg.Dataset // allows omitting dataset in table reference
 
@@ -228,6 +233,10 @@ func (c *Conn) queryContext(ctx context.Context, query string, args []driver.Val
 
 	q.QueryConfig.Labels = c.headersAsLabels()
 
+	if c.cfg.MaxBillableBytes > 0 {
+		q.QueryConfig.MaxBytesBilled = c.cfg.MaxBillableBytes
+	}
+
 	rowsIterator, err := q.Read(ctx)
 	if err != nil {
 		return nil, err

pkg/bigquery/settings.go

@@ -43,6 +43,7 @@ func getConnectionSettings(settings types.BigQuerySettings, queryArgs *Connectio
 		Project:            settings.DefaultProject,
 		Location:           settings.ProcessingLocation,
 		AuthenticationType: settings.AuthenticationType,
+		MaxBillableBytes:   settings.MaxBillableBytes,
 	}
 
 	// We want to set the location to empty string only if query args are set

pkg/bigquery/types/types.go

@@ -14,6 +14,7 @@ type BigQuerySettings struct {
 	TokenUri           string `json:"tokenUri"`
 	QueryPriority      string `json:"queryPriority"`
 	ProcessingLocation string `json:"processingLocation"`
+	MaxBillableBytes   int64  `json:"maxBillableBytes,omitempty"`
 	Updated            time.Time
 	AuthenticationType string `json:"authenticationType"`
 	PrivateKeyPath     string `json:"privateKeyPath"`
@@ -28,6 +29,7 @@ type ConnectionSettings struct {
 	Project            string
 	Dataset            string
 	Headers            map[string][]string
+	MaxBillableBytes   int64
 }
 type TableFieldSchema struct {
 	Name        string       `json:"name"`

src/components/ConfigEditor.tsx

@@ -1,7 +1,7 @@
 import { DataSourcePluginOptionsEditorProps, onUpdateDatasourceJsonDataOptionSelect } from '@grafana/data';
 import { AuthConfig, GOOGLE_AUTH_TYPE_OPTIONS } from '@grafana/google-sdk';
 import { config } from '@grafana/runtime';
-import { Field, SecureSocksProxySettings, Select } from '@grafana/ui';
+import { Field, Input, SecureSocksProxySettings, Select } from '@grafana/ui';
 import React from 'react';
 import { PROCESSING_LOCATIONS } from '../constants';
 import { BigQueryOptions, BigQuerySecureJsonData } from '../types';
@@ -15,6 +15,16 @@ export const BigQueryConfigEditor: React.FC<BigQueryConfigEditorProps> = (props)
   const { options, onOptionsChange } = props;
   const { jsonData } = options;
 
+  const onMaxBillableBytesChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    onOptionsChange({
+      ...options,
+      jsonData: {
+        ...jsonData,
+        maxBillableBytes: Number(event.target.value),
+      },
+    });
+  };
+
   return (
     <>
       <DataSourceDescription
@@ -59,6 +69,30 @@ export const BigQueryConfigEditor: React.FC<BigQueryConfigEditorProps> = (props)
             menuShouldPortal={true}
           />
         </Field>
+        <Field
+          label="Max billable bytes"
+          description={
+            <span>
+              Prevent queries that would process more than this amount of bytes. Read more about max billable bytes{' '}
+              <a
+                href="https://cloud.google.com/bigquery/docs/best-practices-costs"
+                rel="noreferrer"
+                className="external-link"
+                target="_blank"
+              >
+                here
+              </a>
+            </span>
+          }
+        >
+          <Input
+            className="width-30"
+            placeholder="Optional, example 5242880"
+            type={'number'}
+            value={jsonData.maxBillableBytes || ''}
+            onChange={onMaxBillableBytesChange}
+          />
+        </Field>
 
         {config.secureSocksDSProxyEnabled && (
           <SecureSocksProxySettings options={options} onOptionsChange={onOptionsChange} />

src/types.ts

@@ -27,6 +27,7 @@ export interface BigQueryOptions extends DataSourceOptions {
   processingLocation?: string;
   queryPriority?: QueryPriority;
   enableSecureSocksProxy?: boolean;
+  maxBillableBytes?: number;
 }
 
 export interface BigQuerySecureJsonData extends DataSourceSecureJsonData {}