From 5b47370d8697612910b304c52d7a699da6d484c8 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 2 Jun 2026 14:15:23 +0000 Subject: [PATCH 01/25] Add ListDataProducts method to Dataplex source --- internal/sources/dataplex/dataplex.go | 68 ++++++++++++++++++++++----- 1 file changed, 56 insertions(+), 12 deletions(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 31972056f9d4..f75c2fcfd073 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -64,14 +64,15 @@ func (r Config) SourceConfigType() string { func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { // Initializes a Dataplex source - client, dataScanClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project) + client, dataScanClient, dataProductClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project) if err != nil { return nil, err } s := &Source{ - Config: r, - Client: client, - DataScanClient: dataScanClient, + Config: r, + Client: client, + DataScanClient: dataScanClient, + dataProductClient: dataProductClient, } return s, nil @@ -81,8 +82,9 @@ var _ sources.Source = &Source{} type Source struct { Config - Client *dataplexapi.CatalogClient - DataScanClient *dataplexapi.DataScanClient + Client *dataplexapi.CatalogClient + DataScanClient *dataplexapi.DataScanClient + dataProductClient *dataplexapi.DataProductClient } func (s *Source) SourceType() string { @@ -111,30 +113,35 @@ func initDataplexConnection( tracer trace.Tracer, name string, project string, -) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, error) { +) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, *dataplexapi.DataProductClient, error) { ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name) defer span.End() cred, err := google.FindDefaultCredentials(ctx) if err != nil { - return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials for project %q: %w", project, err) + return nil, 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, nil, err + return nil, nil, nil, err } client, err := dataplexapi.NewCatalogClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred)) if err != nil { - return nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err) + return nil, 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 nil, nil, nil, fmt.Errorf("failed to create Dataplex DataScan client for project %q: %w", project, err) } - return client, dataScanClient, nil + + dataProductClient, err := dataplexapi.NewDataProductClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred)) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to create Dataplex DataProduct client for project %q: %w", project, err) + } + return client, dataScanClient, dataProductClient, nil } func (s *Source) LookupEntry(ctx context.Context, name string, view int, aspectTypes []string, entry string) (*dataplexpb.Entry, error) { @@ -293,3 +300,40 @@ func (s *Source) SearchDataQualityScans(ctx context.Context, filter string, page } return results, nil } + +func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { + return s.dataProductClient +} + +func (s *Source) ListDataProducts( + ctx context.Context, + filter string, + pageSize int, + orderBy string, +) ([]*dataplexpb.DataProduct, error) { + parent := fmt.Sprintf("projects/%s/locations/-", s.ProjectID()) + req := &dataplexpb.ListDataProductsRequest{ + Parent: parent, + Filter: filter, + PageSize: int32(pageSize), + OrderBy: orderBy, + } + + it := s.dataProductClient.ListDataProducts(ctx, req) + var results []*dataplexpb.DataProduct + + for range pageSize { + dp, 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 products: code=%s message=%s", st.Code(), st.Message()) + } + return nil, fmt.Errorf("failed to list data products: %w", err) + } + results = append(results, dp) + } + return results, nil +} From 56df565673413d619511c59fe895762650d19b4e Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 2 Jun 2026 16:45:16 +0000 Subject: [PATCH 02/25] Add dataplex-list-data-products tool --- cmd/internal/imports.go | 1 + .../dataplexlistdataproducts.go | 165 +++++++++++++++ .../dataplexlistdataproducts_test.go | 68 +++++++ tests/dataplex/dataplex_integration_test.go | 191 ++++++++++++++++++ 4 files changed, 425 insertions(+) create mode 100644 internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go create mode 100644 internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts_test.go diff --git a/cmd/internal/imports.go b/cmd/internal/imports.go index c5770fb6ce43..8e1f33201cc2 100644 --- a/cmd/internal/imports.go +++ b/cmd/internal/imports.go @@ -113,6 +113,7 @@ import ( _ "github.com/googleapis/mcp-toolbox/internal/tools/couchbase" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataform/dataformcompilelocal" _ "github.com/googleapis/mcp-toolbox/internal/tools/datalineage/datalineagesearchlineage" + _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" _ "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" diff --git a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go new file mode 100644 index 000000000000..958146a58213 --- /dev/null +++ b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go @@ -0,0 +1,165 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexlistdataproducts + +import ( + "context" + "fmt" + "net/http" + + "cloud.google.com/go/dataplex/apiv1/dataplexpb" + "github.com/goccy/go-yaml" + "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" + "github.com/googleapis/mcp-toolbox/internal/sources" + "github.com/googleapis/mcp-toolbox/internal/tools" + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +const resourceType string = "dataplex-list-data-products" + +func init() { + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type compatibleSource interface { + ListDataProducts(ctx context.Context, filter string, pageSize int, orderBy string) ([]*dataplexpb.DataProduct, error) +} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description"` + AuthRequired []string `yaml:"authRequired"` + Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"` + + ScopesRequired []string `yaml:"scopesRequired"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigType() string { + return resourceType +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data products. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching accross all fields at once is not yet supported. E.g. \"display_name:\\\"my-product\\\"\"") + pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Number of returned data products in the page.") + orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Specifies the ordering of results.") + params := parameters.Parameters{filter, pageSize, orderBy} + + t := Tool{ + Config: cfg, + Parameters: params, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: params.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + } + return t, nil +} + +type Tool struct { + Config + Parameters parameters.Parameters + manifest tools.Manifest +} + +func (t Tool) GetName() string { + return t.Name +} + +func (t Tool) GetDescription() string { + return t.Description +} + +func (t Tool) GetAuthRequired() []string { + return t.AuthRequired +} + +func (t Tool) GetAnnotations() *tools.ToolAnnotations { + return tools.GetAnnotationsOrDefault(t.Annotations, tools.NewReadOnlyAnnotations) +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) + if err != nil { + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) + } + paramsMap := params.AsMap() + filter, ok := paramsMap["filter"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'filter' parameter: %v", paramsMap["filter"]), nil) + } + pageSize, ok := paramsMap["pageSize"].(int) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'pageSize' parameter: %v", paramsMap["pageSize"]), nil) + } + orderBy, ok := paramsMap["orderBy"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'orderBy' parameter: %v", paramsMap["orderBy"]), nil) + } + + resp, err := source.ListDataProducts(ctx, filter, pageSize, orderBy) + if err != nil { + return nil, util.ProcessGcpError(err) + } + return resp, nil +} + +func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { + return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { + return false, nil +} + +func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return "Authorization", nil +} + +func (t Tool) GetParameters() parameters.Parameters { + return t.Parameters +} + +func (t Tool) GetScopesRequired() []string { + return t.ScopesRequired +} diff --git a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts_test.go b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts_test.go new file mode 100644 index 000000000000..e6db49596080 --- /dev/null +++ b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexlistdataproducts_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/googleapis/mcp-toolbox/internal/server" + "github.com/googleapis/mcp-toolbox/internal/testutils" + "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" +) + +func TestParseFromYamlDataplexListDataProducts(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + kind: tool + name: example_tool + type: dataplex-list-data-products + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": dataplexlistdataproducts.Config{ + Name: "example_tool", + Type: "dataplex-list-data-products", + Source: "my-instance", + Description: "some description", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + // Parse contents + _, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in)) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } +} diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index f64f5ea2064b..a391187458b9 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -46,6 +46,7 @@ var ( DataplexLookupEntryToolType = "dataplex-lookup-entry" DataplexSearchAspectTypesToolType = "dataplex-search-aspect-types" DataplexSearchDataQualityScansToolType = "dataplex-search-dq-scans" + DataplexListDataProductsToolType = "dataplex-list-data-products" DataplexProject = os.Getenv("DATAPLEX_PROJECT") ) @@ -208,6 +209,19 @@ func initDataplexDataScanConnection(ctx context.Context) (*dataplex.DataScanClie return client, nil } +func initDataplexDataProductConnection(ctx context.Context) (*dataplex.DataProductClient, error) { + cred, err := google.FindDefaultCredentials(ctx) + if err != nil { + return nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err) + } + + client, err := dataplex.NewDataProductClient(ctx, option.WithCredentials(cred)) + if err != nil { + return nil, fmt.Errorf("failed to create Dataplex DataProduct client %w", err) + } + return client, nil +} + func TestDataplexToolEndpoints(t *testing.T) { sourceConfig := getDataplexVars(t) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) @@ -229,6 +243,11 @@ func TestDataplexToolEndpoints(t *testing.T) { t.Fatalf("unable to create Dataplex DataScan connection: %s", err) } + dataplexDataProductClient, err := initDataplexDataProductConnection(ctx) + if err != nil { + t.Fatalf("unable to create Dataplex DataProduct connection: %s", err) + } + // Cleanup older aspecttypes cleanupOldAspectTypes(t, ctx, dataplexClient, 1*time.Hour) @@ -237,15 +256,21 @@ func TestDataplexToolEndpoints(t *testing.T) { tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) aspectTypeId := fmt.Sprintf("param-aspect-type-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) dataScanId := fmt.Sprintf("param-data-scan-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) + dataProductId1 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) + dataProductId2 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) teardownTable1 := setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName) teardownAspectType1 := setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId) teardownDataScan1 := setupDataplexSearchDataQualityScan(t, ctx, dataplexDataScanClient, dataScanId, datasetName, tableName) + teardownDataProduct1 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1) + teardownDataProduct2 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2) time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested defer teardownTable1(t) defer teardownAspectType1(t) defer teardownDataScan1(t) + defer teardownDataProduct1(t) + defer teardownDataProduct2(t) toolsFile := getDataplexToolsConfig(sourceConfig) @@ -269,6 +294,7 @@ func TestDataplexToolEndpoints(t *testing.T) { runDataplexSearchAspectTypesToolInvokeTest(t, aspectTypeId) runDataplexLookupContextToolInvokeTest(t, tableName, datasetName) runDataplexSearchDataQualityScansToolInvokeTest(t, dataScanId, tableName, datasetName) + runDataplexListDataProductsToolInvokeTest(t, dataProductId1, dataProductId2) } func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) { @@ -330,6 +356,48 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C } } +func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataplex.DataProductClient, dataProductId string) func(*testing.T) { + parent := fmt.Sprintf("projects/%s/locations/us-central1", DataplexProject) + ownerEmail := tests.ServiceAccountEmail + if ownerEmail == "" { + ownerEmail = "test-owner@google.com" + } + createReq := &dataplexpb.CreateDataProductRequest{ + Parent: parent, + DataProductId: dataProductId, + DataProduct: &dataplexpb.DataProduct{ + DisplayName: dataProductId, + Description: "Temporary Data Product for MCP Toolbox integration tests", + OwnerEmails: []string{ownerEmail}, + }, + } + + op, err := client.CreateDataProduct(ctx, createReq) + if err != nil { + t.Fatalf("Failed to initiate CreateDataProduct for %s: %v", dataProductId, err) + } + + _, err = op.Wait(ctx) + if err != nil { + t.Fatalf("Failed to wait for CreateDataProduct for %s: %v", dataProductId, err) + } + + return func(t *testing.T) { + deleteReq := &dataplexpb.DeleteDataProductRequest{ + Name: fmt.Sprintf("%s/dataProducts/%s", parent, dataProductId), + } + op, err := client.DeleteDataProduct(ctx, deleteReq) + if err != nil { + t.Errorf("Failed to initiate DeleteDataProduct for %s: %v", dataProductId, err) + return + } + err = op.Wait(ctx) + if err != nil { + t.Logf("Warning: Failed to wait for DeleteDataProduct for %s: %v", dataProductId, err) + } + } +} + func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, aspectTypeId string) func(*testing.T) { parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject) createAspectTypeReq := &dataplexpb.CreateAspectTypeRequest{ @@ -427,6 +495,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any { "description": "Simple dataplex search dq scans tool to test end to end functionality.", "authRequired": []string{"my-google-auth"}, }, + "my-dataplex-list-data-products-tool": map[string]any{ + "type": DataplexListDataProductsToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex list data products tool to test end to end functionality.", + }, + "my-auth-dataplex-list-data-products-tool": map[string]any{ + "type": DataplexListDataProductsToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex list data products tool to test end to end functionality.", + "authRequired": []string{"my-google-auth"}, + }, }, } @@ -459,6 +538,11 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-search-dq-scans-tool", expectedParams: []string{"filter", "data_scan_id", "table_name", "pageSize", "orderBy"}, }, + { + name: "get my-dataplex-list-data-products-tool", + toolName: "my-dataplex-list-data-products-tool", + expectedParams: []string{"filter", "pageSize", "orderBy"}, + }, } for _, tc := range testCases { @@ -1189,3 +1273,110 @@ func runDataplexSearchDataQualityScansToolInvokeTest(t *testing.T, dataScanId st }) } } + +func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 string, dataProductId2 string) { + idToken, err := tests.GetGoogleIdToken(t) + if err != nil { + t.Fatalf("error getting Google ID token: %s", err) + } + + fullDataProductId1 := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId1) + + testCases := []struct { + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantContentKey string + wantValue string + }{ + { + name: "Success - Filter Extracts One Product (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fullDataProductId1, + }, + { + name: "Success - PageSize Limits to One (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"pageSize\":1, \"filter\":\"display_name:\\\"%s\\\" OR display_name:\\\"%s\\\"\"}", dataProductId1, dataProductId2))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + }, + { + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 401, + expectResult: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + for k, v := range tc.requestHeader { + req.Header.Add(k, v) + } + 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 { + t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode) + } + if !tc.expectResult { + return + } + var result map[string]interface{} + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("error parsing response body: %s", err) + } + resultStr, ok := result["result"].(string) + if !ok { + t.Fatalf("expected 'result' field to be a string, got %T", result["result"]) + } + var entries []interface{} + if err := json.Unmarshal([]byte(resultStr), &entries); err != nil { + t.Fatalf("error unmarshalling result string: %v", err) + } + + if len(entries) != 1 { + t.Fatalf("expected exactly one entry, but got %d", len(entries)) + } + entry, ok := entries[0].(map[string]interface{}) + if !ok { + t.Fatalf("expected entry to be a map, got %T", entries[0]) + } + val, ok := entry[tc.wantContentKey].(string) + if !ok { + t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + } + if tc.wantValue != "" && val != tc.wantValue { + t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + } + }) + } +} From 8a316a3685bb6e77cfb4869ebb61ad8728fda294 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 3 Jun 2026 07:14:35 +0000 Subject: [PATCH 03/25] Add documentation for dataplex-list-data-products tool --- docs/KNOWLEDGE_CATALOG_README.md | 4 +- .../integrations/knowledge-catalog/source.md | 11 +++- .../knowledge-catalog-list-data-products.md | 63 +++++++++++++++++++ 3 files changed, 75 insertions(+), 3 deletions(-) create mode 100644 docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md diff --git a/docs/KNOWLEDGE_CATALOG_README.md b/docs/KNOWLEDGE_CATALOG_README.md index b74e1bf9205b..894d4e677401 100644 --- a/docs/KNOWLEDGE_CATALOG_README.md +++ b/docs/KNOWLEDGE_CATALOG_README.md @@ -48,10 +48,12 @@ The Knowledge Catalog MCP server provides the following tools: | Tool Name | Description | |:----------------------|:-----------------------------------------------------------------------------------------------------------------------------| -| `search_entries` | Search for entries in Knowledge Catalog. | +| `search_entries` | Search for entries in Knowledge Catalog. | | `lookup_entry` | Retrieve specific subset of metadata (for example, schema, usage, business overview, and contacts) of a specific data asset. | | `search_aspect_types` | Find aspect types relevant to the query. | | `lookup_context` | Retrieve rich metadata regarding one or more data assets along with their relationships. | +| `search_dq_scans` | Search for Data Quality scans. | +| `list_data_products` | List Data Products for the current project. | ## Custom MCP Server Configuration diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index 626c190e094f..e01f8311ccef 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -373,5 +373,12 @@ This abbreviated syntax works for the qualified predicates except for `label` in 1. Use this tool to retrieve rich metadata regarding one or more data assets along with their relationships. 2. You must provide the `resources` list with full resource names. ### Response -1. Present the requested metadata and relationship information. -``` \ No newline at end of file +1. Present the requested metadata and relationship information. + +## Tool: list_data_products +### Request +1. Use this tool to retrieve all Data Products globally across all locations. +2. You can optionally filter by `display_name` (e.g., "`display_name:\"my-product\"`") or other fields using the Dataplex filter syntax. +### Response +1. Unless asked for a specific data product, respond with all entries returned. +``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md new file mode 100644 index 000000000000..3013b5083940 --- /dev/null +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md @@ -0,0 +1,63 @@ +--- +title: "dataplex-list-data-products" +type: docs +weight: 1 +description: > + A "dataplex-list-data-products" tool allows to list data products. +aliases: + - /integrations/dataplex/tools/dataplex-list-data-products/ +--- + +## About + +A `dataplex-list-data-products` tool lists all Data Products in Knowledge Catalog (formerly known as Dataplex) across all locations (globally). + +`dataplex-list-data-products` optionally accepts the following parameters: + +- `filter` - Filter string to list data products. Use `=` for exact matching and `:` for contains matching. String literals must be enclosed within double quotes. E.g. `display_name:"my-product"`. +- `pageSize` - Number of returned data products in the page. Defaults to `10`. +- `orderBy` - Specifies the ordering of results. + +## Compatible Sources + +{{< compatible-sources >}} + +## Requirements + +### IAM Permissions + +Knowledge Catalog uses [Identity and Access Management (IAM)][iam-overview] to control +user and group access to Knowledge Catalog resources. Toolbox will use your +[Application Default Credentials (ADC)][adc] to authorize and authenticate when +interacting with [Knowledge Catalog][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 [Knowledge Catalog IAM permissions][iam-permissions] +and [Knowledge 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: tool +name: list_data_products +type: dataplex-list-data-products +source: my-dataplex-source +description: Use this tool to list Data Products. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| type | string | true | Must be "dataplex-list-data-products". | +| 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. | From b12df0a03ef4b28b2718d3332728364a41f63751 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 3 Jun 2026 08:42:35 +0000 Subject: [PATCH 04/25] Fix potential resource leak in integration test for Dataplex --- tests/dataplex/dataplex_integration_test.go | 27 ++++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index c2da50c2dd88..09be40ab1310 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -360,7 +360,7 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple parent := fmt.Sprintf("projects/%s/locations/us-central1", DataplexProject) ownerEmail := tests.ServiceAccountEmail if ownerEmail == "" { - ownerEmail = "test-owner@google.com" + t.Errorf("Service account email is required, but tests.ServiceAccountEmail was empty") } createReq := &dataplexpb.CreateDataProductRequest{ Parent: parent, @@ -372,17 +372,7 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple }, } - op, err := client.CreateDataProduct(ctx, createReq) - if err != nil { - t.Fatalf("Failed to initiate CreateDataProduct for %s: %v", dataProductId, err) - } - - _, err = op.Wait(ctx) - if err != nil { - t.Fatalf("Failed to wait for CreateDataProduct for %s: %v", dataProductId, err) - } - - return func(t *testing.T) { + teardown := func(t *testing.T) { deleteReq := &dataplexpb.DeleteDataProductRequest{ Name: fmt.Sprintf("%s/dataProducts/%s", parent, dataProductId), } @@ -396,6 +386,19 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple t.Logf("Warning: Failed to wait for DeleteDataProduct for %s: %v", dataProductId, err) } } + + op, err := client.CreateDataProduct(ctx, createReq) + if err != nil { + t.Fatalf("Failed to initiate CreateDataProduct for %s: %v", dataProductId, err) + } + + _, err = op.Wait(ctx) + if err != nil { + teardown(t) + t.Fatalf("Failed to wait for CreateDataProduct for %s: %v", dataProductId, err) + } + + return teardown } func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, aspectTypeId string) func(*testing.T) { From d0a9e67264ca85e0ce97802290a6f5d225121d97 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 3 Jun 2026 08:48:25 +0000 Subject: [PATCH 05/25] Add defensive parameter checks to Dataplex source method ListDataProducts --- internal/sources/dataplex/dataplex.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 5605346cba3d..d907145d60a8 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -320,6 +320,12 @@ func (s *Source) ListDataProducts( pageSize int, orderBy string, ) ([]*dataplexpb.DataProduct, error) { + if s.dataProductClient == nil { + return nil, fmt.Errorf("dataplex data product client is not initialized") + } + if pageSize <= 0 { + return nil, fmt.Errorf("pageSize must be positive: %d", pageSize) + } parent := fmt.Sprintf("projects/%s/locations/-", s.ProjectID()) req := &dataplexpb.ListDataProductsRequest{ Parent: parent, @@ -331,7 +337,7 @@ func (s *Source) ListDataProducts( it := s.dataProductClient.ListDataProducts(ctx, req) var results []*dataplexpb.DataProduct - for range pageSize { + for len(results) < pageSize { dp, err := it.Next() if err == iterator.Done { break From dfda14a5ab96a8679f5e35697bb400448194250b Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Mon, 8 Jun 2026 04:00:57 +0000 Subject: [PATCH 06/25] Add prebuilt tool for dataplex-list-data-products --- cmd/internal/config_test.go | 2 +- internal/prebuiltconfigs/tools/dataplex.yaml | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/cmd/internal/config_test.go b/cmd/internal/config_test.go index 074da309d1d3..c09b77ea147c 100644 --- a/cmd/internal/config_test.go +++ b/cmd/internal/config_test.go @@ -1832,7 +1832,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", "search_dq_scans"}, + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans", "list_data_products"}, }, }, }, diff --git a/internal/prebuiltconfigs/tools/dataplex.yaml b/internal/prebuiltconfigs/tools/dataplex.yaml index d079dc640b99..90f06af78145 100644 --- a/internal/prebuiltconfigs/tools/dataplex.yaml +++ b/internal/prebuiltconfigs/tools/dataplex.yaml @@ -47,6 +47,12 @@ type: dataplex-search-dq-scans source: dataplex-source description: Use this tool to search for data quality scans in Dataplex. --- +kind: tool +name: list_data_products +type: dataplex-list-data-products +source: dataplex-source +description: Lists Data Products across all locations. +--- kind: toolset name: discovery tools: @@ -55,3 +61,4 @@ tools: - search_aspect_types - lookup_context - search_dq_scans +- list_data_products From d386226305312714bf8f5b3750aa9e9df5cc209b Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Mon, 8 Jun 2026 05:41:57 +0000 Subject: [PATCH 07/25] Add dataplex-get-data-product tool --- cmd/internal/config_test.go | 2 +- cmd/internal/imports.go | 1 + docs/KNOWLEDGE_CATALOG_README.md | 6 +- .../integrations/knowledge-catalog/source.md | 7 + .../knowledge-catalog-get-data-product.md | 48 +++++ go.mod | 10 +- go.sum | 20 +-- internal/prebuiltconfigs/tools/dataplex.yaml | 7 + internal/sources/dataplex/dataplex.go | 64 ++++++- .../dataplexgetdataproduct.go | 164 +++++++++++++++++ .../dataplexgetdataproduct_test.go | 68 +++++++ tests/dataplex/dataplex_integration_test.go | 167 ++++++++++++++++++ 12 files changed, 541 insertions(+), 23 deletions(-) create mode 100644 docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md create mode 100644 internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go create mode 100644 internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct_test.go diff --git a/cmd/internal/config_test.go b/cmd/internal/config_test.go index c09b77ea147c..bc9aa5dcc1b4 100644 --- a/cmd/internal/config_test.go +++ b/cmd/internal/config_test.go @@ -1832,7 +1832,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", "search_dq_scans", "list_data_products"}, + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans", "list_data_products", "get_data_product"}, }, }, }, diff --git a/cmd/internal/imports.go b/cmd/internal/imports.go index 8e1f33201cc2..cd368de9728a 100644 --- a/cmd/internal/imports.go +++ b/cmd/internal/imports.go @@ -113,6 +113,7 @@ import ( _ "github.com/googleapis/mcp-toolbox/internal/tools/couchbase" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataform/dataformcompilelocal" _ "github.com/googleapis/mcp-toolbox/internal/tools/datalineage/datalineagesearchlineage" + _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataproduct" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlookupcontext" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlookupentry" diff --git a/docs/KNOWLEDGE_CATALOG_README.md b/docs/KNOWLEDGE_CATALOG_README.md index 894d4e677401..96ead3248663 100644 --- a/docs/KNOWLEDGE_CATALOG_README.md +++ b/docs/KNOWLEDGE_CATALOG_README.md @@ -7,7 +7,8 @@ The Knowledge Catalog (formerly known as Dataplex) Model Context Protocol (MCP) An editor configured to use the Knowledge Catalog MCP server can use its AI capabilities to help you: - **Search Catalog** - Search for entries in Knowledge Catalog -- **Explore Metadata** - Lookup specific entries and search aspect types +- **Explore Metadata** - Lookup specific entries, search aspect types, and list/retrieve Data Products +- **Data Quality** - Search for data quality scans ## Prerequisites @@ -41,6 +42,8 @@ Once configured, the MCP server will automatically provide Knowledge Catalog cap * "Search for entries related to 'sales' in Knowledge Catalog." * "Look up details for the entry 'projects/my-project/locations/us-central1/entryGroups/my-group/entries/my-entry'." +* "List all Data Products." +* "Get details of the Data Product 'projects/my-project/locations/us-central1/dataProducts/my-product'." ## Server Capabilities @@ -54,6 +57,7 @@ The Knowledge Catalog MCP server provides the following tools: | `lookup_context` | Retrieve rich metadata regarding one or more data assets along with their relationships. | | `search_dq_scans` | Search for Data Quality scans. | | `list_data_products` | List Data Products for the current project. | +| `get_data_product` | Retrieve a specific Data Product. | ## Custom MCP Server Configuration diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index e01f8311ccef..c199999a34eb 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -381,4 +381,11 @@ This abbreviated syntax works for the qualified predicates except for `label` in 2. You can optionally filter by `display_name` (e.g., "`display_name:\"my-product\"`") or other fields using the Dataplex filter syntax. ### Response 1. Unless asked for a specific data product, respond with all entries returned. + +## Tool: get_data_product +### Request +1. Use this tool to retrieve detailed metadata for a specific Data Product. +2. You must provide the full resource name of the Data Product in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +### Response +1. Present the retrieved metadata for the Data Product, including its display name, description, owner emails, asset count, labels, and access groups. ``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md new file mode 100644 index 000000000000..2fe9fce2d9c0 --- /dev/null +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md @@ -0,0 +1,48 @@ +--- +title: "dataplex-get-data-product" +type: docs +weight: 2 +description: > + A "dataplex-get-data-product" tool allows to retrieve a specific Data Product. +aliases: + - /integrations/dataplex/tools/dataplex-get-data-product/ +--- + +## About + +A `dataplex-get-data-product` tool retrieves detailed metadata for a specific Data Product in Knowledge Catalog (formerly known as Dataplex). + +`dataplex-get-data-product` requires the following parameter: + +- `name` - The resource name of the Data Product in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. + +## Compatible Sources + +{{< compatible-sources >}} + +## Requirements + +### IAM Permissions + +To retrieve a data product, your authenticated identity must have the following IAM permissions: +* `dataplex.dataProducts.get` (usually included in `roles/dataplex.viewer` or `roles/dataplex.developer`). + +Refer to the main [Knowledge Catalog Source Requirements](../source.md#requirements) for details on setting up Application Default Credentials (ADC). + +## Example + +```yaml +kind: tool +name: get_data_product +type: dataplex-get-data-product +source: my-dataplex-source +description: Use this tool to retrieve a Data Product. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| type | string | true | Must be "dataplex-get-data-product". | +| 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. | diff --git a/go.mod b/go.mod index 2f5fc844c8ab..bfd0f154dba2 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( cloud.google.com/go/bigtable v1.47.0 cloud.google.com/go/cloudsqlconn v1.21.0 cloud.google.com/go/datacatalog v1.32.0 - cloud.google.com/go/dataplex v1.34.0 + cloud.google.com/go/dataplex v1.35.0 cloud.google.com/go/dataproc/v2 v2.22.0 cloud.google.com/go/firestore v1.22.0 cloud.google.com/go/geminidataanalytics v1.1.0 @@ -73,7 +73,7 @@ require ( go.opentelemetry.io/otel/sdk/metric v1.43.0 go.opentelemetry.io/otel/trace v1.43.0 golang.org/x/oauth2 v0.36.0 - google.golang.org/api v0.279.0 + google.golang.org/api v0.280.0 google.golang.org/genai v1.57.0 google.golang.org/genproto v0.0.0-20260504160031-60b97b32f348 google.golang.org/grpc v1.81.1 @@ -100,7 +100,7 @@ require ( cloud.google.com/go/auth v0.20.0 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect - cloud.google.com/go/iam v1.10.0 // indirect + cloud.google.com/go/iam v1.11.0 // indirect cloud.google.com/go/monitoring v1.29.0 // indirect cloud.google.com/go/trace v1.15.0 // indirect dario.cat/mergo v1.0.2 // indirect @@ -274,8 +274,8 @@ require ( golang.org/x/time v0.15.0 // indirect golang.org/x/tools v0.44.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index a0c384f2b14f..85f432ecca76 100644 --- a/go.sum +++ b/go.sum @@ -21,16 +21,16 @@ cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdB cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= cloud.google.com/go/datacatalog v1.32.0 h1:fyYn8ODkGil5y3zTIqgIhOfzTu1ACaU2o+C750CO6Ac= cloud.google.com/go/datacatalog v1.32.0/go.mod h1:DE272tynQUwheJeQAyVfV+nO8yrdkuDyOgH2LtOrkWM= -cloud.google.com/go/dataplex v1.34.0 h1:WXf+qC/Qhrq6B91HoXYcZJEv1nrLkFpM0HV+JX2SdPs= -cloud.google.com/go/dataplex v1.34.0/go.mod h1:sOazL+Bs/PTxiMHQ5yBboBvEW9qPrpGogx3+RAgfIt8= +cloud.google.com/go/dataplex v1.35.0 h1:EKEhiy/SGYwCH2DZ2r8JEFq1Hx+x+fjJZXRDY3rgPEk= +cloud.google.com/go/dataplex v1.35.0/go.mod h1:B7AFwXU1u3sp7FVQ3IFYnQguGTycJS2mF1voE0lLe1o= cloud.google.com/go/dataproc/v2 v2.22.0 h1:ypUlQKOHMHGv8FQCCNYd0XyM6tAaMDdbcSFBcjYWhbg= cloud.google.com/go/dataproc/v2 v2.22.0/go.mod h1:oARVSa38kAHvSuG+cozsrY2sE6UajGuvOOf9vS+ADHI= cloud.google.com/go/firestore v1.22.0 h1:avooeboIq37vKXobrbPUFhFBxS/c3FqmWoX0xs8dO6E= cloud.google.com/go/firestore v1.22.0/go.mod h1:PaM4i7i7ruALSKmlpHXXZaPObcZw0W7ie5UOPr72iTU= cloud.google.com/go/geminidataanalytics v1.1.0 h1:SpcN1chOo9UrfBFV7EK/BgD9sPdATB4Fm3yITAksoWY= cloud.google.com/go/geminidataanalytics v1.1.0/go.mod h1:YPC+TtyfNM8IdFu/73XwYJVHFyUM6UtZyd//V5kWM2Q= -cloud.google.com/go/iam v1.10.0 h1:cWWt8u8jXv3MzpvBmQgNClvvbVCRukruCJAnoK3fIJY= -cloud.google.com/go/iam v1.10.0/go.mod h1:KP+nKGugNJW4LcLx1uEZcq1ok5sQHFaQehQNl4QDgV4= +cloud.google.com/go/iam v1.11.0 h1:KieQ9Pb+LLPak1O3Rv3GgCxhnmkYf7Xyh0P5HfF1jFM= +cloud.google.com/go/iam v1.11.0/go.mod h1:KP+nKGugNJW4LcLx1uEZcq1ok5sQHFaQehQNl4QDgV4= cloud.google.com/go/logging v1.18.0 h1:KhzZq+1cSkPH9YUaKLLhLtQxIHitVayBmk0sGfoM9+k= cloud.google.com/go/logging v1.18.0/go.mod h1:ZGKnpBaURITh+g/uom2VhbiFoFWvejcrHPDhxFtU/gI= cloud.google.com/go/longrunning v1.0.0 h1:lwzWEYD8+NkYV7dhexOz6kmlvajZA70+bW/xMhRVVdY= @@ -803,8 +803,8 @@ golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhS golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= -google.golang.org/api v0.279.0 h1:hsx2M2OaRcaKtVYK6vXEUnQvdjnend7ZYES+lYaot74= -google.golang.org/api v0.279.0/go.mod h1:B9TqLBwJqVjp1mtt7WeoQwWRwvu/400y5lETOql+giQ= +google.golang.org/api v0.280.0 h1:F4OfEHZhZh6a7uTufJAXXVd/2TQ8EjM4vZH+jX/vFYk= +google.golang.org/api v0.280.0/go.mod h1:oGKmPZRDoD3vdkf6MA7F4VNkR1rxCiuaPSkhsf3EolU= google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= google.golang.org/genai v1.57.0 h1:qTyG2ynz5dQy2jF4CvZdLHHVslhR0heMue+zM1a4GNM= @@ -814,10 +814,10 @@ google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= google.golang.org/genproto v0.0.0-20260504160031-60b97b32f348 h1:JjVGDZYWkJWZcxveJGzfkXC5myDVWAd4dZdgbzrDUv8= google.golang.org/genproto v0.0.0-20260504160031-60b97b32f348/go.mod h1:95PqD4xM+AdOcBGsmgfaofXsiA37uXDtDufVbntT3TU= -google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4 h1:yOzSCGPx+cp5VO7IxvZ9SBFF7j1tZVcNtlHR2iYKtVo= -google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4/go.mod h1:Q9HWtNeE7tM9npdIsEvqXj1QJIvVoeAV3rtXtS715Cw= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 h1:seT2EwLWM78plQ7wcDfuWBc/4FAEAXDDiaSol4ku4qo= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68 h1:PvEgGJf9C/1u5CHkInMg7UFYYUoiaQmW2LbtH0pjB78= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY= diff --git a/internal/prebuiltconfigs/tools/dataplex.yaml b/internal/prebuiltconfigs/tools/dataplex.yaml index 90f06af78145..5d660419b7a3 100644 --- a/internal/prebuiltconfigs/tools/dataplex.yaml +++ b/internal/prebuiltconfigs/tools/dataplex.yaml @@ -53,6 +53,12 @@ type: dataplex-list-data-products source: dataplex-source description: Lists Data Products across all locations. --- +kind: tool +name: get_data_product +type: dataplex-get-data-product +source: dataplex-source +description: Retrieves specific metadata regarding a Data Product. +--- kind: toolset name: discovery tools: @@ -62,3 +68,4 @@ tools: - lookup_context - search_dq_scans - list_data_products +- get_data_product diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index d907145d60a8..226ea3577b86 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -108,6 +108,10 @@ func (s *Source) GetDataScanClient() *dataplexapi.DataScanClient { return s.DataScanClient } +func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { + return s.dataProductClient +} + func initDataplexConnection( ctx context.Context, tracer trace.Tracer, @@ -310,17 +314,13 @@ func (s *Source) SearchDataQualityScans(ctx context.Context, filter string, page return results, nil } -func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { - return s.dataProductClient -} - func (s *Source) ListDataProducts( ctx context.Context, filter string, pageSize int, orderBy string, ) ([]*dataplexpb.DataProduct, error) { - if s.dataProductClient == nil { + if s.GetDataProductClient() == nil { return nil, fmt.Errorf("dataplex data product client is not initialized") } if pageSize <= 0 { @@ -334,7 +334,7 @@ func (s *Source) ListDataProducts( OrderBy: orderBy, } - it := s.dataProductClient.ListDataProducts(ctx, req) + it := s.GetDataProductClient().ListDataProducts(ctx, req) var results []*dataplexpb.DataProduct for len(results) < pageSize { @@ -352,3 +352,55 @@ func (s *Source) ListDataProducts( } return results, nil } + +type AccessGroup struct { + ID string `json:"id"` + DisplayName string `json:"displayName"` + Description string `json:"description"` + GoogleGroup string `json:"googleGroup,omitempty"` + ServiceAccount string `json:"serviceAccount,omitempty"` +} + +type DataProduct struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + Description string `json:"description"` + OwnerEmails []string `json:"ownerEmails"` + AssetCount int32 `json:"assetCount"` + Labels map[string]string `json:"labels"` + AccessGroups []AccessGroup `json:"accessGroups"` +} + +func (s *Source) GetDataProduct(ctx context.Context, name string) (*DataProduct, error) { + if s.GetDataProductClient() == nil { + return nil, fmt.Errorf("dataplex data product client is not initialized") + } + req := &dataplexpb.GetDataProductRequest{ + Name: name, + } + resp, err := s.GetDataProductClient().GetDataProduct(ctx, req) + if err != nil { + return nil, err + } + + accessGroups := []AccessGroup{} + for _, ag := range resp.GetAccessGroups() { + accessGroups = append(accessGroups, AccessGroup{ + ID: ag.GetId(), + DisplayName: ag.GetDisplayName(), + Description: ag.GetDescription(), + GoogleGroup: ag.GetPrincipal().GetGoogleGroup(), + ServiceAccount: ag.GetPrincipal().GetServiceAccount(), + }) + } + + return &DataProduct{ + Name: resp.GetName(), + DisplayName: resp.GetDisplayName(), + Description: resp.GetDescription(), + OwnerEmails: resp.GetOwnerEmails(), + AssetCount: resp.GetAssetCount(), + Labels: resp.GetLabels(), + AccessGroups: accessGroups, + }, nil +} diff --git a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go new file mode 100644 index 000000000000..14b163df0281 --- /dev/null +++ b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go @@ -0,0 +1,164 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexgetdataproduct + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/goccy/go-yaml" + "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" + "github.com/googleapis/mcp-toolbox/internal/sources" + "github.com/googleapis/mcp-toolbox/internal/sources/dataplex" + "github.com/googleapis/mcp-toolbox/internal/tools" + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +const resourceType string = "dataplex-get-data-product" + +func init() { + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type compatibleSource interface { + GetDataProduct(ctx context.Context, name string) (*dataplex.DataProduct, error) +} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description"` + AuthRequired []string `yaml:"authRequired"` + Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"` + + ScopesRequired []string `yaml:"scopesRequired"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigType() string { + return resourceType +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + name := parameters.NewStringParameter("name", "Required. The resource name of the Data Product in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}.") + params := parameters.Parameters{name} + + t := Tool{ + Config: cfg, + Parameters: params, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: params.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + } + return t, nil +} + +type Tool struct { + Config + Parameters parameters.Parameters + manifest tools.Manifest +} + +func (t Tool) GetName() string { + return t.Name +} + +func (t Tool) GetDescription() string { + return t.Description +} + +func (t Tool) GetAuthRequired() []string { + return t.AuthRequired +} + +func (t Tool) GetAnnotations() *tools.ToolAnnotations { + return tools.GetAnnotationsOrDefault(t.Annotations, tools.NewReadOnlyAnnotations) +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) + if err != nil { + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) + } + + paramsMap := params.AsMap() + name, ok := paramsMap["name"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + } + + parts := strings.Split(name, "/") + if len(parts) < 6 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" { + err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}") + return nil, util.NewAgentError(err.Error(), err) + } + + resp, err := source.GetDataProduct(ctx, name) + if err != nil { + return nil, util.ProcessGcpError(err) + } + + return resp, nil +} + +func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { + return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { + return false, nil +} + +func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return "Authorization", nil +} + +func (t Tool) GetParameters() parameters.Parameters { + return t.Parameters +} + +func (t Tool) GetScopesRequired() []string { + return t.ScopesRequired +} diff --git a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct_test.go b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct_test.go new file mode 100644 index 000000000000..49051d963262 --- /dev/null +++ b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexgetdataproduct_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/googleapis/mcp-toolbox/internal/server" + "github.com/googleapis/mcp-toolbox/internal/testutils" + "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataproduct" +) + +func TestParseFromYamlDataplexGetDataProduct(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + kind: tool + name: example_tool + type: dataplex-get-data-product + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": dataplexgetdataproduct.Config{ + Name: "example_tool", + Type: "dataplex-get-data-product", + Source: "my-instance", + Description: "some description", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + // Parse contents + _, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in)) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } +} diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 09be40ab1310..60837ac68677 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -47,6 +47,7 @@ var ( DataplexSearchAspectTypesToolType = "dataplex-search-aspect-types" DataplexSearchDataQualityScansToolType = "dataplex-search-dq-scans" DataplexListDataProductsToolType = "dataplex-list-data-products" + DataplexGetDataProductToolType = "dataplex-get-data-product" DataplexProject = os.Getenv("DATAPLEX_PROJECT") ) @@ -295,6 +296,7 @@ func TestDataplexToolEndpoints(t *testing.T) { runDataplexLookupContextToolInvokeTest(t, tableName, datasetName) runDataplexSearchDataQualityScansToolInvokeTest(t, dataScanId, tableName, datasetName) runDataplexListDataProductsToolInvokeTest(t, dataProductId1, dataProductId2) + runDataplexGetDataProductToolInvokeTest(t, dataProductId1) } func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) { @@ -369,6 +371,18 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple DisplayName: dataProductId, Description: "Temporary Data Product for MCP Toolbox integration tests", OwnerEmails: []string{ownerEmail}, + AccessGroups: map[string]*dataplexpb.DataProduct_AccessGroup{ + "test-group": { + Id: "test-group", + DisplayName: "Test Group", + Description: "Test Group Description", + Principal: &dataplexpb.DataProduct_Principal{ + Type: &dataplexpb.DataProduct_Principal_GoogleGroup{ + GoogleGroup: "test-mcp-group@google.com", + }, + }, + }, + }, }, } @@ -509,6 +523,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any { "description": "Simple dataplex list data products tool to test end to end functionality.", "authRequired": []string{"my-google-auth"}, }, + "my-dataplex-get-data-product-tool": map[string]any{ + "type": DataplexGetDataProductToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex get data product tool to test end to end functionality.", + }, + "my-auth-dataplex-get-data-product-tool": map[string]any{ + "type": DataplexGetDataProductToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex get data product tool to test end to end functionality.", + "authRequired": []string{"my-google-auth"}, + }, }, } @@ -546,6 +571,11 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-list-data-products-tool", expectedParams: []string{"filter", "pageSize", "orderBy"}, }, + { + name: "get my-dataplex-get-data-product-tool", + toolName: "my-dataplex-get-data-product-tool", + expectedParams: []string{"name"}, + }, } for _, tc := range testCases { @@ -1398,3 +1428,140 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri }) } } + +func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) { + idToken, err := tests.GetGoogleIdToken(t) + if err != nil { + t.Fatalf("error getting Google ID token: %s", err) + } + + fullDataProductId := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId) + + testCases := []struct { + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantContentKey string + wantValue string + }{ + { + name: "Success - Get Product (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fullDataProductId, + }, + { + name: "Success - Get Product (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fullDataProductId, + }, + { + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Invalid Name Format", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataProductId))), + wantStatusCode: 500, // Tool validation fails with agent error mapped to 500 + expectResult: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + for k, v := range tc.requestHeader { + req.Header.Add(k, v) + } + 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 { + t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode) + } + if !tc.expectResult { + return + } + var result map[string]interface{} + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("error parsing response body: %s", err) + } + resultStr, ok := result["result"].(string) + if !ok { + t.Fatalf("expected 'result' field to be a string, got %T", result["result"]) + } + var entry map[string]interface{} + if err := json.Unmarshal([]byte(resultStr), &entry); err != nil { + t.Fatalf("error unmarshalling result string: %v", err) + } + + val, ok := entry[tc.wantContentKey].(string) + if !ok { + t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + } + if tc.wantValue != "" && val != tc.wantValue { + t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + } + // Additionally assert key fields are populated + if entry["displayName"] == "" { + t.Errorf("displayName should not be empty") + } + if entry["ownerEmails"] == nil { + t.Errorf("ownerEmails should not be nil") + } + // Assert access groups are mapped correctly + accessGroups, ok := entry["accessGroups"].([]interface{}) + if !ok { + t.Fatalf("expected 'accessGroups' to be a slice, got %T", entry["accessGroups"]) + } + if len(accessGroups) != 1 { + t.Fatalf("expected 1 access group, got %d", len(accessGroups)) + } + ag, ok := accessGroups[0].(map[string]interface{}) + if !ok { + t.Fatalf("expected access group to be a map, got %T", accessGroups[0]) + } + if ag["id"] != "test-group" { + t.Errorf("expected access group id 'test-group', got %q", ag["id"]) + } + if ag["googleGroup"] != "test-mcp-group@google.com" { + t.Errorf("expected googleGroup 'test-mcp-group@google.com', got %q", ag["googleGroup"]) + } + if ag["serviceAccount"] != "" { + t.Errorf("expected serviceAccount to be empty, got %q", ag["serviceAccount"]) + } + }) + } +} From aa84d2530ea13f9e106f6422df3c8e1f4b72493f Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Mon, 8 Jun 2026 06:19:15 +0000 Subject: [PATCH 08/25] Make dataplex-list-data-products tool return only necessary fields --- internal/sources/dataplex/dataplex.go | 18 +++++++++++++++--- .../dataplexlistdataproducts.go | 4 ++-- tests/dataplex/dataplex_integration_test.go | 10 ++++++++++ 3 files changed, 27 insertions(+), 5 deletions(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index d907145d60a8..0fb8d8a99bfd 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -314,12 +314,19 @@ func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { return s.dataProductClient } +type DataProductSummary struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + OwnerEmails []string `json:"ownerEmails"` + AssetCount int32 `json:"assetCount"` +} + func (s *Source) ListDataProducts( ctx context.Context, filter string, pageSize int, orderBy string, -) ([]*dataplexpb.DataProduct, error) { +) ([]*DataProductSummary, error) { if s.dataProductClient == nil { return nil, fmt.Errorf("dataplex data product client is not initialized") } @@ -335,7 +342,7 @@ func (s *Source) ListDataProducts( } it := s.dataProductClient.ListDataProducts(ctx, req) - var results []*dataplexpb.DataProduct + var results []*DataProductSummary for len(results) < pageSize { dp, err := it.Next() @@ -348,7 +355,12 @@ func (s *Source) ListDataProducts( } return nil, fmt.Errorf("failed to list data products: %w", err) } - results = append(results, dp) + results = append(results, &DataProductSummary{ + Name: dp.GetName(), + DisplayName: dp.GetDisplayName(), + OwnerEmails: dp.GetOwnerEmails(), + AssetCount: dp.GetAssetCount(), + }) } return results, nil } diff --git a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go index 958146a58213..fcda1c474dd0 100644 --- a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go +++ b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go @@ -19,10 +19,10 @@ import ( "fmt" "net/http" - "cloud.google.com/go/dataplex/apiv1/dataplexpb" "github.com/goccy/go-yaml" "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" "github.com/googleapis/mcp-toolbox/internal/sources" + "github.com/googleapis/mcp-toolbox/internal/sources/dataplex" "github.com/googleapis/mcp-toolbox/internal/tools" "github.com/googleapis/mcp-toolbox/internal/util" "github.com/googleapis/mcp-toolbox/internal/util/parameters" @@ -45,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T } type compatibleSource interface { - ListDataProducts(ctx context.Context, filter string, pageSize int, orderBy string) ([]*dataplexpb.DataProduct, error) + ListDataProducts(ctx context.Context, filter string, pageSize int, orderBy string) ([]*dataplex.DataProductSummary, error) } type Config struct { diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 09be40ab1310..77fd3acfa09c 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -1395,6 +1395,16 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri if tc.wantValue != "" && val != tc.wantValue { t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) } + // Assert raw SDK fields are cleaned/removed + if _, ok := entry["uid"]; ok { + t.Errorf("expected entry to NOT have 'uid' field, but it was found") + } + if _, ok := entry["etag"]; ok { + t.Errorf("expected entry to NOT have 'etag' field, but it was found") + } + if _, ok := entry["createTime"]; ok { + t.Errorf("expected entry to NOT have 'createTime' field, but it was found") + } }) } } From 3c997dedb0d0b3dee355868dbbbda59381e687af Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 9 Jun 2026 04:07:26 +0000 Subject: [PATCH 09/25] Add dataplex-list-data-assets tool --- cmd/internal/config_test.go | 2 +- cmd/internal/imports.go | 1 + docs/KNOWLEDGE_CATALOG_README.md | 4 +- .../integrations/knowledge-catalog/source.md | 8 + .../knowledge-catalog-list-data-assets.md | 67 ++++++ internal/prebuiltconfigs/tools/dataplex.yaml | 7 + internal/sources/dataplex/dataplex.go | 49 +++++ .../dataplexlistdataassets.go | 178 ++++++++++++++++ .../dataplexlistdataassets_test.go | 68 ++++++ tests/dataplex/dataplex_integration_test.go | 199 +++++++++++++++++- 10 files changed, 578 insertions(+), 5 deletions(-) create mode 100644 docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md create mode 100644 internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go create mode 100644 internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets_test.go diff --git a/cmd/internal/config_test.go b/cmd/internal/config_test.go index bc9aa5dcc1b4..b1db77da7a02 100644 --- a/cmd/internal/config_test.go +++ b/cmd/internal/config_test.go @@ -1832,7 +1832,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", "search_dq_scans", "list_data_products", "get_data_product"}, + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans", "list_data_products", "get_data_product", "list_data_assets"}, }, }, }, diff --git a/cmd/internal/imports.go b/cmd/internal/imports.go index cd368de9728a..ba0f239d0cd0 100644 --- a/cmd/internal/imports.go +++ b/cmd/internal/imports.go @@ -115,6 +115,7 @@ import ( _ "github.com/googleapis/mcp-toolbox/internal/tools/datalineage/datalineagesearchlineage" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataproduct" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" + _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataassets" _ "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" diff --git a/docs/KNOWLEDGE_CATALOG_README.md b/docs/KNOWLEDGE_CATALOG_README.md index 96ead3248663..c70db96822bf 100644 --- a/docs/KNOWLEDGE_CATALOG_README.md +++ b/docs/KNOWLEDGE_CATALOG_README.md @@ -7,7 +7,7 @@ The Knowledge Catalog (formerly known as Dataplex) Model Context Protocol (MCP) An editor configured to use the Knowledge Catalog MCP server can use its AI capabilities to help you: - **Search Catalog** - Search for entries in Knowledge Catalog -- **Explore Metadata** - Lookup specific entries, search aspect types, and list/retrieve Data Products +- **Explore Metadata** - Lookup specific entries, search aspect types, and list/retrieve Data Products and Data Assets - **Data Quality** - Search for data quality scans ## Prerequisites @@ -44,6 +44,7 @@ Once configured, the MCP server will automatically provide Knowledge Catalog cap * "Look up details for the entry 'projects/my-project/locations/us-central1/entryGroups/my-group/entries/my-entry'." * "List all Data Products." * "Get details of the Data Product 'projects/my-project/locations/us-central1/dataProducts/my-product'." +* "List Data Assets for the Data Product 'projects/my-project/locations/us-central1/dataProducts/my-product'." ## Server Capabilities @@ -58,6 +59,7 @@ The Knowledge Catalog MCP server provides the following tools: | `search_dq_scans` | Search for Data Quality scans. | | `list_data_products` | List Data Products for the current project. | | `get_data_product` | Retrieve a specific Data Product. | +| `list_data_assets` | List Data Assets under a Data Product. | ## Custom MCP Server Configuration diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index c199999a34eb..d08717ea897f 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -388,4 +388,12 @@ This abbreviated syntax works for the qualified predicates except for `label` in 2. You must provide the full resource name of the Data Product in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. ### Response 1. Present the retrieved metadata for the Data Product, including its display name, description, owner emails, asset count, labels, and access groups. + +## Tool: list_data_assets +### Request +1. Use this tool to retrieve all Data Assets under a specific Data Product. +2. You must provide the full resource name of the parent Data Product in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +3. You can optionally filter the listed assets using `filter` or limit the response using `pageSize`. +### Response +1. Present the retrieved list of Data Assets, including their names, resources, and labels. ``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md new file mode 100644 index 000000000000..1d8800a27ea7 --- /dev/null +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md @@ -0,0 +1,67 @@ +--- +title: "dataplex-list-data-assets" +type: docs +weight: 2 +description: > + A "dataplex-list-data-assets" tool allows to list Data Assets under a Data Product. +aliases: + - /integrations/dataplex/tools/dataplex-list-data-assets/ +--- + +## About + +A `dataplex-list-data-assets` tool retrieves a list of Data Assets associated with a specific Data Product in Knowledge Catalog (formerly known as Dataplex). + +`dataplex-list-data-assets` requires the following parameter: + +- `name` - The resource name of the parent Data Product in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. + +Optional parameters: + +- `filter` - Filter string to list data assets. +- `pageSize` - Number of returned data assets in the page. +- `orderBy` - Specifies the ordering of results. + +## Compatible Sources + +{{< compatible-sources >}} + +## Requirements + +### IAM Permissions + +Knowledge Catalog uses [Identity and Access Management (IAM)][iam-overview] to control +user and group access to Knowledge Catalog resources. Toolbox will use your +[Application Default Credentials (ADC)][adc] to authorize and authenticate when +interacting with [Knowledge Catalog][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 [Knowledge Catalog IAM permissions][iam-permissions] +and [Knowledge 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: tool +name: list_data_assets +type: dataplex-list-data-assets +source: my-dataplex-source +description: Use this tool to list Data Assets under a Data Product. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| type | string | true | Must be "dataplex-list-data-assets". | +| 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. | diff --git a/internal/prebuiltconfigs/tools/dataplex.yaml b/internal/prebuiltconfigs/tools/dataplex.yaml index 5d660419b7a3..ee0e7d383968 100644 --- a/internal/prebuiltconfigs/tools/dataplex.yaml +++ b/internal/prebuiltconfigs/tools/dataplex.yaml @@ -59,6 +59,12 @@ type: dataplex-get-data-product source: dataplex-source description: Retrieves specific metadata regarding a Data Product. --- +kind: tool +name: list_data_assets +type: dataplex-list-data-assets +source: dataplex-source +description: Lists Data Assets under a Data Product. +--- kind: toolset name: discovery tools: @@ -69,3 +75,4 @@ tools: - search_dq_scans - list_data_products - get_data_product +- list_data_assets diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index d95fb5d3c870..eee316d8a9a1 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -416,3 +416,52 @@ func (s *Source) GetDataProduct(ctx context.Context, name string) (*DataProduct, AccessGroups: accessGroups, }, nil } + +type DataAssetSummary struct { + Name string `json:"name"` + Resource string `json:"resource"` + Labels map[string]string `json:"labels"` +} + +func (s *Source) ListDataAssets( + ctx context.Context, + parent string, + filter string, + pageSize int, + orderBy string, +) ([]*DataAssetSummary, error) { + if s.GetDataProductClient() == nil { + return nil, fmt.Errorf("dataplex data product client is not initialized") + } + if pageSize <= 0 { + return nil, fmt.Errorf("pageSize must be positive: %d", pageSize) + } + req := &dataplexpb.ListDataAssetsRequest{ + Parent: parent, + Filter: filter, + PageSize: int32(pageSize), + OrderBy: orderBy, + } + + it := s.GetDataProductClient().ListDataAssets(ctx, req) + var results []*DataAssetSummary + + for len(results) < pageSize { + asset, 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 assets: code=%s message=%s", st.Code(), st.Message()) + } + return nil, fmt.Errorf("failed to list data assets: %w", err) + } + results = append(results, &DataAssetSummary{ + Name: asset.GetName(), + Resource: asset.GetResource(), + Labels: asset.GetLabels(), + }) + } + return results, nil +} diff --git a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go new file mode 100644 index 000000000000..9c175c7cc962 --- /dev/null +++ b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go @@ -0,0 +1,178 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexlistdataassets + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/goccy/go-yaml" + "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" + "github.com/googleapis/mcp-toolbox/internal/sources" + "github.com/googleapis/mcp-toolbox/internal/sources/dataplex" + "github.com/googleapis/mcp-toolbox/internal/tools" + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +const resourceType string = "dataplex-list-data-assets" + +func init() { + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type compatibleSource interface { + ListDataAssets(ctx context.Context, parent string, filter string, pageSize int, orderBy string) ([]*dataplex.DataAssetSummary, error) +} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description"` + AuthRequired []string `yaml:"authRequired"` + Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"` + + ScopesRequired []string `yaml:"scopesRequired"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigType() string { + return resourceType +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + name := parameters.NewStringParameter("name", "Required. The resource name of the parent Data Product in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}.") + filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data assets. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching accross all fields at once is not yet supported.") + pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Optional. Number of returned data assets in the page.") + orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Optional. Specifies the ordering of results.") + params := parameters.Parameters{name, filter, pageSize, orderBy} + + t := Tool{ + Config: cfg, + Parameters: params, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: params.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + } + return t, nil +} + +type Tool struct { + Config + Parameters parameters.Parameters + manifest tools.Manifest +} + +func (t Tool) GetName() string { + return t.Name +} + +func (t Tool) GetDescription() string { + return t.Description +} + +func (t Tool) GetAuthRequired() []string { + return t.AuthRequired +} + +func (t Tool) GetAnnotations() *tools.ToolAnnotations { + return tools.GetAnnotationsOrDefault(t.Annotations, tools.NewReadOnlyAnnotations) +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) + if err != nil { + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) + } + + paramsMap := params.AsMap() + name, ok := paramsMap["name"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + } + filter, ok := paramsMap["filter"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'filter' parameter: %v", paramsMap["filter"]), nil) + } + pageSize, ok := paramsMap["pageSize"].(int) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'pageSize' parameter: %v", paramsMap["pageSize"]), nil) + } + orderBy, ok := paramsMap["orderBy"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'orderBy' parameter: %v", paramsMap["orderBy"]), nil) + } + + parts := strings.Split(name, "/") + if len(parts) < 6 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" { + err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}") + return nil, util.NewAgentError(err.Error(), err) + } + + resp, err := source.ListDataAssets(ctx, name, filter, pageSize, orderBy) + if err != nil { + return nil, util.ProcessGcpError(err) + } + return resp, nil +} + +func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { + return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { + return false, nil +} + +func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return "Authorization", nil +} + +func (t Tool) GetParameters() parameters.Parameters { + return t.Parameters +} + +func (t Tool) GetScopesRequired() []string { + return t.ScopesRequired +} diff --git a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets_test.go b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets_test.go new file mode 100644 index 000000000000..3984d1b35736 --- /dev/null +++ b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexlistdataassets_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/googleapis/mcp-toolbox/internal/server" + "github.com/googleapis/mcp-toolbox/internal/testutils" + "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataassets" +) + +func TestParseFromYamlDataplexListDataAssets(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + kind: tool + name: example_tool + type: dataplex-list-data-assets + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": dataplexlistdataassets.Config{ + Name: "example_tool", + Type: "dataplex-list-data-assets", + Source: "my-instance", + Description: "some description", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + // Parse contents + _, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in)) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } +} diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 7639d738d916..9f71ce3a294e 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -48,6 +48,7 @@ var ( DataplexSearchDataQualityScansToolType = "dataplex-search-dq-scans" DataplexListDataProductsToolType = "dataplex-list-data-products" DataplexGetDataProductToolType = "dataplex-get-data-product" + DataplexListDataAssetsToolType = "dataplex-list-data-assets" DataplexProject = os.Getenv("DATAPLEX_PROJECT") ) @@ -259,12 +260,14 @@ func TestDataplexToolEndpoints(t *testing.T) { dataScanId := fmt.Sprintf("param-data-scan-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) dataProductId1 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) dataProductId2 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) + dataAssetId := fmt.Sprintf("param-data-asset-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) teardownTable1 := setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName) teardownAspectType1 := setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId) teardownDataScan1 := setupDataplexSearchDataQualityScan(t, ctx, dataplexDataScanClient, dataScanId, datasetName, tableName) teardownDataProduct1 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1) teardownDataProduct2 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2) + teardownDataAsset1 := setupDataplexDataAsset(t, ctx, dataplexDataProductClient, fmt.Sprintf("projects/%s/locations/us/dataProducts/%s", DataplexProject, dataProductId1), dataAssetId, datasetName, tableName) time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested defer teardownTable1(t) @@ -272,6 +275,7 @@ func TestDataplexToolEndpoints(t *testing.T) { defer teardownDataScan1(t) defer teardownDataProduct1(t) defer teardownDataProduct2(t) + defer teardownDataAsset1(t) toolsFile := getDataplexToolsConfig(sourceConfig) @@ -297,6 +301,7 @@ func TestDataplexToolEndpoints(t *testing.T) { runDataplexSearchDataQualityScansToolInvokeTest(t, dataScanId, tableName, datasetName) runDataplexListDataProductsToolInvokeTest(t, dataProductId1, dataProductId2) runDataplexGetDataProductToolInvokeTest(t, dataProductId1) + runDataplexListDataAssetsToolInvokeTest(t, dataProductId1, dataAssetId) } func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) { @@ -359,7 +364,7 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C } func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataplex.DataProductClient, dataProductId string) func(*testing.T) { - parent := fmt.Sprintf("projects/%s/locations/us-central1", DataplexProject) + parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject) ownerEmail := tests.ServiceAccountEmail if ownerEmail == "" { t.Errorf("Service account email is required, but tests.ServiceAccountEmail was empty") @@ -415,6 +420,48 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple return teardown } +func setupDataplexDataAsset(t *testing.T, ctx context.Context, client *dataplex.DataProductClient, parentProductPath string, dataAssetId string, datasetName string, tableName string) func(*testing.T) { + resource := fmt.Sprintf("//bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s", DataplexProject, datasetName, tableName) + createReq := &dataplexpb.CreateDataAssetRequest{ + Parent: parentProductPath, + DataAssetId: dataAssetId, + DataAsset: &dataplexpb.DataAsset{ + Resource: resource, + Labels: map[string]string{ + "env": "test", + }, + }, + } + + teardown := func(t *testing.T) { + deleteReq := &dataplexpb.DeleteDataAssetRequest{ + Name: fmt.Sprintf("%s/dataAssets/%s", parentProductPath, dataAssetId), + } + op, err := client.DeleteDataAsset(ctx, deleteReq) + if err != nil { + t.Errorf("Failed to initiate DeleteDataAsset for %s: %v", dataAssetId, err) + return + } + err = op.Wait(ctx) + if err != nil { + t.Logf("Warning: Failed to wait for DeleteDataAsset for %s: %v", dataAssetId, err) + } + } + + op, err := client.CreateDataAsset(ctx, createReq) + if err != nil { + t.Fatalf("Failed to initiate CreateDataAsset for %s: %v", dataAssetId, err) + } + + _, err = op.Wait(ctx) + if err != nil { + teardown(t) + t.Fatalf("Failed to wait for CreateDataAsset for %s: %v", dataAssetId, err) + } + + return teardown +} + func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, aspectTypeId string) func(*testing.T) { parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject) createAspectTypeReq := &dataplexpb.CreateAspectTypeRequest{ @@ -534,6 +581,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any { "description": "Simple dataplex get data product tool to test end to end functionality.", "authRequired": []string{"my-google-auth"}, }, + "my-dataplex-list-data-assets-tool": map[string]any{ + "type": DataplexListDataAssetsToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex list data assets tool to test end to end functionality.", + }, + "my-auth-dataplex-list-data-assets-tool": map[string]any{ + "type": DataplexListDataAssetsToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex list data assets tool to test end to end functionality.", + "authRequired": []string{"my-google-auth"}, + }, }, } @@ -576,6 +634,11 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-get-data-product-tool", expectedParams: []string{"name"}, }, + { + name: "get my-dataplex-list-data-assets-tool", + toolName: "my-dataplex-list-data-assets-tool", + expectedParams: []string{"name", "filter", "pageSize", "orderBy"}, + }, } for _, tc := range testCases { @@ -1328,7 +1391,7 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri t.Fatalf("error getting Google ID token: %s", err) } - fullDataProductId1 := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId1) + fullDataProductId1 := fmt.Sprintf("projects/%s/locations/us/dataProducts/%s", DataplexProject, dataProductId1) testCases := []struct { name string @@ -1445,7 +1508,7 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) t.Fatalf("error getting Google ID token: %s", err) } - fullDataProductId := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId) + fullDataProductId := fmt.Sprintf("projects/%s/locations/us/dataProducts/%s", DataplexProject, dataProductId) testCases := []struct { name string @@ -1575,3 +1638,133 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) }) } } + +func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, dataAssetId string) { + idToken, err := tests.GetGoogleIdToken(t) + if err != nil { + t.Fatalf("error getting Google ID token: %s", err) + } + + parent := fmt.Sprintf("projects/%s/locations/us/dataProducts/%s", DataplexProject, dataProductId) + + testCases := []struct { + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantContentKey string + wantValue string + }{ + { + name: "Success - List Data Assets (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fmt.Sprintf("%s/dataAssets/%s", parent, dataAssetId), + }, + { + name: "Success - List Data Assets (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fmt.Sprintf("%s/dataAssets/%s", parent, dataAssetId), + }, + { + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Invalid Name Format", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataProductId))), + wantStatusCode: 500, + expectResult: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + for k, v := range tc.requestHeader { + req.Header.Add(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("error when sending a request: %s", err) + } + defer resp.Body.Close() + if resp.StatusCode != tc.wantStatusCode { + t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode) + } + if !tc.expectResult { + return + } + var result map[string]interface{} + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("error parsing response body: %s", err) + } + resultStr, ok := result["result"].(string) + if !ok { + t.Fatalf("expected 'result' field to be a string, got %T", result["result"]) + } + var entries []interface{} + if err := json.Unmarshal([]byte(resultStr), &entries); err != nil { + t.Fatalf("error unmarshalling result string: %v", err) + } + + if len(entries) != 1 { + t.Fatalf("expected exactly one entry, but got %d", len(entries)) + } + entry, ok := entries[0].(map[string]interface{}) + if !ok { + t.Fatalf("expected entry to be a map, got %T", entries[0]) + } + val, ok := entry[tc.wantContentKey].(string) + if !ok { + t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + } + if tc.wantValue != "" && val != tc.wantValue { + t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + } + + // Assert output is cleaned + if entry["resource"] == "" { + t.Errorf("resource should not be empty") + } + if _, ok := entry["uid"]; ok { + t.Errorf("expected entry to NOT have 'uid' field, but it was found") + } + if _, ok := entry["etag"]; ok { + t.Errorf("expected entry to NOT have 'etag' field, but it was found") + } + if _, ok := entry["createTime"]; ok { + t.Errorf("expected entry to NOT have 'createTime' field, but it was found") + } + }) + } +} From 52e0b4e4cfc7203e161a52bba14a0b3bd0860995 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 9 Jun 2026 09:00:37 +0000 Subject: [PATCH 10/25] Add dataplex-get-data-asset tool --- cmd/internal/config_test.go | 2 +- cmd/internal/imports.go | 3 +- docs/KNOWLEDGE_CATALOG_README.md | 2 + .../integrations/knowledge-catalog/source.md | 7 + .../tools/knowledge-catalog-get-data-asset.md | 61 ++++++ internal/prebuiltconfigs/tools/dataplex.yaml | 7 + internal/sources/dataplex/dataplex.go | 29 +++ .../dataplexgetdataasset.go | 164 +++++++++++++++++ .../dataplexgetdataasset_test.go | 68 +++++++ tests/dataplex/dataplex_integration_test.go | 174 ++++++++++++++---- 10 files changed, 482 insertions(+), 35 deletions(-) create mode 100644 docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md create mode 100644 internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go create mode 100644 internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset_test.go diff --git a/cmd/internal/config_test.go b/cmd/internal/config_test.go index b1db77da7a02..3beea4105793 100644 --- a/cmd/internal/config_test.go +++ b/cmd/internal/config_test.go @@ -1832,7 +1832,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", "search_dq_scans", "list_data_products", "get_data_product", "list_data_assets"}, + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans", "list_data_products", "get_data_product", "list_data_assets", "get_data_asset"}, }, }, }, diff --git a/cmd/internal/imports.go b/cmd/internal/imports.go index ba0f239d0cd0..bf05605712b6 100644 --- a/cmd/internal/imports.go +++ b/cmd/internal/imports.go @@ -113,9 +113,10 @@ import ( _ "github.com/googleapis/mcp-toolbox/internal/tools/couchbase" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataform/dataformcompilelocal" _ "github.com/googleapis/mcp-toolbox/internal/tools/datalineage/datalineagesearchlineage" + _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataasset" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataproduct" - _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataassets" + _ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexlistdataproducts" _ "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" diff --git a/docs/KNOWLEDGE_CATALOG_README.md b/docs/KNOWLEDGE_CATALOG_README.md index c70db96822bf..46a394df0443 100644 --- a/docs/KNOWLEDGE_CATALOG_README.md +++ b/docs/KNOWLEDGE_CATALOG_README.md @@ -45,6 +45,7 @@ Once configured, the MCP server will automatically provide Knowledge Catalog cap * "List all Data Products." * "Get details of the Data Product 'projects/my-project/locations/us-central1/dataProducts/my-product'." * "List Data Assets for the Data Product 'projects/my-project/locations/us-central1/dataProducts/my-product'." +* "Get details of the Data Asset 'projects/my-project/locations/us-central1/dataProducts/my-product/dataAssets/my-asset'." ## Server Capabilities @@ -60,6 +61,7 @@ The Knowledge Catalog MCP server provides the following tools: | `list_data_products` | List Data Products for the current project. | | `get_data_product` | Retrieve a specific Data Product. | | `list_data_assets` | List Data Assets under a Data Product. | +| `get_data_asset` | Retrieve specific metadata regarding a Data Asset. | ## Custom MCP Server Configuration diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index d08717ea897f..a19f4f292d00 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -396,4 +396,11 @@ This abbreviated syntax works for the qualified predicates except for `label` in 3. You can optionally filter the listed assets using `filter` or limit the response using `pageSize`. ### Response 1. Present the retrieved list of Data Assets, including their names, resources, and labels. + +## Tool: get_data_asset +### Request +1. Use this tool to retrieve detailed metadata for a specific Data Asset. +2. You must provide the full resource name of the Data Asset in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}`. +### Response +1. Present the retrieved metadata for the Data Asset, including its name, resource, labels, and access group configurations. ``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md new file mode 100644 index 000000000000..0c49088f8f35 --- /dev/null +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md @@ -0,0 +1,61 @@ +--- +title: "dataplex-get-data-asset" +type: docs +weight: 2 +description: > + A "dataplex-get-data-asset" tool retrieve specific metadata regarding a Data Asset. +aliases: + - /integrations/dataplex/tools/dataplex-get-data-asset/ +--- + +## About + +A `dataplex-get-data-asset` tool retrieves detailed metadata for a specific Data Asset in Knowledge Catalog (formerly known as Dataplex). + +`dataplex-get-data-asset` requires the following parameter: + +- `name` - The resource name of the Data Asset in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}`. + +## Compatible Sources + +{{< compatible-sources >}} + +## Requirements + +### IAM Permissions + +Knowledge Catalog uses [Identity and Access Management (IAM)][iam-overview] to control +user and group access to Knowledge Catalog resources. Toolbox will use your +[Application Default Credentials (ADC)][adc] to authorize and authenticate when +interacting with [Knowledge Catalog][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 [Knowledge Catalog IAM permissions][iam-permissions] +and [Knowledge 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: tool +name: get_data_asset +type: dataplex-get-data-asset +source: my-dataplex-source +description: Use this tool to retrieve a Data Asset. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| type | string | true | Must be "dataplex-get-data-asset". | +| 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. | diff --git a/internal/prebuiltconfigs/tools/dataplex.yaml b/internal/prebuiltconfigs/tools/dataplex.yaml index ee0e7d383968..708e37b32029 100644 --- a/internal/prebuiltconfigs/tools/dataplex.yaml +++ b/internal/prebuiltconfigs/tools/dataplex.yaml @@ -65,6 +65,12 @@ type: dataplex-list-data-assets source: dataplex-source description: Lists Data Assets under a Data Product. --- +kind: tool +name: get_data_asset +type: dataplex-get-data-asset +source: dataplex-source +description: Retrieves specific metadata regarding a Data Asset. +--- kind: toolset name: discovery tools: @@ -76,3 +82,4 @@ tools: - list_data_products - get_data_product - list_data_assets +- get_data_asset diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index eee316d8a9a1..d5ab3b179267 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -465,3 +465,32 @@ func (s *Source) ListDataAssets( } return results, nil } + +type DataAsset struct { + Name string `json:"name"` + Resource string `json:"resource"` + Labels map[string]string `json:"labels"` + AccessGroupConfigs map[string]*dataplexpb.DataAsset_AccessGroupConfig `json:"accessGroupConfigs"` +} + +func (s *Source) GetDataAsset(ctx context.Context, name string) (*DataAsset, error) { + if s.GetDataProductClient() == nil { + return nil, fmt.Errorf("dataplex data product client is not initialized") + } + req := &dataplexpb.GetDataAssetRequest{ + Name: name, + } + resp, err := s.GetDataProductClient().GetDataAsset(ctx, req) + if err != nil { + return nil, err + } + + return &DataAsset{ + Name: resp.GetName(), + Resource: resp.GetResource(), + Labels: resp.GetLabels(), + AccessGroupConfigs: resp.GetAccessGroupConfigs(), + }, nil +} + + diff --git a/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go new file mode 100644 index 000000000000..06ad46f4319b --- /dev/null +++ b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go @@ -0,0 +1,164 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexgetdataasset + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/goccy/go-yaml" + "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" + "github.com/googleapis/mcp-toolbox/internal/sources" + "github.com/googleapis/mcp-toolbox/internal/sources/dataplex" + "github.com/googleapis/mcp-toolbox/internal/tools" + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +const resourceType string = "dataplex-get-data-asset" + +func init() { + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type compatibleSource interface { + GetDataAsset(ctx context.Context, name string) (*dataplex.DataAsset, error) +} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description"` + AuthRequired []string `yaml:"authRequired"` + Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"` + + ScopesRequired []string `yaml:"scopesRequired"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigType() string { + return resourceType +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + name := parameters.NewStringParameter("name", "Required. The resource name of the Data Asset in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}.") + params := parameters.Parameters{name} + + t := Tool{ + Config: cfg, + Parameters: params, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: params.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + } + return t, nil +} + +type Tool struct { + Config + Parameters parameters.Parameters + manifest tools.Manifest +} + +func (t Tool) GetName() string { + return t.Name +} + +func (t Tool) GetDescription() string { + return t.Description +} + +func (t Tool) GetAuthRequired() []string { + return t.AuthRequired +} + +func (t Tool) GetAnnotations() *tools.ToolAnnotations { + return tools.GetAnnotationsOrDefault(t.Annotations, tools.NewReadOnlyAnnotations) +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) + if err != nil { + return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err) + } + + paramsMap := params.AsMap() + name, ok := paramsMap["name"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + } + + parts := strings.Split(name, "/") + if len(parts) < 8 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" || parts[6] != "dataAssets" { + err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}") + return nil, util.NewAgentError(err.Error(), err) + } + + resp, err := source.GetDataAsset(ctx, name) + if err != nil { + return nil, util.ProcessGcpError(err) + } + + return resp, nil +} + +func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { + return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { + return false, nil +} + +func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return "Authorization", nil +} + +func (t Tool) GetParameters() parameters.Parameters { + return t.Parameters +} + +func (t Tool) GetScopesRequired() []string { + return t.ScopesRequired +} diff --git a/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset_test.go b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset_test.go new file mode 100644 index 000000000000..21846e21362d --- /dev/null +++ b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataplexgetdataasset_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/googleapis/mcp-toolbox/internal/server" + "github.com/googleapis/mcp-toolbox/internal/testutils" + "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexgetdataasset" +) + +func TestParseFromYamlDataplexGetDataAsset(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + kind: tool + name: example_tool + type: dataplex-get-data-asset + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": dataplexgetdataasset.Config{ + Name: "example_tool", + Type: "dataplex-get-data-asset", + Source: "my-instance", + Description: "some description", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + // Parse contents + _, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in)) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } +} diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 9f71ce3a294e..6c5de9d1c203 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -49,6 +49,7 @@ var ( DataplexListDataProductsToolType = "dataplex-list-data-products" DataplexGetDataProductToolType = "dataplex-get-data-product" DataplexListDataAssetsToolType = "dataplex-list-data-assets" + DataplexGetDataAssetToolType = "dataplex-get-data-asset" DataplexProject = os.Getenv("DATAPLEX_PROJECT") ) @@ -302,6 +303,7 @@ func TestDataplexToolEndpoints(t *testing.T) { runDataplexListDataProductsToolInvokeTest(t, dataProductId1, dataProductId2) runDataplexGetDataProductToolInvokeTest(t, dataProductId1) runDataplexListDataAssetsToolInvokeTest(t, dataProductId1, dataAssetId) + runDataplexGetDataAssetToolInvokeTest(t, dataProductId1, dataAssetId) } func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) { @@ -376,18 +378,6 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple DisplayName: dataProductId, Description: "Temporary Data Product for MCP Toolbox integration tests", OwnerEmails: []string{ownerEmail}, - AccessGroups: map[string]*dataplexpb.DataProduct_AccessGroup{ - "test-group": { - Id: "test-group", - DisplayName: "Test Group", - Description: "Test Group Description", - Principal: &dataplexpb.DataProduct_Principal{ - Type: &dataplexpb.DataProduct_Principal_GoogleGroup{ - GoogleGroup: ownerEmail, - }, - }, - }, - }, }, } @@ -592,6 +582,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any { "description": "Simple dataplex list data assets tool to test end to end functionality.", "authRequired": []string{"my-google-auth"}, }, + "my-dataplex-get-data-asset-tool": map[string]any{ + "type": DataplexGetDataAssetToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex get data asset tool to test end to end functionality.", + }, + "my-auth-dataplex-get-data-asset-tool": map[string]any{ + "type": DataplexGetDataAssetToolType, + "source": "my-dataplex-instance", + "description": "Simple dataplex get data asset tool to test end to end functionality.", + "authRequired": []string{"my-google-auth"}, + }, }, } @@ -639,6 +640,11 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-list-data-assets-tool", expectedParams: []string{"name", "filter", "pageSize", "orderBy"}, }, + { + name: "get my-dataplex-get-data-asset-tool", + toolName: "my-dataplex-get-data-asset-tool", + expectedParams: []string{"name"}, + }, } for _, tc := range testCases { @@ -1614,27 +1620,6 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) if entry["ownerEmails"] == nil { t.Errorf("ownerEmails should not be nil") } - // Assert access groups are mapped correctly - accessGroups, ok := entry["accessGroups"].([]interface{}) - if !ok { - t.Fatalf("expected 'accessGroups' to be a slice, got %T", entry["accessGroups"]) - } - if len(accessGroups) != 1 { - t.Fatalf("expected 1 access group, got %d", len(accessGroups)) - } - ag, ok := accessGroups[0].(map[string]interface{}) - if !ok { - t.Fatalf("expected access group to be a map, got %T", accessGroups[0]) - } - if ag["id"] != "test-group" { - t.Errorf("expected access group id 'test-group', got %q", ag["id"]) - } - if ag["googleGroup"] != tests.ServiceAccountEmail { - t.Errorf("expected googleGroup %q, got %q", tests.ServiceAccountEmail, ag["googleGroup"]) - } - if ag["serviceAccount"] != "" { - t.Errorf("expected serviceAccount to be empty, got %q", ag["serviceAccount"]) - } }) } } @@ -1768,3 +1753,126 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, }) } } + +func runDataplexGetDataAssetToolInvokeTest(t *testing.T, dataProductId string, dataAssetId string) { + idToken, err := tests.GetGoogleIdToken(t) + if err != nil { + t.Fatalf("error getting Google ID token: %s", err) + } + + fullDataAssetId := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s/dataAssets/%s", DataplexProject, dataProductId, dataAssetId) + + testCases := []struct { + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantContentKey string + wantValue string + }{ + { + name: "Success - Get Data Asset (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fullDataAssetId, + }, + { + name: "Success - Get Data Asset (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + wantStatusCode: 200, + expectResult: true, + wantContentKey: "name", + wantValue: fullDataAssetId, + }, + { + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + wantStatusCode: 401, + expectResult: false, + }, + { + name: "Failure - Invalid Name Format", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataAssetId))), + wantStatusCode: 500, + expectResult: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) + if err != nil { + t.Fatalf("unable to create request: %s", err) + } + req.Header.Add("Content-type", "application/json") + for k, v := range tc.requestHeader { + req.Header.Add(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("error when sending a request: %s", err) + } + defer resp.Body.Close() + if resp.StatusCode != tc.wantStatusCode { + t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode) + } + if !tc.expectResult { + return + } + var result map[string]interface{} + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("error parsing response body: %s", err) + } + resultStr, ok := result["result"].(string) + if !ok { + t.Fatalf("expected 'result' field to be a string, got %T", result["result"]) + } + var entry map[string]interface{} + if err := json.Unmarshal([]byte(resultStr), &entry); err != nil { + t.Fatalf("error unmarshalling result string: %v", err) + } + val, ok := entry[tc.wantContentKey].(string) + if !ok { + t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + } + if tc.wantValue != "" && val != tc.wantValue { + t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + } + + // Assert output is cleaned + if entry["resource"] == "" { + t.Errorf("resource should not be empty") + } + if _, ok := entry["uid"]; ok { + t.Errorf("expected entry to NOT have 'uid' field, but it was found") + } + if _, ok := entry["etag"]; ok { + t.Errorf("expected entry to NOT have 'etag' field, but it was found") + } + if _, ok := entry["createTime"]; ok { + t.Errorf("expected entry to NOT have 'createTime' field, but it was found") + } + }) + } +} + From fca8a20e5323d004989eb80c84deb7d8a809ece2 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 16 Jun 2026 03:18:00 +0000 Subject: [PATCH 11/25] Fix existing Dataplex integration test teardown resource leaks and make teardowns concurrent --- tests/dataplex/dataplex_integration_test.go | 57 ++++++++++++++------- 1 file changed, 38 insertions(+), 19 deletions(-) diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 77fd3acfa09c..a8382ee0742b 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -24,6 +24,7 @@ import ( "os" "regexp" "strings" + "sync" "testing" "time" @@ -182,15 +183,17 @@ func setupDataplexSearchDataQualityScan(t *testing.T, ctx context.Context, clien } return func(t *testing.T) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 1*time.Minute) + defer cleanupCancel() deleteDataScanReq := &dataplexpb.DeleteDataScanRequest{ Name: fmt.Sprintf("%s/dataScans/%s", parent, dataScanId), } - op, err := client.DeleteDataScan(ctx, deleteDataScanReq) + op, err := client.DeleteDataScan(cleanupCtx, deleteDataScanReq) if err != nil { t.Errorf("Failed to delete data scan %s: %v", dataScanId, err) return } - if err := op.Wait(ctx); err != nil { + if err := op.Wait(cleanupCtx); err != nil { t.Logf("Warning: Failed to wait for delete data scan %s: %v", dataScanId, err) } } @@ -259,18 +262,26 @@ func TestDataplexToolEndpoints(t *testing.T) { dataProductId1 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) dataProductId2 := fmt.Sprintf("param-data-product-%s", strings.ReplaceAll(uuid.New().String(), "-", "")) - teardownTable1 := setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName) - teardownAspectType1 := setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId) - teardownDataScan1 := setupDataplexSearchDataQualityScan(t, ctx, dataplexDataScanClient, dataScanId, datasetName, tableName) - teardownDataProduct1 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1) - teardownDataProduct2 := setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2) + var teardowns []func(*testing.T) + teardowns = append(teardowns, setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName)) + teardowns = append(teardowns, setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId)) + teardowns = append(teardowns, setupDataplexSearchDataQualityScan(t, ctx, dataplexDataScanClient, dataScanId, datasetName, tableName)) + teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1)) + teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2)) time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested - defer teardownTable1(t) - defer teardownAspectType1(t) - defer teardownDataScan1(t) - defer teardownDataProduct1(t) - defer teardownDataProduct2(t) + // Execute teardowns concurrently using a WaitGroup to minimize overall test cleanup duration + defer func() { + var wg sync.WaitGroup + for _, fn := range teardowns { + wg.Add(1) + go func(cleanup func(*testing.T)) { + defer wg.Done() + cleanup(t) + }(fn) + } + wg.Wait() + }() toolsFile := getDataplexToolsConfig(sourceConfig) @@ -325,14 +336,17 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C } return func(t *testing.T) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 1*time.Minute) + defer cleanupCancel() + // tear down table dropSQL := fmt.Sprintf("drop table %s.%s", datasetName, tableName) - dropJob, err := client.Query(dropSQL).Run(ctx) + dropJob, err := client.Query(dropSQL).Run(cleanupCtx) if err != nil { t.Errorf("Failed to start drop table job for %s: %v", tableName, err) return } - dropStatus, err := dropJob.Wait(ctx) + dropStatus, err := dropJob.Wait(cleanupCtx) if err != nil { t.Errorf("Failed to wait for drop table job for %s: %v", tableName, err) return @@ -343,11 +357,11 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C // tear down dataset datasetToTeardown := client.Dataset(datasetName) - tablesIterator := datasetToTeardown.Tables(ctx) + tablesIterator := datasetToTeardown.Tables(cleanupCtx) _, err = tablesIterator.Next() if err == iterator.Done { - if err := datasetToTeardown.Delete(ctx); err != nil { + if err := datasetToTeardown.Delete(cleanupCtx); err != nil { t.Errorf("Failed to delete dataset %s: %v", datasetName, err) } } else if err != nil { @@ -373,15 +387,17 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple } teardown := func(t *testing.T) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 1*time.Minute) + defer cleanupCancel() deleteReq := &dataplexpb.DeleteDataProductRequest{ Name: fmt.Sprintf("%s/dataProducts/%s", parent, dataProductId), } - op, err := client.DeleteDataProduct(ctx, deleteReq) + op, err := client.DeleteDataProduct(cleanupCtx, deleteReq) if err != nil { t.Errorf("Failed to initiate DeleteDataProduct for %s: %v", dataProductId, err) return } - err = op.Wait(ctx) + err = op.Wait(cleanupCtx) if err != nil { t.Logf("Warning: Failed to wait for DeleteDataProduct for %s: %v", dataProductId, err) } @@ -420,11 +436,14 @@ func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client } return func(t *testing.T) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 1*time.Minute) + defer cleanupCancel() + // tear down aspect type deleteAspectTypeReq := &dataplexpb.DeleteAspectTypeRequest{ Name: fmt.Sprintf("%s/aspectTypes/%s", parent, aspectTypeId), } - if _, err := client.DeleteAspectType(ctx, deleteAspectTypeReq); err != nil { + if _, err := client.DeleteAspectType(cleanupCtx, deleteAspectTypeReq); err != nil { t.Errorf("Failed to delete aspect type %s: %v", aspectTypeId, err) } } From 988b62d2b730e66d04db6c6f60e2ee1269ee7f79 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 16 Jun 2026 06:03:45 +0000 Subject: [PATCH 12/25] Reduce sleep time for aspect type propagation in Dataplex integration test --- tests/dataplex/dataplex_integration_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 015181cdc0c1..c00cb5e200c6 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -270,7 +270,7 @@ func TestDataplexToolEndpoints(t *testing.T) { teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1)) teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2)) - time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested + time.Sleep(1*time.Minute) // wait for table and aspect type to be ingested // Execute teardowns concurrently using a WaitGroup to minimize overall test cleanup duration defer func() { var wg sync.WaitGroup @@ -1517,7 +1517,7 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataProductId))), - wantStatusCode: 500, // Tool validation fails with agent error mapped to 500 + wantStatusCode: 200, // Tool validation returns AgentError which maps to 200 expectResult: false, }, } @@ -1588,8 +1588,8 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) if ag["googleGroup"] != tests.ServiceAccountEmail { t.Errorf("expected googleGroup %q, got %q", tests.ServiceAccountEmail, ag["googleGroup"]) } - if ag["serviceAccount"] != "" { - t.Errorf("expected serviceAccount to be empty, got %q", ag["serviceAccount"]) + if ag["serviceAccount"] != nil { + t.Errorf("expected serviceAccount to be nil, got %q", ag["serviceAccount"]) } }) } From 4336b2f77178ea1d644c13f1f185f9aae052e6df Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 17 Jun 2026 05:27:16 +0000 Subject: [PATCH 13/25] Reduce sleep time for aspect type propagation in Dataplex integration tests --- tests/dataplex/dataplex_integration_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index a8382ee0742b..b07f5c830293 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -269,7 +269,7 @@ func TestDataplexToolEndpoints(t *testing.T) { teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId1)) teardowns = append(teardowns, setupDataplexDataProduct(t, ctx, dataplexDataProductClient, dataProductId2)) - time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested + time.Sleep(1*time.Minute) // wait for table and aspect type to be ingested // Execute teardowns concurrently using a WaitGroup to minimize overall test cleanup duration defer func() { var wg sync.WaitGroup From 35cb5fc605a8b6ffbdec88acd70a1ba77bf16fed Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 17 Jun 2026 05:46:30 +0000 Subject: [PATCH 14/25] Update dataplex-list-data-products tool output format --- internal/sources/dataplex/dataplex.go | 25 ++++-- tests/dataplex/dataplex_integration_test.go | 88 +++++++++++---------- 2 files changed, 64 insertions(+), 49 deletions(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 0fb8d8a99bfd..de4050bd2d24 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -17,6 +17,7 @@ package dataplex import ( "context" "fmt" + "strings" dataplexapi "cloud.google.com/go/dataplex/apiv1" "cloud.google.com/go/dataplex/apiv1/dataplexpb" @@ -315,10 +316,11 @@ func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { } type DataProductSummary struct { - Name string `json:"name"` - DisplayName string `json:"displayName"` - OwnerEmails []string `json:"ownerEmails"` - AssetCount int32 `json:"assetCount"` + LocationID string `json:"locationId"` + DataProductID string `json:"dataProductId"` + DisplayName string `json:"displayName"` + OwnerEmails []string `json:"ownerEmails"` + AssetCount int32 `json:"assetCount"` } func (s *Source) ListDataProducts( @@ -355,11 +357,18 @@ func (s *Source) ListDataProducts( } return nil, fmt.Errorf("failed to list data products: %w", err) } + parts := strings.Split(dp.GetName(), "/") + var locationId, dataProductId string + if len(parts) >= 6 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" { + locationId = parts[3] + dataProductId = parts[5] + } results = append(results, &DataProductSummary{ - Name: dp.GetName(), - DisplayName: dp.GetDisplayName(), - OwnerEmails: dp.GetOwnerEmails(), - AssetCount: dp.GetAssetCount(), + LocationID: locationId, + DataProductID: dataProductId, + DisplayName: dp.GetDisplayName(), + OwnerEmails: dp.GetOwnerEmails(), + AssetCount: dp.GetAssetCount(), }) } return results, nil diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index b07f5c830293..2a0a02ad55b6 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -1317,52 +1317,51 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri t.Fatalf("error getting Google ID token: %s", err) } - fullDataProductId1 := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId1) - testCases := []struct { - name string - api string - requestHeader map[string]string - requestBody io.Reader - wantStatusCode int - expectResult bool - wantContentKey string - wantValue string + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantLocationID string + wantDataProductID string }{ { - name: "Success - Filter Extracts One Product (Authorized)", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": idToken}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fullDataProductId1, + name: "Success - Filter Extracts One Product (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us-central1", + wantDataProductID: dataProductId1, }, { - name: "Success - PageSize Limits to One (Un-authorized)", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-products-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"pageSize\":1, \"filter\":\"display_name:\\\"%s\\\" OR display_name:\\\"%s\\\"\"}", dataProductId1, dataProductId2))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", + name: "Success - PageSize Limits to One (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"pageSize\":1, \"filter\":\"display_name:\\\"%s\\\" OR display_name:\\\"%s\\\"\"}", dataProductId1, dataProductId2))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us-central1", + wantDataProductID: "", }, { - name: "Failure - Invalid Authorization Token", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), - wantStatusCode: 401, - expectResult: false, + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 401, + expectResult: false, }, { - name: "Failure - Without Authorization Token", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), - wantStatusCode: 401, - expectResult: false, + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-products-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"filter\":\"display_name:\\\"%s\\\"\"}", dataProductId1))), + wantStatusCode: 401, + expectResult: false, }, } @@ -1407,12 +1406,19 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri if !ok { t.Fatalf("expected entry to be a map, got %T", entries[0]) } - val, ok := entry[tc.wantContentKey].(string) + locID, ok := entry["locationId"].(string) + if !ok { + t.Fatalf("expected entry to have key 'locationId' as string, but it was not found or not a string in %v", entry) + } + if tc.wantLocationID != "" && locID != tc.wantLocationID { + t.Fatalf("expected locationId to be %q, got %q", tc.wantLocationID, locID) + } + prodID, ok := entry["dataProductId"].(string) if !ok { - t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + t.Fatalf("expected entry to have key 'dataProductId' as string, but it was not found or not a string in %v", entry) } - if tc.wantValue != "" && val != tc.wantValue { - t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + if tc.wantDataProductID != "" && prodID != tc.wantDataProductID { + t.Fatalf("expected dataProductId to be %q, got %q", tc.wantDataProductID, prodID) } // Assert raw SDK fields are cleaned/removed if _, ok := entry["uid"]; ok { From 18fe08822560112e13ae86594de67f812135b413 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 17 Jun 2026 06:23:18 +0000 Subject: [PATCH 15/25] Update dataplex-get-data-product tool API --- .../integrations/knowledge-catalog/source.md | 2 +- .../knowledge-catalog-get-data-product.md | 5 +- internal/sources/dataplex/dataplex.go | 16 ++- .../dataplexgetdataproduct.go | 22 ++--- tests/dataplex/dataplex_integration_test.go | 99 +++++++++---------- 5 files changed, 75 insertions(+), 69 deletions(-) diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index c199999a34eb..11aac2dee320 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -385,7 +385,7 @@ This abbreviated syntax works for the qualified predicates except for `label` in ## Tool: get_data_product ### Request 1. Use this tool to retrieve detailed metadata for a specific Data Product. -2. You must provide the full resource name of the Data Product in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +2. You must provide `locationId` and `dataProductId`. ### Response 1. Present the retrieved metadata for the Data Product, including its display name, description, owner emails, asset count, labels, and access groups. ``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md index 2fe9fce2d9c0..4aed87f97e9d 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-product.md @@ -12,9 +12,10 @@ aliases: A `dataplex-get-data-product` tool retrieves detailed metadata for a specific Data Product in Knowledge Catalog (formerly known as Dataplex). -`dataplex-get-data-product` requires the following parameter: +`dataplex-get-data-product` requires the following parameters: -- `name` - The resource name of the Data Product in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +- `locationId` - The location ID (e.g., `us`, `us-central1`) where the Data Product is located. +- `dataProductId` - The unique ID of the Data Product. ## Compatible Sources diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index ae7ab07dfdce..df906bfd6c19 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -383,7 +383,8 @@ type AccessGroup struct { } type DataProduct struct { - Name string `json:"name"` + LocationID string `json:"locationId"` + DataProductID string `json:"dataProductId"` DisplayName string `json:"displayName"` Description string `json:"description"` OwnerEmails []string `json:"ownerEmails"` @@ -392,10 +393,11 @@ type DataProduct struct { AccessGroups []AccessGroup `json:"accessGroups"` } -func (s *Source) GetDataProduct(ctx context.Context, name string) (*DataProduct, error) { +func (s *Source) GetDataProduct(ctx context.Context, locationId string, dataProductId string) (*DataProduct, error) { if s.GetDataProductClient() == nil { return nil, fmt.Errorf("dataplex data product client is not initialized") } + name := fmt.Sprintf("projects/%s/locations/%s/dataProducts/%s", s.ProjectID(), locationId, dataProductId) req := &dataplexpb.GetDataProductRequest{ Name: name, } @@ -415,8 +417,16 @@ func (s *Source) GetDataProduct(ctx context.Context, name string) (*DataProduct, }) } + parts := strings.Split(resp.GetName(), "/") + var locId, prodId string + if len(parts) >= 6 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" { + locId = parts[3] + prodId = parts[5] + } + return &DataProduct{ - Name: resp.GetName(), + LocationID: locId, + DataProductID: prodId, DisplayName: resp.GetDisplayName(), Description: resp.GetDescription(), OwnerEmails: resp.GetOwnerEmails(), diff --git a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go index 14b163df0281..194671ea7deb 100644 --- a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go +++ b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go @@ -18,7 +18,6 @@ import ( "context" "fmt" "net/http" - "strings" "github.com/goccy/go-yaml" "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" @@ -46,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T } type compatibleSource interface { - GetDataProduct(ctx context.Context, name string) (*dataplex.DataProduct, error) + GetDataProduct(ctx context.Context, locationId string, dataProductId string) (*dataplex.DataProduct, error) } type Config struct { @@ -68,8 +67,9 @@ func (cfg Config) ToolConfigType() string { } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { - name := parameters.NewStringParameter("name", "Required. The resource name of the Data Product in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}.") - params := parameters.Parameters{name} + locationId := parameters.NewStringParameter("locationId", "Required. The location ID (e.g., 'us', 'us-central1') where the Data Product is located.") + dataProductId := parameters.NewStringParameter("dataProductId", "Required. The unique ID of the Data Product.") + params := parameters.Parameters{locationId, dataProductId} t := Tool{ Config: cfg, @@ -116,18 +116,16 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para } paramsMap := params.AsMap() - name, ok := paramsMap["name"].(string) + locationId, ok := paramsMap["locationId"].(string) if !ok { - return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'locationId' parameter: %v", paramsMap["locationId"]), nil) } - - parts := strings.Split(name, "/") - if len(parts) < 6 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" { - err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}") - return nil, util.NewAgentError(err.Error(), err) + dataProductId, ok := paramsMap["dataProductId"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'dataProductId' parameter: %v", paramsMap["dataProductId"]), nil) } - resp, err := source.GetDataProduct(ctx, name) + resp, err := source.GetDataProduct(ctx, locationId, dataProductId) if err != nil { return nil, util.ProcessGcpError(err) } diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 7f39d4b7aad8..2f636ccfaaab 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -593,7 +593,7 @@ func runDataplexToolGetTest(t *testing.T) { { name: "get my-dataplex-get-data-product-tool", toolName: "my-dataplex-get-data-product-tool", - expectedParams: []string{"name"}, + expectedParams: []string{"locationId", "dataProductId"}, }, } @@ -1470,61 +1470,51 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) t.Fatalf("error getting Google ID token: %s", err) } - fullDataProductId := fmt.Sprintf("projects/%s/locations/us-central1/dataProducts/%s", DataplexProject, dataProductId) - testCases := []struct { - name string - api string - requestHeader map[string]string - requestBody io.Reader - wantStatusCode int - expectResult bool - wantContentKey string - wantValue string + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantLocationID string + wantDataProductID string }{ { - name: "Success - Get Product (Authorized)", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": idToken}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fullDataProductId, - }, - { - name: "Success - Get Product (Un-authorized)", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fullDataProductId, + name: "Success - Get Product (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us-central1\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us-central1", + wantDataProductID: dataProductId, }, { - name: "Failure - Invalid Authorization Token", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), - wantStatusCode: 401, - expectResult: false, + name: "Success - Get Product (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us-central1\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us-central1", + wantDataProductID: dataProductId, }, { - name: "Failure - Without Authorization Token", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataProductId))), - wantStatusCode: 401, - expectResult: false, + name: "Failure - Invalid Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us-central1\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 401, + expectResult: false, }, { - name: "Failure - Invalid Name Format", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-product-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataProductId))), - wantStatusCode: 200, // Tool validation returns AgentError which maps to 200 - expectResult: false, + name: "Failure - Without Authorization Token", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-product-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us-central1\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 401, + expectResult: false, }, } @@ -1562,12 +1552,19 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string) t.Fatalf("error unmarshalling result string: %v", err) } - val, ok := entry[tc.wantContentKey].(string) + locID, ok := entry["locationId"].(string) + if !ok { + t.Fatalf("expected entry to have key 'locationId' as string, but it was not found or not a string in %v", entry) + } + if tc.wantLocationID != "" && locID != tc.wantLocationID { + t.Fatalf("expected locationId to be %q, got %q", tc.wantLocationID, locID) + } + prodID, ok := entry["dataProductId"].(string) if !ok { - t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + t.Fatalf("expected entry to have key 'dataProductId' as string, but it was not found or not a string in %v", entry) } - if tc.wantValue != "" && val != tc.wantValue { - t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + if tc.wantDataProductID != "" && prodID != tc.wantDataProductID { + t.Fatalf("expected dataProductId to be %q, got %q", tc.wantDataProductID, prodID) } // Additionally assert key fields are populated if entry["displayName"] == "" { From 2227d05b6a0a2853a43b5ea378eb5a08328b028e Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 17 Jun 2026 08:48:17 +0000 Subject: [PATCH 16/25] Update dataplex-list-data-assets tool API --- .../integrations/knowledge-catalog/source.md | 2 +- .../knowledge-catalog-list-data-assets.md | 5 +- internal/sources/dataplex/dataplex.go | 27 ++++-- .../dataplexlistdataassets.go | 24 +++-- tests/dataplex/dataplex_integration_test.go | 92 ++++++++++--------- 5 files changed, 83 insertions(+), 67 deletions(-) diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index bef67de65f52..343e17ecbac2 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -392,7 +392,7 @@ This abbreviated syntax works for the qualified predicates except for `label` in ## Tool: list_data_assets ### Request 1. Use this tool to retrieve all Data Assets under a specific Data Product. -2. You must provide the full resource name of the parent Data Product in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +2. You must provide `locationId` and `dataProductId`. 3. You can optionally filter the listed assets using `filter` or limit the response using `pageSize`. ### Response 1. Present the retrieved list of Data Assets, including their names, resources, and labels. diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md index 1d8800a27ea7..8585ab40d237 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-assets.md @@ -12,9 +12,10 @@ aliases: A `dataplex-list-data-assets` tool retrieves a list of Data Assets associated with a specific Data Product in Knowledge Catalog (formerly known as Dataplex). -`dataplex-list-data-assets` requires the following parameter: +`dataplex-list-data-assets` requires the following parameters: -- `name` - The resource name of the parent Data Product in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}`. +- `locationId` - The location ID (e.g., `us`, `us-central1`) where the Data Product is located. +- `dataProductId` - The unique ID of the parent Data Product. Optional parameters: diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index a1378a85e737..7fa9041348e2 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -437,14 +437,17 @@ func (s *Source) GetDataProduct(ctx context.Context, locationId string, dataProd } type DataAssetSummary struct { - Name string `json:"name"` - Resource string `json:"resource"` - Labels map[string]string `json:"labels"` + LocationID string `json:"locationId"` + DataProductID string `json:"dataProductId"` + DataAsset string `json:"dataAsset"` + ResourceUri string `json:"resourceUri"` + Labels map[string]string `json:"labels"` } func (s *Source) ListDataAssets( ctx context.Context, - parent string, + locationId string, + dataProductId string, filter string, pageSize int, orderBy string, @@ -455,6 +458,7 @@ func (s *Source) ListDataAssets( if pageSize <= 0 { return nil, fmt.Errorf("pageSize must be positive: %d", pageSize) } + parent := fmt.Sprintf("projects/%s/locations/%s/dataProducts/%s", s.ProjectID(), locationId, dataProductId) req := &dataplexpb.ListDataAssetsRequest{ Parent: parent, Filter: filter, @@ -476,10 +480,19 @@ func (s *Source) ListDataAssets( } return nil, fmt.Errorf("failed to list data assets: %w", err) } + parts := strings.Split(asset.GetName(), "/") + var locId, prodId, assetId string + if len(parts) >= 8 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" && parts[6] == "dataAssets" { + locId = parts[3] + prodId = parts[5] + assetId = parts[7] + } results = append(results, &DataAssetSummary{ - Name: asset.GetName(), - Resource: asset.GetResource(), - Labels: asset.GetLabels(), + LocationID: locId, + DataProductID: prodId, + DataAsset: assetId, + ResourceUri: asset.GetResource(), + Labels: asset.GetLabels(), }) } return results, nil diff --git a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go index 9c175c7cc962..3c5cbe096e00 100644 --- a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go +++ b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go @@ -18,7 +18,6 @@ import ( "context" "fmt" "net/http" - "strings" "github.com/goccy/go-yaml" "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" @@ -46,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T } type compatibleSource interface { - ListDataAssets(ctx context.Context, parent string, filter string, pageSize int, orderBy string) ([]*dataplex.DataAssetSummary, error) + ListDataAssets(ctx context.Context, locationId string, dataProductId string, filter string, pageSize int, orderBy string) ([]*dataplex.DataAssetSummary, error) } type Config struct { @@ -68,11 +67,12 @@ func (cfg Config) ToolConfigType() string { } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { - name := parameters.NewStringParameter("name", "Required. The resource name of the parent Data Product in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}.") + locationId := parameters.NewStringParameter("locationId", "Required. The location ID (e.g., 'us', 'us-central1') where the Data Product is located.") + dataProductId := parameters.NewStringParameter("dataProductId", "Required. The unique ID of the parent Data Product.") filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data assets. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching accross all fields at once is not yet supported.") pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Optional. Number of returned data assets in the page.") orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Optional. Specifies the ordering of results.") - params := parameters.Parameters{name, filter, pageSize, orderBy} + params := parameters.Parameters{locationId, dataProductId, filter, pageSize, orderBy} t := Tool{ Config: cfg, @@ -119,9 +119,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para } paramsMap := params.AsMap() - name, ok := paramsMap["name"].(string) + locationId, ok := paramsMap["locationId"].(string) if !ok { - return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'locationId' parameter: %v", paramsMap["locationId"]), nil) + } + dataProductId, ok := paramsMap["dataProductId"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'dataProductId' parameter: %v", paramsMap["dataProductId"]), nil) } filter, ok := paramsMap["filter"].(string) if !ok { @@ -136,13 +140,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para return nil, util.NewAgentError(fmt.Sprintf("error casting 'orderBy' parameter: %v", paramsMap["orderBy"]), nil) } - parts := strings.Split(name, "/") - if len(parts) < 6 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" { - err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}") - return nil, util.NewAgentError(err.Error(), err) - } - - resp, err := source.ListDataAssets(ctx, name, filter, pageSize, orderBy) + resp, err := source.ListDataAssets(ctx, locationId, dataProductId, filter, pageSize, orderBy) if err != nil { return nil, util.ProcessGcpError(err) } diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 79152ceb3cf3..d7e84d0d96e4 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -668,7 +668,7 @@ func runDataplexToolGetTest(t *testing.T) { { name: "get my-dataplex-list-data-assets-tool", toolName: "my-dataplex-list-data-assets-tool", - expectedParams: []string{"name", "filter", "pageSize", "orderBy"}, + expectedParams: []string{"locationId", "dataProductId", "filter", "pageSize", "orderBy"}, }, } @@ -1679,43 +1679,44 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, t.Fatalf("error getting Google ID token: %s", err) } - parent := fmt.Sprintf("projects/%s/locations/us/dataProducts/%s", DataplexProject, dataProductId) - testCases := []struct { - name string - api string - requestHeader map[string]string - requestBody io.Reader - wantStatusCode int - expectResult bool - wantContentKey string - wantValue string + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantLocationID string + wantDataProductID string + wantDataAsset string }{ { - name: "Success - List Data Assets (Authorized)", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": idToken}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fmt.Sprintf("%s/dataAssets/%s", parent, dataAssetId), + name: "Success - List Data Assets (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us", + wantDataProductID: dataProductId, + wantDataAsset: dataAssetId, }, { - name: "Success - List Data Assets (Un-authorized)", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fmt.Sprintf("%s/dataAssets/%s", parent, dataAssetId), + name: "Success - List Data Assets (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us", + wantDataProductID: dataProductId, + wantDataAsset: dataAssetId, }, { name: "Failure - Invalid Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), wantStatusCode: 401, expectResult: false, }, @@ -1723,18 +1724,10 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, name: "Failure - Without Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", parent))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), wantStatusCode: 401, expectResult: false, }, - { - name: "Failure - Invalid Name Format", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataProductId))), - wantStatusCode: 200, - expectResult: false, - }, } for _, tc := range testCases { @@ -1778,18 +1771,29 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, if !ok { t.Fatalf("expected entry to be a map, got %T", entries[0]) } - val, ok := entry[tc.wantContentKey].(string) + locID, ok := entry["locationId"].(string) if !ok { - t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + t.Fatalf("expected entry to have key 'locationId' as string, but it was not found or not a string in %v", entry) + } + if tc.wantLocationID != "" && locID != tc.wantLocationID { + t.Fatalf("expected locationId to be %q, got %q", tc.wantLocationID, locID) + } + prodID, ok := entry["dataProductId"].(string) + if !ok { + t.Fatalf("expected entry to have key 'dataProductId' as string, but it was not found or not a string in %v", entry) } - if tc.wantValue != "" && val != tc.wantValue { - t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + if tc.wantDataProductID != "" && prodID != tc.wantDataProductID { + t.Fatalf("expected dataProductId to be %q, got %q", tc.wantDataProductID, prodID) + } + assetID, ok := entry["dataAsset"].(string) + if !ok { + t.Fatalf("expected entry to have key 'dataAsset' as string, but it was not found or not a string in %v", entry) + } + if tc.wantDataAsset != "" && assetID != tc.wantDataAsset { + t.Fatalf("expected dataAsset to be %q, got %q", tc.wantDataAsset, assetID) } // Assert output is cleaned - if entry["resource"] == "" { - t.Errorf("resource should not be empty") - } if _, ok := entry["uid"]; ok { t.Errorf("expected entry to NOT have 'uid' field, but it was found") } From 9fb3726de4206ad4b75ffc6b65b69d232f26fcae Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 17 Jun 2026 09:48:01 +0000 Subject: [PATCH 17/25] Update dataplex-get-data-asset tool API --- .../integrations/knowledge-catalog/source.md | 2 +- .../tools/knowledge-catalog-get-data-asset.md | 6 +- internal/sources/dataplex/dataplex.go | 35 ++++-- .../dataplexgetdataasset.go | 28 +++-- tests/dataplex/dataplex_integration_test.go | 115 +++++++++--------- 5 files changed, 104 insertions(+), 82 deletions(-) diff --git a/docs/en/integrations/knowledge-catalog/source.md b/docs/en/integrations/knowledge-catalog/source.md index f3e13bf0cc35..6d50d4409e88 100644 --- a/docs/en/integrations/knowledge-catalog/source.md +++ b/docs/en/integrations/knowledge-catalog/source.md @@ -400,7 +400,7 @@ This abbreviated syntax works for the qualified predicates except for `label` in ## Tool: get_data_asset ### Request 1. Use this tool to retrieve detailed metadata for a specific Data Asset. -2. You must provide the full resource name of the Data Asset in the format `projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}`. +2. You must provide `locationId`, `dataProductId`, and `dataAssetId`. ### Response 1. Present the retrieved metadata for the Data Asset, including its name, resource, labels, and access group configurations. ``` diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md index 0c49088f8f35..e30e38cf00fd 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-data-asset.md @@ -12,9 +12,11 @@ aliases: A `dataplex-get-data-asset` tool retrieves detailed metadata for a specific Data Asset in Knowledge Catalog (formerly known as Dataplex). -`dataplex-get-data-asset` requires the following parameter: +`dataplex-get-data-asset` requires the following parameters: -- `name` - The resource name of the Data Asset in the following form: `projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}`. +- `locationId` - The location ID (e.g., `us`, `us-central1`) where the Data Product is located. +- `dataProductId` - The unique ID of the parent Data Product. +- `dataAssetId` - The unique ID of the Data Asset. ## Compatible Sources diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 413e976fb54a..12881a538f86 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -378,8 +378,8 @@ type AccessGroup struct { ID string `json:"id"` DisplayName string `json:"displayName"` Description string `json:"description"` - GoogleGroup string `json:"googleGroup,omitempty"` - ServiceAccount string `json:"serviceAccount,omitempty"` + GoogleGroup string `json:"googleGroup"` + ServiceAccount string `json:"serviceAccount"` } type DataProduct struct { @@ -439,8 +439,8 @@ func (s *Source) GetDataProduct(ctx context.Context, locationId string, dataProd type DataAssetSummary struct { LocationID string `json:"locationId"` DataProductID string `json:"dataProductId"` - DataAsset string `json:"dataAsset"` - ResourceUri string `json:"resourceUri"` + DataAssetID string `json:"dataAssetId"` + ResourceURI string `json:"resourceUri"` Labels map[string]string `json:"labels"` } @@ -490,8 +490,8 @@ func (s *Source) ListDataAssets( results = append(results, &DataAssetSummary{ LocationID: locId, DataProductID: prodId, - DataAsset: assetId, - ResourceUri: asset.GetResource(), + DataAssetID: assetId, + ResourceURI: asset.GetResource(), Labels: asset.GetLabels(), }) } @@ -499,16 +499,19 @@ func (s *Source) ListDataAssets( } type DataAsset struct { - Name string `json:"name"` - Resource string `json:"resource"` + LocationID string `json:"locationId"` + DataProductID string `json:"dataProductId"` + DataAssetID string `json:"dataAssetId"` + ResourceURI string `json:"resourceUri"` Labels map[string]string `json:"labels"` AccessGroupConfigs map[string]*dataplexpb.DataAsset_AccessGroupConfig `json:"accessGroupConfigs"` } -func (s *Source) GetDataAsset(ctx context.Context, name string) (*DataAsset, error) { +func (s *Source) GetDataAsset(ctx context.Context, locationId string, dataProductId string, dataAssetId string) (*DataAsset, error) { if s.GetDataProductClient() == nil { return nil, fmt.Errorf("dataplex data product client is not initialized") } + name := fmt.Sprintf("projects/%s/locations/%s/dataProducts/%s/dataAssets/%s", s.ProjectID(), locationId, dataProductId, dataAssetId) req := &dataplexpb.GetDataAssetRequest{ Name: name, } @@ -517,9 +520,19 @@ func (s *Source) GetDataAsset(ctx context.Context, name string) (*DataAsset, err return nil, err } + parts := strings.Split(resp.GetName(), "/") + var locId, prodId, assetId string + if len(parts) >= 8 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" && parts[6] == "dataAssets" { + locId = parts[3] + prodId = parts[5] + assetId = parts[7] + } + return &DataAsset{ - Name: resp.GetName(), - Resource: resp.GetResource(), + LocationID: locId, + DataProductID: prodId, + DataAssetID: assetId, + ResourceURI: resp.GetResource(), Labels: resp.GetLabels(), AccessGroupConfigs: resp.GetAccessGroupConfigs(), }, nil diff --git a/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go index 06ad46f4319b..a7d99aa4b46a 100644 --- a/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go +++ b/internal/tools/dataplex/dataplexgetdataasset/dataplexgetdataasset.go @@ -18,7 +18,6 @@ import ( "context" "fmt" "net/http" - "strings" "github.com/goccy/go-yaml" "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" @@ -46,7 +45,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T } type compatibleSource interface { - GetDataAsset(ctx context.Context, name string) (*dataplex.DataAsset, error) + GetDataAsset(ctx context.Context, locationId string, dataProductId string, dataAssetId string) (*dataplex.DataAsset, error) } type Config struct { @@ -68,8 +67,10 @@ func (cfg Config) ToolConfigType() string { } func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { - name := parameters.NewStringParameter("name", "Required. The resource name of the Data Asset in the following form: projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}.") - params := parameters.Parameters{name} + locationId := parameters.NewStringParameter("locationId", "Required. The location ID (e.g., 'us', 'us-central1') where the Data Product is located.") + dataProductId := parameters.NewStringParameter("dataProductId", "Required. The unique ID of the parent Data Product.") + dataAssetId := parameters.NewStringParameter("dataAssetId", "Required. The unique ID of the Data Asset.") + params := parameters.Parameters{locationId, dataProductId, dataAssetId} t := Tool{ Config: cfg, @@ -116,22 +117,23 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para } paramsMap := params.AsMap() - name, ok := paramsMap["name"].(string) + locationId, ok := paramsMap["locationId"].(string) if !ok { - return nil, util.NewAgentError(fmt.Sprintf("error casting 'name' parameter: %v", paramsMap["name"]), nil) + return nil, util.NewAgentError(fmt.Sprintf("error casting 'locationId' parameter: %v", paramsMap["locationId"]), nil) } - - parts := strings.Split(name, "/") - if len(parts) < 8 || parts[0] != "projects" || parts[2] != "locations" || parts[4] != "dataProducts" || parts[6] != "dataAssets" { - err = fmt.Errorf("invalid name format: must be in the form projects/{project}/locations/{location}/dataProducts/{dataProduct}/dataAssets/{dataAsset}") - return nil, util.NewAgentError(err.Error(), err) + dataProductId, ok := paramsMap["dataProductId"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'dataProductId' parameter: %v", paramsMap["dataProductId"]), nil) + } + dataAssetId, ok := paramsMap["dataAssetId"].(string) + if !ok { + return nil, util.NewAgentError(fmt.Sprintf("error casting 'dataAssetId' parameter: %v", paramsMap["dataAssetId"]), nil) } - resp, err := source.GetDataAsset(ctx, name) + resp, err := source.GetDataAsset(ctx, locationId, dataProductId, dataAssetId) if err != nil { return nil, util.ProcessGcpError(err) } - return resp, nil } diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 896feb37b526..a4beb22d5981 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -674,7 +674,7 @@ func runDataplexToolGetTest(t *testing.T) { { name: "get my-dataplex-get-data-asset-tool", toolName: "my-dataplex-get-data-asset-tool", - expectedParams: []string{"name"}, + expectedParams: []string{"locationId", "dataProductId", "dataAssetId"}, }, } @@ -1673,35 +1673,35 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, expectResult bool wantLocationID string wantDataProductID string - wantDataAsset string + wantDataAssetID string }{ { name: "Success - List Data Assets (Authorized)", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": idToken}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 200, expectResult: true, wantLocationID: "us", wantDataProductID: dataProductId, - wantDataAsset: dataAssetId, + wantDataAssetID: dataAssetId, }, { name: "Success - List Data Assets (Un-authorized)", api: "http://127.0.0.1:5000/api/tool/my-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 200, expectResult: true, wantLocationID: "us", wantDataProductID: dataProductId, - wantDataAsset: dataAssetId, + wantDataAssetID: dataAssetId, }, { name: "Failure - Invalid Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 401, expectResult: false, }, @@ -1709,7 +1709,7 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, name: "Failure - Without Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-list-data-assets-tool/invoke", requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\"}", dataProductId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 401, expectResult: false, }, @@ -1770,12 +1770,12 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string, if tc.wantDataProductID != "" && prodID != tc.wantDataProductID { t.Fatalf("expected dataProductId to be %q, got %q", tc.wantDataProductID, prodID) } - assetID, ok := entry["dataAsset"].(string) + assetID, ok := entry["dataAssetId"].(string) if !ok { - t.Fatalf("expected entry to have key 'dataAsset' as string, but it was not found or not a string in %v", entry) + t.Fatalf("expected entry to have key 'dataAssetId' as string, but it was not found or not a string in %v", entry) } - if tc.wantDataAsset != "" && assetID != tc.wantDataAsset { - t.Fatalf("expected dataAsset to be %q, got %q", tc.wantDataAsset, assetID) + if tc.wantDataAssetID != "" && assetID != tc.wantDataAssetID { + t.Fatalf("expected dataAssetId to be %q, got %q", tc.wantDataAssetID, assetID) } // Assert output is cleaned @@ -1798,43 +1798,44 @@ func runDataplexGetDataAssetToolInvokeTest(t *testing.T, dataProductId string, d t.Fatalf("error getting Google ID token: %s", err) } - fullDataAssetId := fmt.Sprintf("projects/%s/locations/us/dataProducts/%s/dataAssets/%s", DataplexProject, dataProductId, dataAssetId) - testCases := []struct { - name string - api string - requestHeader map[string]string - requestBody io.Reader - wantStatusCode int - expectResult bool - wantContentKey string - wantValue string + name string + api string + requestHeader map[string]string + requestBody io.Reader + wantStatusCode int + expectResult bool + wantLocationID string + wantDataProductID string + wantDataAssetID string }{ { - name: "Success - Get Data Asset (Authorized)", - api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", - requestHeader: map[string]string{"my-google-auth_token": idToken}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fullDataAssetId, + name: "Success - Get Data Asset (Authorized)", + api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{"my-google-auth_token": idToken}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us", + wantDataProductID: dataProductId, + wantDataAssetID: dataAssetId, }, { - name: "Success - Get Data Asset (Un-authorized)", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-asset-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), - wantStatusCode: 200, - expectResult: true, - wantContentKey: "name", - wantValue: fullDataAssetId, + name: "Success - Get Data Asset (Un-authorized)", + api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-asset-tool/invoke", + requestHeader: map[string]string{}, + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), + wantStatusCode: 200, + expectResult: true, + wantLocationID: "us", + wantDataProductID: dataProductId, + wantDataAssetID: dataAssetId, }, { name: "Failure - Invalid Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": "invalid_token"}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 401, expectResult: false, }, @@ -1842,18 +1843,10 @@ func runDataplexGetDataAssetToolInvokeTest(t *testing.T, dataProductId string, d name: "Failure - Without Authorization Token", api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-get-data-asset-tool/invoke", requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"%s\"}", fullDataAssetId))), + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"locationId\":\"us\",\"dataProductId\":\"%s\",\"dataAssetId\":\"%s\"}", dataProductId, dataAssetId))), wantStatusCode: 401, expectResult: false, }, - { - name: "Failure - Invalid Name Format", - api: "http://127.0.0.1:5000/api/tool/my-dataplex-get-data-asset-tool/invoke", - requestHeader: map[string]string{}, - requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"invalid-name-%s\"}", dataAssetId))), - wantStatusCode: 200, - expectResult: false, - }, } for _, tc := range testCases { @@ -1889,18 +1882,30 @@ func runDataplexGetDataAssetToolInvokeTest(t *testing.T, dataProductId string, d if err := json.Unmarshal([]byte(resultStr), &entry); err != nil { t.Fatalf("error unmarshalling result string: %v", err) } - val, ok := entry[tc.wantContentKey].(string) + + locID, ok := entry["locationId"].(string) + if !ok { + t.Fatalf("expected entry to have key 'locationId' as string, but it was not found or not a string in %v", entry) + } + if tc.wantLocationID != "" && locID != tc.wantLocationID { + t.Fatalf("expected locationId to be %q, got %q", tc.wantLocationID, locID) + } + prodID, ok := entry["dataProductId"].(string) if !ok { - t.Fatalf("expected entry to have key '%s' as string, but it was not found or not a string in %v", tc.wantContentKey, entry) + t.Fatalf("expected entry to have key 'dataProductId' as string, but it was not found or not a string in %v", entry) } - if tc.wantValue != "" && val != tc.wantValue { - t.Fatalf("expected entry %s to be %q, got %q", tc.wantContentKey, tc.wantValue, val) + if tc.wantDataProductID != "" && prodID != tc.wantDataProductID { + t.Fatalf("expected dataProductId to be %q, got %q", tc.wantDataProductID, prodID) + } + assetID, ok := entry["dataAssetId"].(string) + if !ok { + t.Fatalf("expected entry to have key 'dataAssetId' as string, but it was not found or not a string in %v", entry) + } + if tc.wantDataAssetID != "" && assetID != tc.wantDataAssetID { + t.Fatalf("expected dataAssetId to be %q, got %q", tc.wantDataAssetID, assetID) } // Assert output is cleaned - if entry["resource"] == "" { - t.Errorf("resource should not be empty") - } if _, ok := entry["uid"]; ok { t.Errorf("expected entry to NOT have 'uid' field, but it was found") } From 59658a7c2181ec8fb8ec5c8822581324f6b032d1 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Thu, 18 Jun 2026 05:20:11 +0000 Subject: [PATCH 18/25] Update Dataplex data-products toolset --- internal/prebuiltconfigs/tools/dataplex.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/internal/prebuiltconfigs/tools/dataplex.yaml b/internal/prebuiltconfigs/tools/dataplex.yaml index 0b33511b61cc..0dc437afcf3a 100644 --- a/internal/prebuiltconfigs/tools/dataplex.yaml +++ b/internal/prebuiltconfigs/tools/dataplex.yaml @@ -240,6 +240,10 @@ tools: kind: toolset name: data-products tools: +- search_entries +- lookup_entry +- search_aspect_types +- lookup_context - list_data_products --- kind: toolset From 469590fb5562e8496ba2bd55310bafe22dc6165e Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Thu, 18 Jun 2026 05:21:31 +0000 Subject: [PATCH 19/25] Reference Data Products usage guide in dataplex-list-data-products tool doc --- .../tools/knowledge-catalog-list-data-products.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md index 3013b5083940..e8ee6507ab44 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md @@ -12,12 +12,16 @@ aliases: A `dataplex-list-data-products` tool lists all Data Products in Knowledge Catalog (formerly known as Dataplex) across all locations (globally). +View the [Data Products usage guide][usage-guide] for more information. + `dataplex-list-data-products` optionally accepts the following parameters: - `filter` - Filter string to list data products. Use `=` for exact matching and `:` for contains matching. String literals must be enclosed within double quotes. E.g. `display_name:"my-product"`. - `pageSize` - Number of returned data products in the page. Defaults to `10`. - `orderBy` - Specifies the ordering of results. +[usage-guide]: https://docs.cloud.google.com/dataplex/docs/use-data-products + ## Compatible Sources {{< compatible-sources >}} From c4b1d5e2ac5a8392687c39d20ec7366c4f4c508c Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Thu, 18 Jun 2026 15:03:18 +0000 Subject: [PATCH 20/25] Fix failing prebuilt configs test --- cmd/internal/config_test.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cmd/internal/config_test.go b/cmd/internal/config_test.go index 63c954eef318..62e64509dc1b 100644 --- a/cmd/internal/config_test.go +++ b/cmd/internal/config_test.go @@ -1858,7 +1858,11 @@ func TestPrebuiltTools(t *testing.T) { wantToolset: server.ToolsetConfigs{ "discovery": tools.ToolsetConfig{ Name: "discovery", - ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans", "list_data_products"}, + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans"}, + }, + "data-products": tools.ToolsetConfig{ + Name: "data-products", + ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "list_data_products"}, }, "enrich": tools.ToolsetConfig{ Name: "enrich", From c2f10b37224f30c2de4710760792364057f034a1 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Thu, 18 Jun 2026 15:33:42 +0000 Subject: [PATCH 21/25] Update dataplex-list-data-products tool doc formatting --- .../knowledge-catalog-list-data-products.md | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md index e8ee6507ab44..a6e4c606623c 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md @@ -14,12 +14,6 @@ A `dataplex-list-data-products` tool lists all Data Products in Knowledge Catalo View the [Data Products usage guide][usage-guide] for more information. -`dataplex-list-data-products` optionally accepts the following parameters: - -- `filter` - Filter string to list data products. Use `=` for exact matching and `:` for contains matching. String literals must be enclosed within double quotes. E.g. `display_name:"my-product"`. -- `pageSize` - Number of returned data products in the page. Defaults to `10`. -- `orderBy` - Specifies the ordering of results. - [usage-guide]: https://docs.cloud.google.com/dataplex/docs/use-data-products ## Compatible Sources @@ -48,6 +42,16 @@ applying IAM permissions and roles to an identity. [iam-roles]: https://cloud.google.com/dataplex/docs/iam-roles [dataplex-docs]: https://cloud.google.com/dataplex +## Parameters + +The `dataplex-list-data-products` tool has the following optional parameters: + +| **field** | **type** | **required** | **description** | +| --------- | :------: | :----------: | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| filter | string | false | Filter string to list data products. Use `=` for exact matching and `:` for contains matching. String literals must be enclosed within double quotes. E.g. `display_name:"my-product"`. | +| pageSize | integer | false | Number of returned data products in the page. Defaults to `10`. | +| orderBy | string | false | Specifies the ordering of results. | + ## Example ```yaml From 19866797d701ca434fa8cbee125165d468ec1e23 Mon Sep 17 00:00:00 2001 From: Tejas Singh <88895455+theantagonist9509@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:33:52 +0530 Subject: [PATCH 22/25] Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- internal/sources/dataplex/dataplex.go | 24 +++++++++---------- .../dataplexgetdataproduct.go | 8 +++---- .../dataplexlistdataproducts.go | 2 +- tests/dataplex/dataplex_integration_test.go | 4 ++-- 4 files changed, 19 insertions(+), 19 deletions(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 12b4c556ba18..9b38872ff1ff 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -401,14 +401,14 @@ func (s *Source) ListDataProducts( return nil, fmt.Errorf("failed to list data products: %w", err) } parts := strings.Split(dp.GetName(), "/") - var locId, prodId string + var locID, prodID string if len(parts) >= 6 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" { - locId = parts[3] - prodId = parts[5] + locID = parts[3] + prodID = parts[5] } results = append(results, &DataProductSummary{ - LocationID: locId, - DataProductID: prodId, + LocationID: locID, + DataProductID: prodID, DisplayName: dp.GetDisplayName(), OwnerEmails: dp.GetOwnerEmails(), AssetCount: dp.GetAssetCount(), @@ -436,11 +436,11 @@ type DataProduct struct { AccessGroups []AccessGroup `json:"accessGroups"` } -func (s *Source) GetDataProduct(ctx context.Context, locationId string, dataProductId string) (*DataProduct, error) { +func (s *Source) GetDataProduct(ctx context.Context, locationID string, dataProductID string) (*DataProduct, error) { if s.GetDataProductClient() == nil { return nil, fmt.Errorf("dataplex data product client is not initialized") } - name := fmt.Sprintf("projects/%s/locations/%s/dataProducts/%s", s.ProjectID(), locationId, dataProductId) + name := fmt.Sprintf("projects/%s/locations/%s/dataProducts/%s", s.ProjectID(), locationID, dataProductID) req := &dataplexpb.GetDataProductRequest{ Name: name, } @@ -461,15 +461,15 @@ func (s *Source) GetDataProduct(ctx context.Context, locationId string, dataProd } parts := strings.Split(resp.GetName(), "/") - var locId, prodId string + var locID, prodID string if len(parts) >= 6 && parts[0] == "projects" && parts[2] == "locations" && parts[4] == "dataProducts" { - locId = parts[3] - prodId = parts[5] + locID = parts[3] + prodID = parts[5] } return &DataProduct{ - LocationID: locId, - DataProductID: prodId, + LocationID: locID, + DataProductID: prodID, DisplayName: resp.GetDisplayName(), Description: resp.GetDescription(), OwnerEmails: resp.GetOwnerEmails(), diff --git a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go index c66e252d443b..f1d71ad80cb6 100644 --- a/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go +++ b/internal/tools/dataplex/dataplexgetdataproduct/dataplexgetdataproduct.go @@ -43,7 +43,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T } type compatibleSource interface { - GetDataProduct(ctx context.Context, locationId string, dataProductId string) (*dataplex.DataProduct, error) + GetDataProduct(ctx context.Context, locationID string, dataProductID string) (*dataplex.DataProduct, error) } type Config struct { @@ -97,16 +97,16 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para } paramsMap := params.AsMap() - locationId, ok := paramsMap["locationId"].(string) + locationID, ok := paramsMap["locationId"].(string) if !ok { return nil, util.NewAgentError(fmt.Sprintf("error casting 'locationId' parameter: %v", paramsMap["locationId"]), nil) } - dataProductId, ok := paramsMap["dataProductId"].(string) + dataProductID, ok := paramsMap["dataProductId"].(string) if !ok { return nil, util.NewAgentError(fmt.Sprintf("error casting 'dataProductId' parameter: %v", paramsMap["dataProductId"]), nil) } - resp, err := source.GetDataProduct(ctx, locationId, dataProductId) + resp, err := source.GetDataProduct(ctx, locationID, dataProductID) if err != nil { return nil, util.ProcessGcpError(err) } diff --git a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go index 36a465bae11e..6a4456761e16 100644 --- a/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go +++ b/internal/tools/dataplex/dataplexlistdataproducts/dataplexlistdataproducts.go @@ -61,7 +61,7 @@ func (cfg Config) ToolConfigType() string { } func (cfg Config) Initialize() (tools.Tool, error) { - filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data products. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching accross all fields at once is not yet supported. E.g. \"display_name:\\\"my-product\\\"\"") + filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data products. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching across all fields at once is not yet supported. E.g. \"display_name:\\\"my-product\\\"\"") pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Number of returned data products in the page.") orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Specifies the ordering of results.") params := parameters.Parameters{filter, pageSize, orderBy} diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 4ec48de37ad2..c9783c792301 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -226,7 +226,7 @@ func initDataplexDataScanConnection(ctx context.Context) (*dataplex.DataScanClie } func initDataplexDataProductConnection(ctx context.Context) (*dataplex.DataProductClient, error) { - cred, err := google.FindDefaultCredentials(ctx) + cred, err := google.FindDefaultCredentials(ctx, sources.CloudPlatformScope) if err != nil { return nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err) } @@ -391,7 +391,7 @@ func setupDataplexDataProduct(t *testing.T, ctx context.Context, client *dataple parent := fmt.Sprintf("projects/%s/locations/us-central1", DataplexProject) ownerEmail := tests.ServiceAccountEmail if ownerEmail == "" { - t.Errorf("Service account email is required, but tests.ServiceAccountEmail was empty") + t.Fatalf("Service account email is required, but tests.ServiceAccountEmail was empty") } createReq := &dataplexpb.CreateDataProductRequest{ Parent: parent, From 7cab2ba4256c4ce3d494a1239882050d3f4f6239 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 23 Jun 2026 09:25:25 +0000 Subject: [PATCH 23/25] Revert accidental removal of error check --- internal/sources/dataplex/dataplex.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index 51a503a32eb4..102ed1cd38c8 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -638,6 +638,9 @@ func (s *Source) GetOperation(ctx context.Context, opName string) (map[string]an Name: opName, } op, err := s.DataScanClient.LROClient.GetOperation(ctx, req) + if err != nil { + return nil, err + } bytes, err := protojson.Marshal(op) if err != nil { From 73f81fb41087a993fbeaf1eb32c70bba4a9e4325 Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Tue, 23 Jun 2026 09:37:02 +0000 Subject: [PATCH 24/25] Fix typo in tool parameter annotation --- internal/sources/dataplex/dataplex.go | 6 +++--- .../dataplexlistdataassets/dataplexlistdataassets.go | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/sources/dataplex/dataplex.go b/internal/sources/dataplex/dataplex.go index fdb6169f681d..3678cd6abec5 100644 --- a/internal/sources/dataplex/dataplex.go +++ b/internal/sources/dataplex/dataplex.go @@ -86,7 +86,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So Config: r, Client: client, DataScanClient: dataScanClient, - dataProductClient: dataProductClient, + DataProductClient: dataProductClient, } return s, nil @@ -98,7 +98,7 @@ type Source struct { Config Client *dataplexapi.CatalogClient DataScanClient *dataplexapi.DataScanClient - dataProductClient *dataplexapi.DataProductClient + DataProductClient *dataplexapi.DataProductClient } func (s *Source) SourceType() string { @@ -123,7 +123,7 @@ func (s *Source) GetDataScanClient() *dataplexapi.DataScanClient { } func (s *Source) GetDataProductClient() *dataplexapi.DataProductClient { - return s.dataProductClient + return s.DataProductClient } func initDataplexConnection( diff --git a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go index 9e94d8c98a31..41769251653b 100644 --- a/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go +++ b/internal/tools/dataplex/dataplexlistdataassets/dataplexlistdataassets.go @@ -63,7 +63,7 @@ func (cfg Config) ToolConfigType() string { func (cfg Config) Initialize() (tools.Tool, error) { locationId := parameters.NewStringParameter("locationId", "Required. The location ID (e.g., 'us', 'us-central1') where the Data Product is located.") dataProductId := parameters.NewStringParameter("dataProductId", "Required. The unique ID of the parent Data Product.") - filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data assets. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching accross all fields at once is not yet supported.") + filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to list data assets. Based on the AIP-160 proposal. Use '=' for exact, and ':' for contains matching. String literals must be enclosed within \"\". Matching across all fields at once is not yet supported.") pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Optional. Number of returned data assets in the page.") orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Optional. Specifies the ordering of results.") params := parameters.Parameters{locationId, dataProductId, filter, pageSize, orderBy} From f7c451a4a220fd20a6b31777464dba9c1615d2aa Mon Sep 17 00:00:00 2001 From: Tejas Singh Date: Wed, 24 Jun 2026 03:43:25 +0000 Subject: [PATCH 25/25] Address PR comments --- .../prebuilt-configs/knowledge-catalog.md | 2 ++ .../tools/knowledge-catalog-list-data-products.md | 2 -- tests/dataplex/dataplex_integration_test.go | 10 +++++----- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/en/integrations/knowledge-catalog/prebuilt-configs/knowledge-catalog.md b/docs/en/integrations/knowledge-catalog/prebuilt-configs/knowledge-catalog.md index 57985a93414b..aef0efd1314d 100644 --- a/docs/en/integrations/knowledge-catalog/prebuilt-configs/knowledge-catalog.md +++ b/docs/en/integrations/knowledge-catalog/prebuilt-configs/knowledge-catalog.md @@ -21,6 +21,7 @@ aliases: * `search_aspect_types`: Finds aspect types relevant to the query. * `lookup_context`: Retrieves rich metadata regarding one or more data assets along with their relationships. * `search_dq_scans`: Search for data quality scans in Dataplex. + * `list_data_products`: Lists Data Products across all locations. * `generate_data_insights`: Creates a new Dataplex Data Documentation scan template and triggers the run. * `get_data_insights`: Retrieves the final generated data insights for a completed scan. * `generate_data_profile`: Creates a new Dataplex Data Profile scan template and triggers the run. @@ -33,5 +34,6 @@ aliases: * `get_run_status`: Retrieves the execution status of the latest background job run. * **Toolsets:** * `discovery`: Metadata discovery and search toolset (`search_entries`, `lookup_entry`, `search_aspect_types`, `lookup_context`, `search_dq_scans`). + * `data-products`: Data Products and Data Assets curation and management toolset (`search_entries`, `lookup_entry`, `search_aspect_types`, `lookup_context`, `list_data_products`). * `enrich`: Metadata enrichment pipeline orchestration and execution toolset (`search_entries`, `lookup_entry`, `lookup_context`, `generate_data_insights`, `get_data_insights`, `generate_data_profile`, `get_data_profile`, `discover_metadata`, `get_discovery_results`, `check_data_quality`, `get_data_quality_results`, `get_operation`, `get_run_status`). diff --git a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md index a6e4c606623c..572a268fc691 100644 --- a/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md +++ b/docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-list-data-products.md @@ -4,8 +4,6 @@ type: docs weight: 1 description: > A "dataplex-list-data-products" tool allows to list data products. -aliases: - - /integrations/dataplex/tools/dataplex-list-data-products/ --- ## About diff --git a/tests/dataplex/dataplex_integration_test.go b/tests/dataplex/dataplex_integration_test.go index 567c97a7f654..b0c09cd92ac5 100644 --- a/tests/dataplex/dataplex_integration_test.go +++ b/tests/dataplex/dataplex_integration_test.go @@ -647,6 +647,11 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-search-dq-scans-tool", expectedParams: []string{"filter", "dataScanId", "resourcePath", "pageSize", "orderBy"}, }, + { + name: "get my-dataplex-list-data-products-tool", + toolName: "my-dataplex-list-data-products-tool", + expectedParams: []string{"filter", "pageSize", "orderBy"}, + }, { name: "get my-dataplex-generate-data-profile-tool", toolName: "my-dataplex-generate-data-profile-tool", @@ -697,11 +702,6 @@ func runDataplexToolGetTest(t *testing.T) { toolName: "my-dataplex-get-data-quality-results-tool", expectedParams: []string{"scanId", "location"}, }, - { - name: "get my-dataplex-list-data-products-tool", - toolName: "my-dataplex-list-data-products-tool", - expectedParams: []string{"filter", "pageSize", "orderBy"}, - }, } for _, tc := range testCases {