From fa8220d08aeeb49821f236ad09bd5568c425d8a9 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Wed, 8 Oct 2025 00:07:54 +0000 Subject: [PATCH 1/6] feat(tools/bigquery-get-dataset-info)!: add allowed dataset support --- .../bigquery/bigquery-get-dataset-info.md | 17 +++- .../bigquerygetdatasetinfo.go | 84 +++++++++++++++---- tests/bigquery/bigquery_integration_test.go | 57 +++++++++++++ 3 files changed, 137 insertions(+), 21 deletions(-) diff --git a/docs/en/resources/tools/bigquery/bigquery-get-dataset-info.md b/docs/en/resources/tools/bigquery/bigquery-get-dataset-info.md index 4f1f2cb716b1..59627c57cee0 100644 --- a/docs/en/resources/tools/bigquery/bigquery-get-dataset-info.md +++ b/docs/en/resources/tools/bigquery/bigquery-get-dataset-info.md @@ -15,10 +15,19 @@ It's compatible with the following sources: - [bigquery](../../sources/bigquery.md) -`bigquery-get-dataset-info` takes a `dataset` parameter to specify the dataset -on the given source. It also optionally accepts a `project` parameter to -define the Google Cloud project ID. If the `project` parameter is not provided, -the tool defaults to using the project defined in the source configuration. +`bigquery-get-dataset-info` accepts the following parameters: +- **`dataset`** (required): Specifies the dataset for which to retrieve metadata. +- **`project`** (optional): Defines the Google Cloud project ID. If not provided, + the tool defaults to the project from the source configuration. + +The tool's behavior regarding these parameters is influenced by the +`allowedDatasets` restriction on the `bigquery` source: +- **Without `allowedDatasets` restriction:** The tool can retrieve metadata for + any dataset specified by the `dataset` and `project` parameters. +- **With `allowedDatasets` restriction:** Before retrieving metadata, the tool + verifies that the requested dataset is in the allowed list. If it is not, the + request is denied. If only one dataset is specified in the `allowedDatasets` + list, it will be used as the default value for the `dataset` parameter. ## Example diff --git a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index e6f7b54ddcf7..85ae303e8cec 100644 --- a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -17,6 +17,7 @@ package bigquerygetdatasetinfo import ( "context" "fmt" + "strings" bigqueryapi "cloud.google.com/go/bigquery" yaml "github.com/goccy/go-yaml" @@ -48,6 +49,8 @@ type compatibleSource interface { BigQueryClient() *bigqueryapi.Client BigQueryClientCreator() bigqueryds.BigqueryClientCreator UseClientAuthorization() bool + IsDatasetAllowed(projectID, datasetID string) bool + BigQueryAllowedDatasets() []string } // validate compatible sources are still compatible @@ -83,23 +86,48 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } - projectParameter := tools.NewStringParameterWithDefault(projectKey, s.BigQueryProject(), "The Google Cloud project ID containing the dataset.") - datasetParameter := tools.NewStringParameter(datasetKey, "The dataset to get metadata information.") + defaultProjectID := s.BigQueryClient().Project() + projectDescription := "The Google Cloud project ID containing the dataset." + datasetDescription := "The dataset to get metadata information." + var datasetParameter tools.Parameter + allowedDatasets := s.BigQueryAllowedDatasets() + if len(allowedDatasets) > 0 { + if len(allowedDatasets) == 1 { + parts := strings.Split(allowedDatasets[0], ".") + defaultProjectID = parts[0] + datasetID := parts[1] + projectDescription += fmt.Sprintf(" Must be `%s`.", defaultProjectID) + datasetDescription += fmt.Sprintf(" Must be `%s`.", datasetID) + datasetParameter = tools.NewStringParameterWithDefault(datasetKey, datasetID, datasetDescription) + } else { + datasetIDs := []string{} + for _, ds := range allowedDatasets { + datasetIDs = append(datasetIDs, fmt.Sprintf("`%s`", ds)) + } + datasetDescription += fmt.Sprintf(" Must be one of the following: %s.", strings.Join(datasetIDs, ", ")) + datasetParameter = tools.NewStringParameter(datasetKey, datasetDescription) + } + } else { + datasetParameter = tools.NewStringParameter(datasetKey, datasetDescription) + } + + projectParameter := tools.NewStringParameterWithDefault(projectKey, defaultProjectID, projectDescription) parameters := tools.Parameters{projectParameter, datasetParameter} mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) // finish tool setup t := Tool{ - Name: cfg.Name, - Kind: kind, - Parameters: parameters, - AuthRequired: cfg.AuthRequired, - UseClientOAuth: s.UseClientAuthorization(), - ClientCreator: s.BigQueryClientCreator(), - Client: s.BigQueryClient(), - manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired}, - mcpManifest: mcpManifest, + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientAuthorization(), + ClientCreator: s.BigQueryClientCreator(), + Client: s.BigQueryClient(), + IsDatasetAllowed: s.IsDatasetAllowed, + manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired}, + mcpManifest: mcpManifest, } return t, nil } @@ -114,11 +142,12 @@ type Tool struct { UseClientOAuth bool `yaml:"useClientOAuth"` Parameters tools.Parameters `yaml:"parameters"` - Client *bigqueryapi.Client - ClientCreator bigqueryds.BigqueryClientCreator - Statement string - manifest tools.Manifest - mcpManifest tools.McpManifest + Client *bigqueryapi.Client + ClientCreator bigqueryds.BigqueryClientCreator + Statement string + IsDatasetAllowed func(projectID, datasetID string) bool + manifest tools.Manifest + mcpManifest tools.McpManifest } func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { @@ -147,11 +176,32 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken return nil, fmt.Errorf("error creating client from OAuth access token: %w", err) } } + + // Handle fully-qualified dataset ID in the dataset parameter. + if parts := strings.Split(datasetId, "."); len(parts) == 2 { + projectId = parts[0] + datasetId = parts[1] + } + + if !t.IsDatasetAllowed(projectId, datasetId) { + return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) + } + dsHandle := bqClient.DatasetInProject(projectId, datasetId) + // Handle fully-qualified dataset ID in the dataset parameter. + if parts := strings.Split(datasetId, "."); len(parts) == 2 { + projectId = parts[0] + datasetId = parts[1] + } + + if !t.IsDatasetAllowed(projectId, datasetId) { + return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) + } + metadata, err := dsHandle.Metadata(ctx) if err != nil { - return nil, fmt.Errorf("failed to get metadata for dataset %s (in project %s): %w", datasetId, bqClient.Project(), err) + return nil, fmt.Errorf("failed to get metadata for dataset %s (in project %s): %w", datasetId, projectId, err) } return metadata, nil diff --git a/tests/bigquery/bigquery_integration_test.go b/tests/bigquery/bigquery_integration_test.go index 7ff75b091f4a..9cd4afc5e862 100644 --- a/tests/bigquery/bigquery_integration_test.go +++ b/tests/bigquery/bigquery_integration_test.go @@ -274,6 +274,11 @@ func TestBigQueryToolWithDatasetRestriction(t *testing.T) { "source": "my-instance", "description": "Tool to list table within a dataset", }, + "get-dataset-info-restricted": map[string]any{ + "kind": "bigquery-get-dataset-info", + "source": "my-instance", + "description": "Tool to get dataset info", + }, "execute-sql-restricted": map[string]any{ "kind": "bigquery-execute-sql", "source": "my-instance", @@ -318,12 +323,15 @@ func TestBigQueryToolWithDatasetRestriction(t *testing.T) { runListDatasetIdsWithRestriction(t, allowedDatasetName1, allowedDatasetName2) runListTableIdsWithRestriction(t, allowedDatasetName1, disallowedDatasetName, allowedTableName1, allowedForecastTableName1) runListTableIdsWithRestriction(t, allowedDatasetName2, disallowedDatasetName, allowedTableName2, allowedForecastTableName2) + runGetDatasetInfoWithRestriction(t, allowedDatasetName1, disallowedDatasetName) + runGetDatasetInfoWithRestriction(t, allowedDatasetName2, disallowedDatasetName) runExecuteSqlWithRestriction(t, allowedTableNameParam1, disallowedTableNameParam) runExecuteSqlWithRestriction(t, allowedTableNameParam2, disallowedTableNameParam) runConversationalAnalyticsWithRestriction(t, allowedDatasetName1, disallowedDatasetName, allowedTableName1, disallowedTableName) runConversationalAnalyticsWithRestriction(t, allowedDatasetName2, disallowedDatasetName, allowedTableName2, disallowedTableName) runForecastWithRestriction(t, allowedForecastTableFullName1, disallowedForecastTableFullName) runForecastWithRestriction(t, allowedForecastTableFullName2, disallowedForecastTableFullName) + } // getBigQueryParamToolInfo returns statements and param for my-tool for bigquery kind @@ -2206,6 +2214,55 @@ func runListTableIdsWithRestriction(t *testing.T, allowedDatasetName, disallowed } } +func runGetDatasetInfoWithRestriction(t *testing.T, allowedDatasetName, disallowedDatasetName string) { + testCases := []struct { + name string + dataset string + wantStatusCode int + wantInError string + }{ + { + name: "invoke on allowed dataset", + dataset: allowedDatasetName, + wantStatusCode: http.StatusOK, + }, + { + name: "invoke on disallowed dataset", + dataset: disallowedDatasetName, + wantStatusCode: http.StatusBadRequest, + wantInError: fmt.Sprintf("access denied to dataset '%s'", disallowedDatasetName), + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + body := bytes.NewBuffer([]byte(fmt.Sprintf(`{"dataset":"%s"}`, tc.dataset))) + req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:5000/api/tool/get-dataset-info-restricted/invoke", body) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("unable to send request: %s", err) + } + defer resp.Body.Close() + + if resp.StatusCode != tc.wantStatusCode { + bodyBytes, _ := io.ReadAll(resp.Body) + t.Fatalf("unexpected status code: got %d, want %d. Body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes)) + } + + if tc.wantInError != "" { + bodyBytes, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(bodyBytes), tc.wantInError) { + t.Errorf("unexpected error message: got %q, want to contain %q", string(bodyBytes), tc.wantInError) + } + } + }) + } +} + func runExecuteSqlWithRestriction(t *testing.T, allowedTableFullName, disallowedTableFullName string) { allowedTableParts := strings.Split(strings.Trim(allowedTableFullName, "`"), ".") if len(allowedTableParts) != 3 { From 99fb7d79cfb57b61a1503c8037cd0546afafda18 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Wed, 8 Oct 2025 23:52:08 +0000 Subject: [PATCH 2/6] fix merge --- tests/bigquery/bigquery_integration_test.go | 159 -------------------- 1 file changed, 159 deletions(-) diff --git a/tests/bigquery/bigquery_integration_test.go b/tests/bigquery/bigquery_integration_test.go index 9cdbe4784293..bbf8c33ffff4 100644 --- a/tests/bigquery/bigquery_integration_test.go +++ b/tests/bigquery/bigquery_integration_test.go @@ -499,165 +499,6 @@ func TestBigQueryWriteModeProtected(t *testing.T) { runBigQueryWriteModeProtectedTest(t, permanentDatasetName) } -func TestBigQueryWriteModeAllowed(t *testing.T) { - sourceConfig := getBigQueryVars(t) - sourceConfig["writeMode"] = "allowed" - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - datasetName := fmt.Sprintf("temp_toolbox_test_allowed_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) - - client, err := initBigQueryConnection(BigqueryProject) - if err != nil { - t.Fatalf("unable to create BigQuery connection: %s", err) - } - - dataset := client.Dataset(datasetName) - if err := dataset.Create(ctx, &bigqueryapi.DatasetMetadata{Name: datasetName}); err != nil { - t.Fatalf("Failed to create dataset %q: %v", datasetName, err) - } - defer func() { - if err := dataset.DeleteWithContents(ctx); err != nil { - t.Logf("failed to cleanup dataset %s: %v", datasetName, err) - } - }() - - toolsFile := map[string]any{ - "sources": map[string]any{ - "my-instance": sourceConfig, - }, - "tools": map[string]any{ - "my-exec-sql-tool": map[string]any{ - "kind": "bigquery-execute-sql", - "source": "my-instance", - "description": "Tool to execute sql", - }, - }, - } - - cmd, cleanup, err := tests.StartCmd(ctx, toolsFile) - if err != nil { - t.Fatalf("command initialization returned an error: %s", err) - } - defer cleanup() - - waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out) - if err != nil { - t.Logf("toolbox command logs: \n%s", out) - t.Fatalf("toolbox didn't start successfully: %s", err) - } - - runBigQueryWriteModeAllowedTest(t, datasetName) -} - -func TestBigQueryWriteModeBlocked(t *testing.T) { - sourceConfig := getBigQueryVars(t) - sourceConfig["writeMode"] = "blocked" - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - datasetName := fmt.Sprintf("temp_toolbox_test_blocked_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) - tableName := fmt.Sprintf("param_table_blocked_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) - tableNameParam := fmt.Sprintf("`%s.%s.%s`", BigqueryProject, datasetName, tableName) - - client, err := initBigQueryConnection(BigqueryProject) - if err != nil { - t.Fatalf("unable to create BigQuery connection: %s", err) - } - createParamTableStmt, insertParamTableStmt, _, _, _, _, paramTestParams := getBigQueryParamToolInfo(tableNameParam) - teardownTable := setupBigQueryTable(t, ctx, client, createParamTableStmt, insertParamTableStmt, datasetName, tableNameParam, paramTestParams) - defer teardownTable(t) - - toolsFile := map[string]any{ - "sources": map[string]any{"my-instance": sourceConfig}, - "tools": map[string]any{ - "my-exec-sql-tool": map[string]any{"kind": "bigquery-execute-sql", "source": "my-instance", "description": "Tool to execute sql"}, - }, - } - - cmd, cleanup, err := tests.StartCmd(ctx, toolsFile) - if err != nil { - t.Fatalf("command initialization returned an error: %s", err) - } - defer cleanup() - - waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out) - if err != nil { - t.Logf("toolbox command logs: \n%s", out) - t.Fatalf("toolbox didn't start successfully: %s", err) - } - - runBigQueryWriteModeBlockedTest(t, tableNameParam, datasetName) -} - -func TestBigQueryWriteModeProtected(t *testing.T) { - sourceConfig := getBigQueryVars(t) - sourceConfig["writeMode"] = "protected" - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - permanentDatasetName := fmt.Sprintf("perm_dataset_protected_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) - client, err := initBigQueryConnection(BigqueryProject) - if err != nil { - t.Fatalf("unable to create BigQuery connection: %s", err) - } - dataset := client.Dataset(permanentDatasetName) - if err := dataset.Create(ctx, &bigqueryapi.DatasetMetadata{Name: permanentDatasetName}); err != nil { - t.Fatalf("Failed to create dataset %q: %v", permanentDatasetName, err) - } - defer func() { - if err := dataset.DeleteWithContents(ctx); err != nil { - t.Logf("failed to cleanup dataset %s: %v", permanentDatasetName, err) - } - }() - - toolsFile := map[string]any{ - "sources": map[string]any{"my-instance": sourceConfig}, - "tools": map[string]any{ - "my-exec-sql-tool": map[string]any{"kind": "bigquery-execute-sql", "source": "my-instance", "description": "Tool to execute sql"}, - "my-sql-tool-protected": map[string]any{ - "kind": "bigquery-sql", - "source": "my-instance", - "description": "Tool to query from the session", - "statement": "SELECT * FROM my_shared_temp_table", - }, - "my-forecast-tool-protected": map[string]any{ - "kind": "bigquery-forecast", - "source": "my-instance", - "description": "Tool to forecast from session temp table", - }, - "my-analyze-contribution-tool-protected": map[string]any{ - "kind": "bigquery-analyze-contribution", - "source": "my-instance", - "description": "Tool to analyze contribution from session temp table", - }, - }, - } - - cmd, cleanup, err := tests.StartCmd(ctx, toolsFile) - if err != nil { - t.Fatalf("command initialization returned an error: %s", err) - } - defer cleanup() - - waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out) - if err != nil { - t.Logf("toolbox command logs: \n%s", out) - t.Fatalf("toolbox didn't start successfully: %s", err) - } - - runBigQueryWriteModeProtectedTest(t, permanentDatasetName) -} - // getBigQueryParamToolInfo returns statements and param for my-tool for bigquery kind func getBigQueryParamToolInfo(tableName string) (string, string, string, string, string, string, []bigqueryapi.QueryParameter) { createStatement := fmt.Sprintf(` From f99b0c4cd59497775a1a76fb46bef1b5bcdb363d Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Wed, 8 Oct 2025 23:56:13 +0000 Subject: [PATCH 3/6] dedup --- .../bigquerygetdatasetinfo.go | 40 ++++--------------- 1 file changed, 8 insertions(+), 32 deletions(-) diff --git a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index 85ae303e8cec..e56032525c11 100644 --- a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -24,6 +24,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/sources" bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery" "github.com/googleapis/genai-toolbox/internal/tools" + bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon" ) const kind string = "bigquery-get-dataset-info" @@ -88,30 +89,15 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) defaultProjectID := s.BigQueryClient().Project() projectDescription := "The Google Cloud project ID containing the dataset." - datasetDescription := "The dataset to get metadata information." + datasetDescription := "The dataset to get metadata information. Can be in `project.dataset` format." var datasetParameter tools.Parameter - allowedDatasets := s.BigQueryAllowedDatasets() - if len(allowedDatasets) > 0 { - if len(allowedDatasets) == 1 { - parts := strings.Split(allowedDatasets[0], ".") - defaultProjectID = parts[0] - datasetID := parts[1] - projectDescription += fmt.Sprintf(" Must be `%s`.", defaultProjectID) - datasetDescription += fmt.Sprintf(" Must be `%s`.", datasetID) - datasetParameter = tools.NewStringParameterWithDefault(datasetKey, datasetID, datasetDescription) - } else { - datasetIDs := []string{} - for _, ds := range allowedDatasets { - datasetIDs = append(datasetIDs, fmt.Sprintf("`%s`", ds)) - } - datasetDescription += fmt.Sprintf(" Must be one of the following: %s.", strings.Join(datasetIDs, ", ")) - datasetParameter = tools.NewStringParameter(datasetKey, datasetDescription) - } - } else { - datasetParameter = tools.NewStringParameter(datasetKey, datasetDescription) - } + var projectParameter tools.Parameter - projectParameter := tools.NewStringParameterWithDefault(projectKey, defaultProjectID, projectDescription) + projectParameter, datasetParameter = bqutil.InitializeDatasetParameters( + s.BigQueryAllowedDatasets(), + defaultProjectID, + projectKey, datasetKey, + projectDescription, datasetDescription) parameters := tools.Parameters{projectParameter, datasetParameter} mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) @@ -189,16 +175,6 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken dsHandle := bqClient.DatasetInProject(projectId, datasetId) - // Handle fully-qualified dataset ID in the dataset parameter. - if parts := strings.Split(datasetId, "."); len(parts) == 2 { - projectId = parts[0] - datasetId = parts[1] - } - - if !t.IsDatasetAllowed(projectId, datasetId) { - return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) - } - metadata, err := dsHandle.Metadata(ctx) if err != nil { return nil, fmt.Errorf("failed to get metadata for dataset %s (in project %s): %w", datasetId, projectId, err) From ddf119cac16d856eeb6123a7e7174668faf849c9 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Wed, 8 Oct 2025 23:58:00 +0000 Subject: [PATCH 4/6] remove code --- .../bigquerygetdatasetinfo/bigquerygetdatasetinfo.go | 7 ------- 1 file changed, 7 deletions(-) diff --git a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index e56032525c11..b9c40d8e3079 100644 --- a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -17,7 +17,6 @@ package bigquerygetdatasetinfo import ( "context" "fmt" - "strings" bigqueryapi "cloud.google.com/go/bigquery" yaml "github.com/goccy/go-yaml" @@ -163,12 +162,6 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken } } - // Handle fully-qualified dataset ID in the dataset parameter. - if parts := strings.Split(datasetId, "."); len(parts) == 2 { - projectId = parts[0] - datasetId = parts[1] - } - if !t.IsDatasetAllowed(projectId, datasetId) { return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId) } From 2a007a96a26958b095b0bd2b0086b6f20e1756a6 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Thu, 9 Oct 2025 00:13:25 +0000 Subject: [PATCH 5/6] fix test --- tests/bigquery/bigquery_integration_test.go | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/bigquery/bigquery_integration_test.go b/tests/bigquery/bigquery_integration_test.go index bbf8c33ffff4..944842531599 100644 --- a/tests/bigquery/bigquery_integration_test.go +++ b/tests/bigquery/bigquery_integration_test.go @@ -2506,8 +2506,23 @@ func runListDatasetIdsWithRestriction(t *testing.T, allowedDatasetName1, allowed if !ok { t.Fatalf("unable to find result in response body") } - if got != tc.wantResult { - t.Errorf("unexpected result: got %q, want %q", got, tc.wantResult) + + var gotSlice, wantSlice []string + if err := json.Unmarshal([]byte(got), &gotSlice); err != nil { + t.Fatalf("failed to unmarshal 'got' result: %v", err) + } + if err := json.Unmarshal([]byte(tc.wantResult), &wantSlice); err != nil { + t.Fatalf("failed to unmarshal 'want' result: %v", err) + } + + sort.Strings(gotSlice) + sort.Strings(wantSlice) + + sortedGot, _ := json.Marshal(gotSlice) + sortedWant, _ := json.Marshal(wantSlice) + + if string(sortedGot) != string(sortedWant) { + t.Errorf("unexpected result: got %q, want %q", string(sortedGot), string(sortedWant)) } }) } From c96de948207a5d906fc8d1f692b357fb4c8ae59f Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Thu, 9 Oct 2025 19:52:16 +0000 Subject: [PATCH 6/6] fix --- .../bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go index b9c40d8e3079..6d3920c4ac2d 100644 --- a/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go +++ b/internal/tools/bigquery/bigquerygetdatasetinfo/bigquerygetdatasetinfo.go @@ -86,7 +86,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) } - defaultProjectID := s.BigQueryClient().Project() + defaultProjectID := s.BigQueryProject() projectDescription := "The Google Cloud project ID containing the dataset." datasetDescription := "The dataset to get metadata information. Can be in `project.dataset` format." var datasetParameter tools.Parameter