feat(source/bigquery): add optional write mode config (googleapis#1157)

Summary
Adds an optional write_mode configuration to the BigQuery source,
enhancing security by controlling the types of SQL statements that can
be executed to prevent unauthorized data modification.

Key Changes
Added writeMode Configuration: A new write_mode field is added to the
BigQuery source, supporting three modes:

allowed (Default): Permits all SQL statements.

blocked: Allows only SELECT queries.

protected: Enables session-based execution, restricting write operations
(like CREATE TABLE) to the session's temporary dataset, thus protecting
permanent datasets. Note: at the moment, this won't work with
useClientOAuth, will fix this in the future.

These restrictions primarily apply to the bigquery-execute-sql tool and
the session may be used in other tools.

Author: Genesis929

docs/en/resources/sources/bigquery.md

@@ -119,6 +119,7 @@ sources:
     kind: "bigquery"
     project: "my-project-id"
     # location: "US" # Optional: Specifies the location for query jobs.
+    # writeMode: "allowed" # One of: allowed, blocked, protected. Defaults to "allowed".
     # allowedDatasets: # Optional: Restricts tool access to a specific list of datasets.
     #   - "my_dataset_1"
     #   - "other_project.my_dataset_2"
@@ -133,6 +134,7 @@ sources:
     project: "my-project-id"
     useClientOAuth: true
     # location: "US" # Optional: Specifies the location for query jobs.
+    # writeMode: "allowed" # One of: allowed, blocked, protected. Defaults to "allowed".
     # allowedDatasets: # Optional: Restricts tool access to a specific list of datasets.
     #   - "my_dataset_1"
     #   - "other_project.my_dataset_2"
@@ -145,5 +147,6 @@ sources:
 | kind            |  string  |     true     | Must be "bigquery".                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
 | project         |  string  |     true     | Id of the Google Cloud project to use for billing and as the default project for BigQuery resources.                                                                                                                                                                                                                                                                                                                                                                                                                |
 | location        |  string  |    false     | Specifies the location (e.g., 'us', 'asia-northeast1') in which to run the query job. This location must match the location of any tables referenced in the query. Defaults to the table's location or 'US' if the location cannot be determined. [Learn More](https://cloud.google.com/bigquery/docs/locations)                                                                                                                                                                                                    |
+| writeMode       |  string  |    false     | Controls the write behavior for tools. `allowed` (default): All queries are permitted. `blocked`: Only `SELECT` statements are allowed for the `bigquery-execute-sql` tool. `protected`: Enables session-based execution where all tools associated with this source instance share the same [BigQuery session](https://cloud.google.com/bigquery/docs/sessions-intro). This allows for stateful operations using temporary tables (e.g., `CREATE TEMP TABLE`). For `bigquery-execute-sql`, `SELECT` statements can be used on all tables, but write operations are restricted to the session's temporary dataset. For tools like `bigquery-sql`, `bigquery-forecast`, and `bigquery-analyze-contribution`, the `writeMode` restrictions do not apply, but they will operate within the shared session. **Note:** The `protected` mode cannot be used with `useClientOAuth: true`. It is also not recommended for multi-user server environments, as all users would share the same session. A session is terminated automatically after 24 hours of inactivity or after 7 days, whichever comes first. A new session is created on the next request, and any temporary data from the previous session will be lost. |
 | allowedDatasets | []string |    false     | An optional list of dataset IDs that tools using this source are allowed to access. If provided, any tool operation attempting to access a dataset not in this list will be rejected. To enforce this, two types of operations are also disallowed: 1) Dataset-level operations (e.g., `CREATE SCHEMA`), and 2) operations where table access cannot be statically analyzed (e.g., `EXECUTE IMMEDIATE`, `CREATE PROCEDURE`). If a single dataset is provided, it will be treated as the default for prebuilt tools. |
-| useClientOAuth  |   bool   |    false     | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries.                                                                                                                                                                                                                                                                                                                                                                                                            |
+| useClientOAuth  |   bool   |    false     | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. **Note:** This cannot be used with `writeMode: protected`.                                                                                                                                                                                                                                                                                                                                                |

docs/en/resources/tools/bigquery/bigquery-analyze-contribution.md

@@ -39,6 +39,13 @@ It's compatible with the following sources:
   insights. Can be `'NO_PRUNING'` or `'PRUNE_REDUNDANT_INSIGHTS'`. Defaults to
   `'PRUNE_REDUNDANT_INSIGHTS'`.
 
