Merge branch 'main' into remove-sse-access-control

Author: Yuan325

.ci/continuous.release.cloudbuild.yaml

@@ -47,7 +47,7 @@ steps:
         tar -xf zig.tar.xz -C /zig-tools --strip-components=1
 
   - id: "install-macos-sdk"
-    name: golang:1
+    name: "gcr.io/cloud-builders/gcloud:latest"
     waitFor: ['-']
     volumes:
       - name: 'macos-sdk'
@@ -57,8 +57,8 @@ steps:
         set -e
         apt-get update && apt-get install -y xz-utils
         echo "Downloading macOS 14.5 SDK..."
-        curl -fL -o sdk.tar.xz https://github.com/alexey-lysiuk/macos-sdk/releases/download/14.5/MacOSX14.5.tar.xz
-        
+        gcloud storage cp gs://${_ASSETS_BUCKET}/MacOSX14.5.tar.xz sdk.tar.xz
+
         mkdir -p /macos-sdk/MacOSX14.5.sdk
         echo "Unpacking macOS 14.5 SDK..."
         tar -xf sdk.tar.xz -C /macos-sdk/MacOSX14.5.sdk --strip-components=1
@@ -310,4 +310,5 @@ substitutions:
   _AR_REPO_NAME: toolbox-dev
   _OLD_BUCKET_NAME: genai-toolbox-dev
   _BUCKET_NAME: mcp-toolbox-for-databases-dev
+  _ASSETS_BUCKET: toolbox-build-assets
   _DOCKER_URI: ${_AR_HOSTNAME}/${PROJECT_ID}/${_AR_REPO_NAME}/toolbox

.ci/versioned.release.cloudbuild.yaml

@@ -63,7 +63,7 @@ steps:
         tar -xf zig.tar.xz -C /zig-tools --strip-components=1
 
   - id: "install-macos-sdk"
-    name: golang:1
+    name: "gcr.io/cloud-builders/gcloud:latest"
     waitFor: ['-']
     volumes:
       - name: 'macos-sdk'
@@ -73,8 +73,8 @@ steps:
         set -e
         apt-get update && apt-get install -y xz-utils
         echo "Downloading macOS 14.5 SDK..."
-        curl -fL -o sdk.tar.xz https://github.com/alexey-lysiuk/macos-sdk/releases/download/14.5/MacOSX14.5.tar.xz
-        
+        gcloud storage cp gs://${_ASSETS_BUCKET}/MacOSX14.5.tar.xz sdk.tar.xz
+
         mkdir -p /macos-sdk/MacOSX14.5.sdk
         echo "Unpacking macOS 14.5 SDK..."
         tar -xf sdk.tar.xz -C /macos-sdk/MacOSX14.5.sdk --strip-components=1
@@ -375,5 +375,6 @@ substitutions:
   _AR_REPO_NAME: toolbox
   _OLD_BUCKET_NAME: genai-toolbox
   _BUCKET_NAME: mcp-toolbox-for-databases
+  _ASSETS_BUCKET: toolbox-build-assets
   _DOCKER_URI: ${_AR_HOSTNAME}/${PROJECT_ID}/${_AR_REPO_NAME}/toolbox
   _PUSH_LATEST: "false" # Substituted in trigger

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