feat(bigquery): support authorized views with dataset restrictions

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

@@ -39,6 +39,7 @@ layer of security by controlling which datasets can be accessed:
   It will reject the query if it attempts to access any table outside the
   allowed `datasets` list. To enforce this restriction, the following operations
   are also disallowed:
+
   - **Dataset-level operations** (e.g., `CREATE SCHEMA`, `ALTER SCHEMA`).
   - **Unanalyzable operations** where the accessed tables cannot be determined
     statically (e.g., `EXECUTE IMMEDIATE`, `CREATE PROCEDURE`, `CALL`).

internal/tools/bigquery/bigqueryanalyzecontribution/bigqueryanalyzecontribution.go

@@ -116,7 +116,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 		return nil, util.NewAgentError(fmt.Sprintf("unable to cast input_data parameter %s", paramsMap["input_data"]), nil)
 	}
 
-	bqClient, _, err := source.RetrieveClientAndService(accessToken)
+	bqClient, restService, err := source.RetrieveClientAndService(accessToken)
 	if err != nil {
 		return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
 	}
@@ -176,6 +176,26 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 	var inputDataSource string
 	trimmedUpperInputData := strings.TrimSpace(strings.ToUpper(inputData))
 	if strings.HasPrefix(trimmedUpperInputData, "SELECT") || strings.HasPrefix(trimmedUpperInputData, "WITH") {
+		if len(source.BigQueryAllowedDatasets()) > 0 {
+			var connProps []*bigqueryapi.ConnectionProperty
+			session, err := source.BigQuerySession()(ctx)
+			if err != nil {
+				return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
+			}
+			if session != nil {
+				connProps = []*bigqueryapi.ConnectionProperty{
+					{Key: "session_id", Value: session.ID},
+				}
+			}
+
+			dryRunJob, validationErr := bqutil.ValidateQueryAgainstAllowedDatasets(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, inputData, nil, connProps, source, source.GetMaximumBytesBilled(), false)
+			if validationErr != nil {
+				return nil, validationErr
+			}
+			if dryRunJob.Statistics.Query.StatementType != "SELECT" {
+				return nil, util.NewAgentError(fmt.Sprintf("the 'input_data' parameter only supports a table ID or a SELECT query. The provided query has statement type '%s'", dryRunJob.Statistics.Query.StatementType), nil)
+			}
+		}
 		inputDataSource = fmt.Sprintf("(%s)", inputData)
 	} else {
 		if !bqutil.ValidTableID(inputData) {
@@ -225,27 +245,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 		// If not in protected mode, create a session for this invocation.
 		createModelQuery.CreateSession = true
 	}
+
 	if len(source.BigQueryAllowedDatasets()) > 0 {
-		createModelQuery.DryRun = true
-		dryRunJob, err := createModelQuery.Run(ctx)
-		if err != nil {
-			return nil, util.ProcessGcpError(err)
-		}
-		status := dryRunJob.LastStatus()
-		if status.Statistics != nil {
-			if qStats, ok := status.Statistics.Details.(*bigqueryapi.QueryStatistics); ok {
-				for _, tableRef := range qStats.ReferencedTables {
-					if !source.IsDatasetAllowed(tableRef.ProjectID, tableRef.DatasetID) {
-						return nil, util.NewAgentError(fmt.Sprintf("query accesses dataset '%s.%s', which is not in the allowed list", tableRef.ProjectID, tableRef.DatasetID), nil)
-					}
-				}
-			} else {
-				return nil, util.NewAgentError("could not get query statistics details during dry run validation", nil)
-			}
-		} else {
-			return nil, util.NewAgentError("could not dry run model creation query to validate allowed datasets", nil)
+		_, validationErr := bqutil.ValidateQueryAgainstAllowedDatasets(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, createModelSQL, nil, createModelQuery.ConnectionProperties, source, source.GetMaximumBytesBilled(), true)
+		if validationErr != nil {
+			return nil, validationErr
 		}
-		createModelQuery.DryRun = false
 	}
 
 	createModelJob, err := createModelQuery.Run(ctx)

internal/tools/bigquery/bigqueryanalyzecontribution/bigqueryanalyzecontribution_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/googleapis/mcp-toolbox/internal/tools/bigquery/bigqueryanalyzecontribution"
 	"github.com/googleapis/mcp-toolbox/internal/tools/bigquery/bigquerycommon"
 	"github.com/googleapis/mcp-toolbox/internal/util/parameters"
+	bigqueryrestapi "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/option"
 )
 
@@ -319,9 +320,15 @@ func TestInvokeAllowedDatasetsValidation(t *testing.T) {
 		t.Fatalf("failed to create mocked BigQuery client: %v", err)
 	}
 
+	restService, err := bigqueryrestapi.NewService(ctx, option.WithEndpoint(mockServer.URL), option.WithoutAuthentication())
+	if err != nil {
+		t.Fatalf("failed to create mocked BigQuery REST service: %v", err)
+	}
+
 	// 3. Define mock source that returns this client and allowed datasets configuration
 	testSrc := &bigquerycommon.MockSource{
 		Client:          bqClient,
+		Service:         restService,
 		AllowedDatasets: []string{"allowed_dataset"},
 	}
 
@@ -348,7 +355,7 @@ func TestInvokeAllowedDatasetsValidation(t *testing.T) {
 
 	// 4. Set up parameters
 	data := map[string]any{
-		"input_data":          "allowed_dataset.my_table",
+		"input_data":          "SELECT * FROM unauthorized_dataset.some_table",
 		"contribution_metric": "SUM(metric)",
 		"is_test_col":         "is_test",
 		"dimension_id_cols":   []any{"dim1"},
@@ -370,7 +377,7 @@ func TestInvokeAllowedDatasetsValidation(t *testing.T) {
 		t.Fatal("expected Invoke to return an error due to out-of-allowlist dataset reference, but got nil")
 	}
 
-	expectedErr := "query accesses dataset 'test-project.unauthorized_dataset', which is not in the allowed list"
+	expectedErr := "access to dataset 'test-project.unauthorized_dataset' is not allowed"
 	if !strings.Contains(err.Error(), expectedErr) {
 		t.Errorf("expected error to contain %q, got: %v", expectedErr, err)
 	}