+The behavior of this tool is influenced by the `writeMode` setting on its `bigquery` source:
+
+- **`allowed` (default) and `blocked`:** These modes do not impose any special restrictions on the `bigquery-analyze-contribution` tool.
+- **`protected`:** This mode enables session-based execution. The tool will operate within the same BigQuery session as other
+  tools using the same source. This allows the `input_data` parameter to be a query that references temporary resources (e.g., 
+  `TEMP` tables) created within that session.
+
 
 ## Example
 

docs/en/resources/tools/bigquery/bigquery-execute-sql.md

@@ -20,8 +20,15 @@ It's compatible with the following sources:
 - **`dry_run`** (optional): If set to `true`, the query is validated but not run,
   returning information about the execution instead. Defaults to `false`.
 
+The behavior of this tool is influenced by the `writeMode` setting on its `bigquery` source:
+
+- **`allowed` (default):** All SQL statements are permitted.
+- **`blocked`:** Only `SELECT` statements are allowed. Any other type of statement (e.g., `INSERT`, `UPDATE`, `CREATE`) will be rejected.
+- **`protected`:** This mode enables session-based execution. `SELECT` statements can be used on all tables, while write operations are allowed only for the session's temporary dataset (e.g., `CREATE TEMP TABLE ...`). This prevents modifications to permanent datasets while allowing stateful, multi-step operations within a secure session.
+
 The tool's behavior is influenced by the `allowedDatasets` restriction on the
-`bigquery` source:
+`bigquery` source. Similar to `writeMode`, this setting provides an additional layer of security by controlling which datasets can be accessed:
+
 - **Without `allowedDatasets` restriction:** The tool can execute any valid GoogleSQL
   query.
 - **With `allowedDatasets` restriction:** Before execution, the tool performs a dry run
@@ -33,6 +40,8 @@ The tool's behavior is influenced by the `allowedDatasets` restriction on the
   - **Unanalyzable operations** where the accessed tables cannot be determined
     statically (e.g., `EXECUTE IMMEDIATE`, `CREATE PROCEDURE`, `CALL`).
 
