Merge branch 'main' into fix/bigquery-execute-sql-mcp-errors

Author: duwenxin99

cmd/internal/config_test.go

@@ -1803,7 +1803,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"discovery": tools.ToolsetConfig{
 					Name:      "discovery",
-					ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context"},
+					ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans"},
 				},
 			},
 		},

cmd/internal/imports.go

@@ -95,6 +95,7 @@ import (
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlookupcontext"
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlookupentry"
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
+	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexsearchdqscans"
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexsearchentries"
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataproc/dataprocgetcluster"
 	_ "github.com/googleapis/mcp-toolbox/internal/tools/dataproc/dataprocgetjob"

docs/en/integrations/bigquery/source.md

@@ -113,6 +113,7 @@ project: "my-project-id"
 #   - "https://www.googleapis.com/auth/bigquery"
 #   - "https://www.googleapis.com/auth/drive.readonly"
 # maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
+# maximumBytesBilled: 10737418240 # Optional: Per-query bytes scanned cap (in bytes).
 ```
 
 Initialize a BigQuery source that uses the client's access token:
@@ -133,6 +134,7 @@ useClientOAuth: true
 #   - "https://www.googleapis.com/auth/bigquery"
 #   - "https://www.googleapis.com/auth/drive.readonly"
 # maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
+# maximumBytesBilled: 10737418240 # Optional: Per-query bytes scanned cap (in bytes).
 ```
 
 ## Reference
@@ -148,3 +150,4 @@ useClientOAuth: true
 | scopes                    | []string |    false     | A list of OAuth 2.0 scopes to use for the credentials. If not provided, default scopes are used.                                                                                                                                                                                                                                                                                                                                                                                                                     |
 | impersonateServiceAccount |  string  |    false     | Service account email to impersonate when making BigQuery and Dataplex API calls. The authenticated principal must have the `roles/iam.serviceAccountTokenCreator` role on the target service account. [Learn More](https://cloud.google.com/iam/docs/service-account-impersonation)                                                                                                                                                                                                                                |
 | maxQueryResultRows             |   int    |    false     | The maximum number of rows to return from a query. Defaults to 50. |
+| maximumBytesBilled             |  int64   |    false     | The maximum bytes billed per query. When set, queries that exceed this limit fail before executing. |

docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-search-dq-scans.md

@@ -0,0 +1,64 @@
+---
+title: "dataplex-search-dq-scans"
+type: docs
+weight: 1
+description: >
+  A "dataplex-search-dq-scans" tool allows to search for data quality scans based on the provided parameters.
+aliases:
+- /resources/tools/dataplex-search-dq-scans
+---
+
+## About
+
+A `dataplex-search-dq-scans` tool returns data quality scans that match the given criteria.
+It's compatible with the following sources:
+
+- [dataplex](../../sources/dataplex.md)
+
+`dataplex-search-dq-scans` accepts the following optional parameters:
+
+- `filter` - Filter string to search/filter data quality scans. E.g. "display_name = \"my-scan\"".
+- `data_scan_id` - The resource name of the data scan to filter by: projects/{project}/locations/{locationId}/dataScans/{dataScanId}.
+- `table_name` - The name of the table to filter by. Maps to data.entity in the filter string. E.g. "//bigquery.googleapis.com/projects/P/datasets/D/tables/T".
+- `pageSize` - Number of returned data quality scans in the page. Defaults to `10`.
+- `orderBy` - Specifies the ordering of results.
+
+## Requirements
+
+### IAM Permissions
+
+Dataplex uses [Identity and Access Management (IAM)][iam-overview] to control
+user and group access to Dataplex resources. Toolbox will use your
+[Application Default Credentials (ADC)][adc] to authorize and authenticate when
+interacting with [Dataplex][dataplex-docs].
+
+In addition to [setting the ADC for your server][set-adc], you need to ensure
+the IAM identity has been given the correct IAM permissions for the tasks you
+intend to perform. See [Dataplex Universal Catalog IAM permissions][iam-permissions]
+and [Dataplex Universal Catalog IAM roles][iam-roles] for more information on
+applying IAM permissions and roles to an identity.
+
+[iam-overview]: https://cloud.google.com/dataplex/docs/iam-and-access-control
+[adc]: https://cloud.google.com/docs/authentication#adc
+[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc
+[iam-permissions]: https://cloud.google.com/dataplex/docs/iam-permissions
+[iam-roles]: https://cloud.google.com/dataplex/docs/iam-roles
+[dataplex-docs]: https://cloud.google.com/dataplex
+
+## Example
+
+```yaml
+kind: tools
+name: dataplex-search-dq-scans
+type: dataplex-search-dq-scans
+source: my-dataplex-source
+description: Use this tool to search for data quality scans.
+```
+
+## Reference
+
+| **field**   | **type** | **required** | **description**                                    |
+|-------------|:--------:|:------------:|----------------------------------------------------|
+| type        |  string  |     true     | Must be "dataplex-search-dq-scans".                |
+| source      |  string  |     true     | Name of the source the tool should execute on.     |
+| description |  string  |     true     | Description of the tool that is passed to the LLM. |

internal/prebuiltconfigs/tools/dataplex.yaml

@@ -34,10 +34,15 @@ tools:
     kind: dataplex-lookup-context
     source: dataplex-source
     description: Retrieves rich metadata regarding one or more data assets along with their relationships.
+  search_dq_scans:
+    kind: dataplex-search-dq-scans
+    source: dataplex-source
+    description: Use this tool to search for data quality scans in Dataplex.
 
 toolsets:
   discovery:
     - search_entries
     - lookup_entry
     - search_aspect_types
     - lookup_context
+    - search_dq_scans

internal/sources/bigquery/bigquery.go

@@ -90,6 +90,7 @@ type Config struct {
 	ImpersonateServiceAccount string              `yaml:"impersonateServiceAccount"`
 	Scopes                    StringOrStringSlice `yaml:"scopes"`
 	MaxQueryResultRows        int                 `yaml:"maxQueryResultRows"`
+	MaximumBytesBilled        int64               `yaml:"maximumBytesBilled" validate:"gte=0"`
 }
 
 // StringOrStringSlice is a custom type that can unmarshal both a single string
@@ -157,6 +158,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		RestService:         restService,
 		TokenSource:         tokenSource,
 		MaxQueryResultRows:  r.MaxQueryResultRows,
+		MaximumBytesBilled:  r.MaximumBytesBilled,
 		ClientCreator:       clientCreator,
 		AuthTokenHeaderName: "Authorization",
 	}
@@ -291,6 +293,7 @@ type Source struct {
 	TokenSource               oauth2.TokenSource
 	AuthTokenHeaderName       string
 	MaxQueryResultRows        int
+	MaximumBytesBilled        int64
 	ClientCreator             BigqueryClientCreator
 	AllowedDatasets           map[string]struct{}
 	sessionMutex              sync.Mutex
@@ -472,6 +475,10 @@ func (s *Source) GetMaxQueryResultRows() int {
 	return s.MaxQueryResultRows
 }
 
+func (s *Source) GetMaximumBytesBilled() int64 {
+	return s.MaximumBytesBilled
+}
+
 func (s *Source) BigQueryClientCreator() BigqueryClientCreator {
 	return s.ClientCreator
 }
@@ -572,6 +579,9 @@ func (s *Source) RunSQL(ctx context.Context, bqClient *bigqueryapi.Client, state
 	if connProps != nil {
 		query.ConnectionProperties = connProps
 	}
+	if s.MaximumBytesBilled > 0 {
+		query.MaxBytesBilled = s.MaximumBytesBilled
+	}
 
 	// This block handles SELECT statements, which return a row set.
 	// We iterate through the results, convert each row into a map of

internal/sources/bigquery/bigquery_test.go

@@ -237,6 +237,26 @@ func TestParseFromYamlBigQuery(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "with maximum bytes billed example",
+			in: `
+			kind: source
+			name: my-instance
+			type: bigquery
+			project: my-project
+			location: us
+			maximumBytesBilled: 10737418240
+			`,
+			want: map[string]sources.SourceConfig{
+				"my-instance": bigquery.Config{
+					Name:               "my-instance",
+					Type:               bigquery.SourceType,
+					Project:            "my-project",
+					Location:           "us",
+					MaximumBytesBilled: 10737418240,
+				},
+			},
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -279,6 +299,17 @@ func TestFailParseFromYaml(t *testing.T) {
 			`,
 			err: "error unmarshaling source: unable to parse source \"my-instance\" as \"bigquery\": Key: 'Config.Project' Error:Field validation for 'Project' failed on the 'required' tag",
 		},
+		{
+			desc: "negative maximum bytes billed",
+			in: `
+			kind: source
+			name: my-instance
+			type: bigquery
+			project: my-project
+			maximumBytesBilled: -1
+			`,
+			err: "error unmarshaling source: unable to parse source \"my-instance\" as \"bigquery\": [1:21] Key: 'Config.MaximumBytesBilled' Error:Field validation for 'MaximumBytesBilled' failed on the 'gte' tag\n>  1 | maximumBytesBilled: -1\n                           ^\n   2 | name: my-instance\n   3 | project: my-project\n   4 | type: bigquery",
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -347,6 +378,59 @@ func TestInitialize_MaxQueryResultRows(t *testing.T) {
 	}
 }
 
+func TestInitialize_MaximumBytesBilled(t *testing.T) {
+	ctx, err := testutils.ContextWithNewLogger()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	ctx = util.WithUserAgent(ctx, "test-agent")
+	tracer := noop.NewTracerProvider().Tracer("")
+
+	tcs := []struct {
+		desc string
+		cfg  bigquery.Config
+		want int64
+	}{
+		{
+			desc: "default value",
+			cfg: bigquery.Config{
+				Name:           "test-default",
+				Type:           bigquery.SourceType,
+				Project:        "test-project",
+				UseClientOAuth: "true",
+			},
+			want: 0,
+		},
+		{
+			desc: "configured value",
+			cfg: bigquery.Config{
+				Name:               "test-configured",
+				Type:               bigquery.SourceType,
+				Project:            "test-project",
+				UseClientOAuth:     "true",
+				MaximumBytesBilled: 10737418240,
+			},
+			want: 10737418240,
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			src, err := tc.cfg.Initialize(ctx, tracer)
+			if err != nil {
+				t.Fatalf("Initialize failed: %v", err)
+			}
+			bqSrc, ok := src.(*bigquery.Source)
+			if !ok {
+				t.Fatalf("Expected *bigquery.Source, got %T", src)
+			}
+			if bqSrc.MaximumBytesBilled != tc.want {
+				t.Errorf("MaximumBytesBilled = %d, want %d", bqSrc.MaximumBytesBilled, tc.want)
+			}
+		})
+	}
+}
+
 func TestNormalizeValue(t *testing.T) {
 	tests := []struct {
 		name     string

internal/sources/dataplex/dataplex.go

@@ -64,13 +64,14 @@ func (r Config) SourceConfigType() string {
 
 func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
 	// Initializes a Dataplex source
-	client, err := initDataplexConnection(ctx, tracer, r.Name, r.Project)
+	client, dataScanClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project)
 	if err != nil {
 		return nil, err
 	}
 	s := &Source{
-		Config: r,
-		Client: client,
+		Config:         r,
+		Client:         client,
+		DataScanClient: dataScanClient,
 	}
 
 	return s, nil
@@ -80,7 +81,8 @@ var _ sources.Source = &Source{}
 
 type Source struct {
 	Config
-	Client *dataplexapi.CatalogClient
+	Client         *dataplexapi.CatalogClient
+	DataScanClient *dataplexapi.DataScanClient
 }
 
 func (s *Source) SourceType() string {
@@ -100,30 +102,39 @@ func (s *Source) CatalogClient() *dataplexapi.CatalogClient {
 	return s.Client
 }
 
+func (s *Source) GetDataScanClient() *dataplexapi.DataScanClient {
+	return s.DataScanClient
+}
+
 func initDataplexConnection(
 	ctx context.Context,
 	tracer trace.Tracer,
 	name string,
 	project string,
-) (*dataplexapi.CatalogClient, error) {
+) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, error) {
 	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
 	defer span.End()
 
 	cred, err := google.FindDefaultCredentials(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err)
+		return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials for project %q: %w", project, err)
 	}
 
 	userAgent, err := util.UserAgentFromContext(ctx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	client, err := dataplexapi.NewCatalogClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred))
 	if err != nil {
-		return nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
+		return nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
+	}
+
+	dataScanClient, err := dataplexapi.NewDataScanClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred))
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create Dataplex DataScan client for project %q: %w", project, err)
 	}
-	return client, nil
+	return client, dataScanClient, nil
 }
 
 func (s *Source) LookupEntry(ctx context.Context, name string, view int, aspectTypes []string, entry string) (*dataplexpb.Entry, error) {
@@ -256,3 +267,29 @@ func (s *Source) LookupContext(ctx context.Context, name string, resources []str
 	}
 	return result, nil
 }
+
+func (s *Source) SearchDataQualityScans(ctx context.Context, filter string, pageSize int, orderBy string) ([]*dataplexpb.DataScan, error) {
+	req := &dataplexpb.ListDataScansRequest{
+		Parent:   fmt.Sprintf("projects/%s/locations/-", s.ProjectID()),
+		Filter:   filter,
+		PageSize: int32(pageSize),
+		OrderBy:  orderBy,
+	}
+
+	it := s.GetDataScanClient().ListDataScans(ctx, req)
+	var results []*dataplexpb.DataScan
+	for {
+		scan, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			if st, ok := grpcstatus.FromError(err); ok {
+				return nil, fmt.Errorf("failed to list data scans: code=%s message=%s", st.Code(), st.Message())
+			}
+			return nil, fmt.Errorf("failed to list data scans: %w", err)
+		}
+		results = append(results, scan)
+	}
+	return results, nil
+}

internal/tools/bigquery/bigqueryanalyzecontribution/bigqueryanalyzecontribution.go

@@ -53,6 +53,7 @@ type compatibleSource interface {
 	BigQueryClient() *bigqueryapi.Client
 	UseClientAuthorization() bool
 	GetAuthTokenHeaderName() string
+	GetMaximumBytesBilled() int64
 	IsDatasetAllowed(projectID, datasetID string) bool
 	BigQueryAllowedDatasets() []string
 	BigQuerySession() bigqueryds.BigQuerySessionProvider
@@ -219,7 +220,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 					{Key: "session_id", Value: session.ID},
 				}
 			}
-			dryRunJob, err := bqutil.DryRunQuery(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, inputData, nil, connProps)
+			dryRunJob, err := bqutil.DryRunQuery(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, inputData, nil, connProps, source.GetMaximumBytesBilled())
 			if err != nil {
 				return nil, util.ProcessGcpError(err)
 			}