+> **Note:** This tool is intended for developer assistant workflows with human-in-the-loop and shouldn't be used for production agents.
+
 ## Example
 
 ```yaml

docs/en/resources/tools/bigquery/bigquery-forecast.md

@@ -33,12 +33,19 @@ query based on the provided parameters:
 - **horizon** (integer, optional): The number of future time steps you want to
   predict. It defaults to 10 if not specified.
 
-The tool's behavior regarding these parameters is influenced by the `allowedDatasets` restriction on the `bigquery` source:
+The behavior of this tool is influenced by the `writeMode` setting on its `bigquery` source:
+
+- **`allowed` (default) and `blocked`:** These modes do not impose any special restrictions on the `bigquery-forecast` tool.
+- **`protected`:** This mode enables session-based execution. The tool will operate within the same BigQuery session as other
+  tools using the same source. This allows the `history_data` parameter to be a query that references temporary resources (e.g., 
+  `TEMP` tables) created within that session.
+
+The tool's behavior is also influenced by the `allowedDatasets` restriction on the `bigquery` source:
+
 - **Without `allowedDatasets` restriction:** The tool can use any table or query for the `history_data` parameter.
-- **With `allowedDatasets` restriction:** The tool verifies that the `history_data` parameter only accesses tables 
-within the allowed datasets. If `history_data` is a table ID, the tool checks if the table's dataset is in the 
-allowed list. If `history_data` is a query, the tool performs a dry run to analyze the query and rejects it 
-if it accesses any table outside the allowed list.
+- **With `allowedDatasets` restriction:** The tool verifies that the `history_data` parameter only accesses tables within the allowed datasets.
+  - If `history_data` is a table ID, the tool checks if the table's dataset is in the allowed list.
+  - If `history_data` is a query, the tool performs a dry run to analyze the query and rejects it if it accesses any table outside the allowed list.
 
 ## Example
 

docs/en/resources/tools/bigquery/bigquery-sql.md

@@ -15,6 +15,11 @@ the following sources:
 
 - [bigquery](../../sources/bigquery.md)
 
+The behavior of this tool is influenced by the `writeMode` setting on its `bigquery` source:
+
+- **`allowed` (default) and `blocked`:** These modes do not impose any restrictions on the `bigquery-sql` tool. The pre-defined SQL statement will be executed as-is.
+- **`protected`:** This mode enables session-based execution. The tool will operate within the same BigQuery session as other tools using the same source, allowing it to interact with temporary resources like `TEMP` tables created within that session.
+
 ### GoogleSQL
 
 BigQuery uses [GoogleSQL][bigquery-googlesql] for querying data. The integration

internal/sources/bigquery/bigquery.go

@@ -20,6 +20,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	bigqueryapi "cloud.google.com/go/bigquery"
 	dataplexapi "cloud.google.com/go/dataplex/apiv1"
@@ -36,11 +37,22 @@ import (
 
 const SourceKind string = "bigquery"
 
+const (
+	// No write operations are allowed.
+	WriteModeBlocked string = "blocked"
+	// Only protected write operations are allowed in a BigQuery session.
+	WriteModeProtected string = "protected"
+	// All write operations are allowed.
+	WriteModeAllowed string = "allowed"
+)
+
 // validate interface
 var _ sources.SourceConfig = Config{}
 
 type BigqueryClientCreator func(tokenString string, wantRestService bool) (*bigqueryapi.Client, *bigqueryrestapi.Service, error)
 
+type BigQuerySessionProvider func(ctx context.Context) (*Session, error)
+
 type DataplexClientCreator func(tokenString string) (*dataplexapi.CatalogClient, error)
 
 func init() {
@@ -63,6 +75,7 @@ type Config struct {
 	Kind            string   `yaml:"kind" validate:"required"`
 	Project         string   `yaml:"project" validate:"required"`
 	Location        string   `yaml:"location"`
+	WriteMode       string   `yaml:"writeMode"`
 	AllowedDatasets []string `yaml:"allowedDatasets"`
 	UseClientOAuth  bool     `yaml:"useClientOAuth"`
 }
@@ -73,6 +86,14 @@ func (r Config) SourceConfigKind() string {
 }
 
 func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
+	if r.WriteMode == "" {
+		r.WriteMode = WriteModeAllowed
+	}
+
+	if r.WriteMode == WriteModeProtected && r.UseClientOAuth {
+		return nil, fmt.Errorf("writeMode 'protected' cannot be used with useClientOAuth 'true'")
+	}
+
 	var client *bigqueryapi.Client
 	var restService *bigqueryrestapi.Service
 	var tokenSource oauth2.TokenSource
@@ -133,9 +154,15 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		TokenSource:        tokenSource,
 		MaxQueryResultRows: 50,
 		ClientCreator:      clientCreator,
+		WriteMode:          r.WriteMode,
 		AllowedDatasets:    allowedDatasets,
 		UseClientOAuth:     r.UseClientOAuth,
 	}
+	s.SessionProvider = s.newBigQuerySessionProvider()
+
+	if r.WriteMode != WriteModeAllowed && r.WriteMode != WriteModeBlocked && r.WriteMode != WriteModeProtected {
+		return nil, fmt.Errorf("invalid writeMode %q: must be one of %q, %q, or %q", r.WriteMode, WriteModeAllowed, WriteModeProtected, WriteModeBlocked)
+	}
 	s.makeDataplexCatalogClient = s.lazyInitDataplexClient(ctx, tracer)
 	return s, nil
 
@@ -156,7 +183,19 @@ type Source struct {
 	ClientCreator             BigqueryClientCreator
 	AllowedDatasets           map[string]struct{}
 	UseClientOAuth            bool
+	WriteMode                 string
+	sessionMutex              sync.Mutex
 	makeDataplexCatalogClient func() (*dataplexapi.CatalogClient, DataplexClientCreator, error)
+	SessionProvider           BigQuerySessionProvider
+	Session                   *Session
+}
+
+type Session struct {
+	ID           string
+	ProjectID    string
+	DatasetID    string
+	CreationTime time.Time
+	LastUsed     time.Time
 }
 
 func (s *Source) SourceKind() string {
@@ -172,6 +211,103 @@ func (s *Source) BigQueryRestService() *bigqueryrestapi.Service {
 	return s.RestService
 }
 
+func (s *Source) BigQueryWriteMode() string {
+	return s.WriteMode
+}
+
+func (s *Source) BigQuerySession() BigQuerySessionProvider {
+	return s.SessionProvider
+}
+
+func (s *Source) newBigQuerySessionProvider() BigQuerySessionProvider {
+	return func(ctx context.Context) (*Session, error) {
+		if s.WriteMode != WriteModeProtected {
+			return nil, nil
+		}
+
+		s.sessionMutex.Lock()
+		defer s.sessionMutex.Unlock()
+
+		logger, err := util.LoggerFromContext(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get logger from context: %w", err)
+		}
+
+		if s.Session != nil {
+			// Absolute 7-day lifetime check.
+			const sessionMaxLifetime = 7 * 24 * time.Hour
+			// This assumes a single task will not exceed 30 minutes, preventing it from failing mid-execution.
+			const refreshThreshold = 30 * time.Minute
+			if time.Since(s.Session.CreationTime) > (sessionMaxLifetime - refreshThreshold) {
+				logger.DebugContext(ctx, "Session is approaching its 7-day maximum lifetime. Creating a new one.")
+			} else {
+				job := &bigqueryrestapi.Job{
+					Configuration: &bigqueryrestapi.JobConfiguration{
+						DryRun: true,
+						Query: &bigqueryrestapi.JobConfigurationQuery{
+							Query:                "SELECT 1",
+							UseLegacySql:         new(bool),
+							ConnectionProperties: []*bigqueryrestapi.ConnectionProperty{{Key: "session_id", Value: s.Session.ID}},
+						},
+					},
+				}
+				_, err := s.RestService.Jobs.Insert(s.Project, job).Do()
+				if err == nil {
+					s.Session.LastUsed = time.Now()
+					return s.Session, nil
+				}
+				logger.DebugContext(ctx, "Session validation failed (likely expired), creating a new one.", "error", err)
+			}
+		}
+
+		// Create a new session if one doesn't exist, it has passed its 7-day lifetime,
+		// or it failed the validation dry run.
+
+		creationTime := time.Now()
+		job := &bigqueryrestapi.Job{
+			JobReference: &bigqueryrestapi.JobReference{
+				ProjectId: s.Project,
+				Location:  s.Location,
+			},
+			Configuration: &bigqueryrestapi.JobConfiguration{
+				DryRun: true,
+				Query: &bigqueryrestapi.JobConfigurationQuery{
+					Query:         "SELECT 1",
+					CreateSession: true,
+				},
+			},
+		}
+
+		createdJob, err := s.RestService.Jobs.Insert(s.Project, job).Do()
+		if err != nil {
+			return nil, fmt.Errorf("failed to create new session: %w", err)
+		}
+
+		var sessionID, sessionDatasetID, projectID string
+		if createdJob.Status != nil && createdJob.Statistics.SessionInfo != nil {
+			sessionID = createdJob.Statistics.SessionInfo.SessionId
+		} else {
+			return nil, fmt.Errorf("failed to get session ID from new session job")
+		}
+
+		if createdJob.Configuration != nil && createdJob.Configuration.Query != nil && createdJob.Configuration.Query.DestinationTable != nil {
+			sessionDatasetID = createdJob.Configuration.Query.DestinationTable.DatasetId
+			projectID = createdJob.Configuration.Query.DestinationTable.ProjectId
+		} else {
+			return nil, fmt.Errorf("failed to get session dataset ID from new session job")
+		}
+
+		s.Session = &Session{
+			ID:           sessionID,
+			ProjectID:    projectID,
+			DatasetID:    sessionDatasetID,
+			CreationTime: creationTime,
+			LastUsed:     creationTime,
+		}
+		return s.Session, nil
+	}
+}
+
 func (s *Source) UseClientAuthorization() bool {
 	return s.UseClientOAuth
 }
@@ -257,7 +393,7 @@ func initBigQueryConnection(
 	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
 	defer span.End()
 
-	cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope)
+	cred, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform")
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err)
 	}

internal/sources/bigquery/bigquery_test.go

@@ -37,14 +37,34 @@ func TestParseFromYamlBigQuery(t *testing.T) {
 				my-instance:
 					kind: bigquery
 					project: my-project
-					location: us
+			`,
+			want: server.SourceConfigs{
+				"my-instance": bigquery.Config{
+					Name:      "my-instance",
+					Kind:      bigquery.SourceKind,
+					Project:   "my-project",
+					Location:  "",
+					WriteMode: "",
+				},
+			},
+		},
+		{
+			desc: "all fields specified",
+			in: `
+			sources:
+				my-instance:
+					kind: bigquery
+					project: my-project
+					location: asia
+					writeMode: blocked
 			`,
 			want: server.SourceConfigs{
 				"my-instance": bigquery.Config{
 					Name:           "my-instance",
 					Kind:           bigquery.SourceKind,
 					Project:        "my-project",
-					Location:       "us",
+					Location:       "asia",
+					WriteMode:      "blocked",
 					UseClientOAuth: false,
 				},
 			},