feat(tools/dataplex-search-aspect-types): Add support for dataplex-search-aspect-types tool (#1061)

Added support for search aspect types tool in Dataplex.
Fixes #1056

---------

Co-authored-by: Averi Kitsch <akitsch@google.com>
This commit is contained in:
Anuj Jhunjhunwala
2025-08-14 22:34:00 +02:00
committed by GitHub
parent 31ed87861d
commit d940187c85
12 changed files with 581 additions and 52 deletions

View File

@@ -52,6 +52,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/bigtable"
_ "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoredeletedocuments"

View File

@@ -1250,7 +1250,7 @@ func TestPrebuiltTools(t *testing.T) {
wantToolset: server.ToolsetConfigs{
"dataplex-tools": tools.ToolsetConfig{
Name: "dataplex-tools",
ToolNames: []string{"dataplex_search_entries", "dataplex_lookup_entry"},
ToolNames: []string{"dataplex_search_entries", "dataplex_lookup_entry", "dataplex_search_aspect_types"},
},
},
},

View File

@@ -218,22 +218,45 @@ Search syntax supports the following qualifiers:
- "label.foo" - Matches BigQuery resources that have a label whose key equals foo as a string.
- "type=TYPE" - Matches resources of a specific entry type or its type alias.
- "projectid:bar" - Matches resources within Google Cloud projects that match bar as a substring in the ID.
- "parent:x" - Matches x as a substring of the hierarchical path of a resource. The parent path is a fully_qualified_name of the parent resource.
- "parent:x" - Matches x as a substring of the hierarchical path of a resource. It supports same syntax as `name` predicate.
- "orgid=number" - Matches resources within a Google Cloud organization with the exact ID value of the number.
- "system=SYSTEM" - Matches resources from a specified system. For example, system=bigquery matches BigQuery resources.
- "location=LOCATION" - Matches resources in a specified location with an exact name. For example, location=us-central1 matches assets hosted in Iowa. BigQuery Omni assets support this qualifier by using the BigQuery Omni location name. For example, location=aws-us-east-1 matches BigQuery Omni assets in Northern Virginia.
- "createtime" -
Finds resources that were created within, before, or after a given date or time. For example "createtime:2019-01-01" matches resources created on 2019-01-01.
- "updatetime" - Finds resources that were updated within, before, or after a given date or time. For example "updatetime>2019-01-01" matches resources updated after 2019-01-01.
- "fully_qualified_name:x" - Matches x as a substring of fully_qualified_name.
- "fully_qualified_name=x" - Matches x as fully_qualified_name.
### Aspect Search
To search for entries based on their attached aspects, use the following query syntax.
aspect:x Matches x as a substring of the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
aspect=x Matches x as the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
aspect:xOPERATORvalue
Searches for aspect field values. Matches x as a substring of the full path to the aspect type and field name of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID.FIELD_NAME
The list of supported {OPERATOR}s depends on the type of field in the aspect, as follows:
- String: = (exact match) and : (substring)
- All number types: =, :, <, >, <=, >=, =>, =<
- Enum: =
- Datetime: same as for numbers, but the values to compare are treated as datetimes instead of numbers
- Boolean: =
Only top-level fields of the aspect are searchable. For example, all of the following queries match entries where the value of the is-enrolled field in the employee-info aspect type is true. Other entries that match on the substring are also returned.
- aspect:example-project.us-central1.employee-info.is-enrolled=true
- aspect:example-project.us-central1.employee=true
- aspect:employee=true
Example:-
You can use following filters
- dataplex-types.global.bigquery-table.type={BIGLAKE_TABLE, BIGLAKE_OBJECT_TABLE, EXTERNAL_TABLE, TABLE}
- dataplex-types.global.storage.type={STRUCTURED, UNSTRUCTURED}
### Logical operators
A query can consist of several predicates with logical operators. If you don't specify an operator, logical AND is implied. For example, foo bar returns resources that match both predicate foo and predicate bar.
Logical AND and logical OR are supported. For example, foo OR bar.
You can negate a predicate with a - (hyphen) or NOT prefix. For example, -name:foo returns resources with names that don't match the predicate foo.
Logical operators aren't case-sensitive. For example, both or and OR are acceptable.
Logical operators are case-sensitive. `OR` and `AND` are acceptable whereas `or` and `and` are not.
### Request
1. Always try to rewrite the prompt using search syntax.
@@ -287,7 +310,7 @@ Logical operators aren't case-sensitive. For example, both or and OR are accepta
## Tool: dataplex_lookup_entry
### Request
1. Always try to limit the size of the response by specifying `aspect_types` parameter. Make sure to include to select view=CUSTOM when using aspect_types parameter.
1. Always try to limit the size of the response by specifying `aspect_types` parameter. Make sure to include to select view=CUSTOM when using aspect_types parameter. If you do not know the name of the aspect type, use the `dataplex_search_aspect_types` tool.
2. If you do not know the name of the entry, use `dataplex_search_entries` tool
### Response
1. Unless asked for a specific aspect, respond with all aspects attached to the entry.
@@ -298,4 +321,4 @@ Logical operators aren't case-sensitive. For example, both or and OR are accepta
| **field** | **type** | **required** | **description** |
|-----------|:--------:|:------------:|----------------------------------------------------------------------------------|
| kind | string | true | Must be "dataplex". |
| project | string | true | ID of the GCP project used for quota and billing purposes (e.g. "my-project-id").|
| project | string | true | ID of the GCP project used for quota and billing purposes (e.g. "my-project-id").|

View File

@@ -0,0 +1,62 @@
---
title: "dataplex-search-aspect-types"
type: docs
weight: 1
description: >
A "dataplex-search-aspect-types" tool allows to to find aspect types relevant to the query.
aliases:
- /resources/tools/dataplex-search-aspect-types
---
## About
A `dataplex-search-aspect-types` tool allows to fetch the metadata template of aspect types based on search query.
It's compatible with the following sources:
- [dataplex](../../sources/dataplex.md)
`dataplex-search-aspect-types` accepts following parameters optionally:
- `query` - Narrows down the search of aspect types to value of this parameter. If not provided, it fetches all aspect types available to the user.
- `pageSize` - Number of returned aspect types in the search page. Defaults to `5`.
- `orderBy` - Specifies the ordering of results. Supported values are: relevance (default), last_modified_timestamp, last_modified_timestamp asc.
## Requirements
### IAM Permissions
Dataplex uses [Identity and Access Management (IAM)][iam-overview] to control
user and group access to Dataplex resources. Toolbox will use your
[Application Default Credentials (ADC)][adc] to authorize and authenticate when
interacting with [Dataplex][dataplex-docs].
In addition to [setting the ADC for your server][set-adc], you need to ensure
the IAM identity has been given the correct IAM permissions for the tasks you
intend to perform. See [Dataplex Universal Catalog IAM permissions][iam-permissions]
and [Dataplex Universal Catalog IAM roles][iam-roles] for more information on
applying IAM permissions and roles to an identity.
[iam-overview]: https://cloud.google.com/dataplex/docs/iam-and-access-control
[adc]: https://cloud.google.com/docs/authentication#adc
[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc
[iam-permissions]: https://cloud.google.com/dataplex/docs/iam-permissions
[iam-roles]: https://cloud.google.com/dataplex/docs/iam-roles
[dataplex-docs]: https://cloud.google.com/dataplex
## Example
```yaml
tools:
dataplex-search-aspect-types:
kind: dataplex-search-aspect-types
source: my-dataplex-source
description: Use this tool to find aspect types relevant to the query.
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "dataplex-search-aspect-types". |
| 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. |

View File

@@ -17,22 +17,11 @@ It's compatible with the following sources:
- [dataplex](../../sources/dataplex.md)
`dataplex-search-entries` takes a required `query` parameter based on which
entries are filtered and returned to the user and a required `name` parameter
which is constructed using source's project if user does not provide it
explicitly and has the following format: projects/{project}/locations/global. It
also optionally accepts following parameters:
entries are filtered and returned to the user. It also optionally accepts following parameters:
- `pageSize` - Number of results in the search page. Defaults to `5`.
- `pageToken` - Page token received from a previous locations.searchEntries
call.
- `orderBy` - Specifies the ordering of results. Supported values are: relevance
(default), last_modified_timestamp, last_modified_timestamp asc
- `semanticSearch` - Specifies whether the search should understand the meaning
and intent behind the query, rather than just matching keywords. Defaults to
`true`.
- `scope` - The scope under which the search should be operating. Since this
parameter is not exposed to the toolbox user, it defaults to the organization
where the project provided in name is located.
(default), last_modified_timestamp, last_modified_timestamp asc.
## Requirements

2
go.mod
View File

@@ -14,6 +14,7 @@ require (
cloud.google.com/go/spanner v1.84.1
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.29.0
github.com/cenkalti/backoff/v5 v5.0.2
github.com/couchbase/gocb/v2 v2.10.1
github.com/couchbase/tools-common/http v1.0.9
github.com/fsnotify/fsnotify v1.9.0
@@ -73,7 +74,6 @@ require (
github.com/PuerkitoBio/goquery v1.10.3 // indirect
github.com/ajg/form v1.5.1 // indirect
github.com/apache/arrow/go/v15 v15.0.2 // indirect
github.com/cenkalti/backoff/v5 v5.0.2 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect
github.com/couchbase/gocbcore/v10 v10.7.1 // indirect

View File

@@ -12,8 +12,13 @@ tools:
kind: dataplex-lookup-entry
source: dataplex-source
description: Use this tool to retrieve a specific entry from Dataplex Catalog.
dataplex_search_aspect_types:
kind: dataplex-search-aspect-types
source: dataplex-source
description: Use this tool to find aspect types relevant to the query.
toolsets:
dataplex-tools:
- dataplex_search_entries
- dataplex_lookup_entry
- dataplex_lookup_entry
- dataplex_search_aspect_types

View File

@@ -106,7 +106,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
InputSchema: parameters.McpManifest(),
}
t := &Tool{
t := Tool{
Name: cfg.Name,
Kind: kind,
Parameters: parameters,
@@ -132,11 +132,11 @@ type Tool struct {
mcpManifest tools.McpManifest
}
func (t *Tool) Authorized(verifiedAuthServices []string) bool {
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
paramsMap := params.AsMap()
viewMap := map[int]dataplexpb.EntryView{
1: dataplexpb.EntryView_BASIC,
@@ -167,17 +167,17 @@ func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error
return result, nil
}
func (t *Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
// Parse parameters from the provided data
return tools.ParseParams(t.Parameters, data, claims)
}
func (t *Tool) Manifest() tools.Manifest {
func (t Tool) Manifest() tools.Manifest {
// Returns the tool manifest
return t.manifest
}
func (t *Tool) McpManifest() tools.McpManifest {
func (t Tool) McpManifest() tools.McpManifest {
// Returns the tool MCP manifest
return t.mcpManifest
}

View File

@@ -0,0 +1,196 @@
// Copyright 2025 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 dataplexsearchaspecttypes
import (
"context"
"fmt"
dataplexapi "cloud.google.com/go/dataplex/apiv1"
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
"github.com/cenkalti/backoff/v5"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
dataplexds "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
"github.com/googleapis/genai-toolbox/internal/tools"
)
const kind string = "dataplex-search-aspect-types"
func init() {
if !tools.Register(kind, newConfig) {
panic(fmt.Sprintf("tool kind %q already registered", kind))
}
}
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 {
CatalogClient() *dataplexapi.CatalogClient
ProjectID() string
}
// validate compatible sources are still compatible
var _ compatibleSource = &dataplexds.Source{}
var compatibleSources = [...]string{dataplexds.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description"`
AuthRequired []string `yaml:"authRequired"`
}
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// Initialize the search configuration with the provided sources
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is compatible
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
query := tools.NewStringParameter("query", "The query against which aspect type should be matched.")
pageSize := tools.NewIntParameterWithDefault("pageSize", 5, "Number of returned aspect types in the search page.")
orderBy := tools.NewStringParameterWithDefault("orderBy", "relevance", "Specifies the ordering of results. Supported values are: relevance, last_modified_timestamp, last_modified_timestamp asc")
parameters := tools.Parameters{query, pageSize, orderBy}
mcpManifest := tools.McpManifest{
Name: cfg.Name,
Description: cfg.Description,
InputSchema: parameters.McpManifest(),
}
t := Tool{
Name: cfg.Name,
Kind: kind,
Parameters: parameters,
AuthRequired: cfg.AuthRequired,
CatalogClient: s.CatalogClient(),
ProjectID: s.ProjectID(),
manifest: tools.Manifest{
Description: cfg.Description,
Parameters: parameters.Manifest(),
AuthRequired: cfg.AuthRequired,
},
mcpManifest: mcpManifest,
}
return t, nil
}
type Tool struct {
Name string
Kind string
Parameters tools.Parameters
AuthRequired []string
CatalogClient *dataplexapi.CatalogClient
ProjectID string
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
// Invoke the tool with the provided parameters
paramsMap := params.AsMap()
query, _ := paramsMap["query"].(string)
pageSize := int32(paramsMap["pageSize"].(int))
orderBy, _ := paramsMap["orderBy"].(string)
// Create SearchEntriesRequest with the provided parameters
req := &dataplexpb.SearchEntriesRequest{
Query: query + " type=projects/dataplex-types/locations/global/entryTypes/aspecttype",
Name: fmt.Sprintf("projects/%s/locations/global", t.ProjectID),
PageSize: pageSize,
OrderBy: orderBy,
SemanticSearch: true,
}
// Perform the search using the CatalogClient - this will return an iterator
it := t.CatalogClient.SearchEntries(ctx, req)
if it == nil {
return nil, fmt.Errorf("failed to create search entries iterator for project %q", t.ProjectID)
}
// Create an instance of exponential backoff with default values for retrying GetAspectType calls
// InitialInterval, RandomizationFactor, Multiplier, MaxInterval = 500 ms, 0.5, 1.5, 60 s
getAspectBackOff := backoff.NewExponentialBackOff()
// Iterate through the search results and call GetAspectType for each result using the resource name
var results []*dataplexpb.AspectType
for {
entry, err := it.Next()
if err != nil {
break
}
resourceName := entry.DataplexEntry.GetEntrySource().Resource
getAspectTypeReq := &dataplexpb.GetAspectTypeRequest{
Name: resourceName,
}
operation := func() (*dataplexpb.AspectType, error) {
aspectType, err := t.CatalogClient.GetAspectType(ctx, getAspectTypeReq)
if err != nil {
return nil, fmt.Errorf("failed to get aspect type for entry %q: %w", resourceName, err)
}
return aspectType, nil
}
// Retry the GetAspectType operation with exponential backoff
aspectType, err := backoff.Retry(ctx, operation, backoff.WithBackOff(getAspectBackOff))
if err != nil {
return nil, fmt.Errorf("failed to get aspect type after retries for entry %q: %w", resourceName, err)
}
results = append(results, aspectType)
}
return results, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
// Parse parameters from the provided data
return tools.ParseParams(t.Parameters, data, claims)
}
func (t Tool) Manifest() tools.Manifest {
// Returns the tool manifest
return t.manifest
}
func (t Tool) McpManifest() tools.McpManifest {
// Returns the tool MCP manifest
return t.mcpManifest
}

View File

@@ -0,0 +1,73 @@
// Copyright 2025 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 dataplexsearchaspecttypes_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
)
func TestParseFromYamlDataplexSearchAspectTypes(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: `
tools:
example_tool:
kind: dataplex-search-aspect-types
source: my-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": dataplexsearchaspecttypes.Config{
Name: "example_tool",
Kind: "dataplex-search-aspect-types",
Source: "my-instance",
Description: "some description",
AuthRequired: []string{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}

View File

@@ -81,10 +81,8 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
query := tools.NewStringParameter("query", "The query against which entries in scope should be matched.")
pageSize := tools.NewIntParameterWithDefault("pageSize", 5, "Number of results in the search page.")
pageToken := tools.NewStringParameterWithDefault("pageToken", "", "Page token received from a previous locations.searchEntries call. Provide this to retrieve the subsequent page.")
orderBy := tools.NewStringParameterWithDefault("orderBy", "relevance", "Specifies the ordering of results. Supported values are: relevance, last_modified_timestamp, last_modified_timestamp asc")
semanticSearch := tools.NewBooleanParameterWithDefault("semanticSearch", true, "Whether to use semantic search for the query. If true, the query will be processed using semantic search capabilities.")
parameters := tools.Parameters{query, pageSize, pageToken, orderBy, semanticSearch}
parameters := tools.Parameters{query, pageSize, orderBy}
mcpManifest := tools.McpManifest{
Name: cfg.Name,
@@ -92,7 +90,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
InputSchema: parameters.McpManifest(),
}
t := &Tool{
t := Tool{
Name: cfg.Name,
Kind: kind,
Parameters: parameters,
@@ -120,25 +118,22 @@ type Tool struct {
mcpManifest tools.McpManifest
}
func (t *Tool) Authorized(verifiedAuthServices []string) bool {
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
paramsMap := params.AsMap()
query, _ := paramsMap["query"].(string)
pageSize, _ := paramsMap["pageSize"].(int32)
pageToken, _ := paramsMap["pageToken"].(string)
pageSize := int32(paramsMap["pageSize"].(int))
orderBy, _ := paramsMap["orderBy"].(string)
semanticSearch, _ := paramsMap["semanticSearch"].(bool)
req := &dataplexpb.SearchEntriesRequest{
Query: query,
Name: fmt.Sprintf("projects/%s/locations/global", t.ProjectID),
PageSize: pageSize,
PageToken: pageToken,
OrderBy: orderBy,
SemanticSearch: semanticSearch,
SemanticSearch: true,
}
it := t.CatalogClient.SearchEntries(ctx, req)
@@ -157,17 +152,17 @@ func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error
return results, nil
}
func (t *Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
// Parse parameters from the provided data
return tools.ParseParams(t.Parameters, data, claims)
}
func (t *Tool) Manifest() tools.Manifest {
func (t Tool) Manifest() tools.Manifest {
// Returns the tool manifest
return t.manifest
}
func (t *Tool) McpManifest() tools.McpManifest {
func (t Tool) McpManifest() tools.McpManifest {
// Returns the tool MCP manifest
return t.mcpManifest
}

View File

@@ -28,6 +28,8 @@ import (
"time"
bigqueryapi "cloud.google.com/go/bigquery"
dataplex "cloud.google.com/go/dataplex/apiv1"
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
@@ -38,10 +40,11 @@ import (
)
var (
DataplexSourceKind = "dataplex"
DataplexSearchEntriesToolKind = "dataplex-search-entries"
DataplexLookupEntryToolKind = "dataplex-lookup-entry"
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
DataplexSourceKind = "dataplex"
DataplexSearchEntriesToolKind = "dataplex-search-entries"
DataplexLookupEntryToolKind = "dataplex-lookup-entry"
DataplexSearchAspectTypesToolKind = "dataplex-search-aspect-types"
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
)
func getDataplexVars(t *testing.T) map[string]any {
@@ -69,6 +72,19 @@ func initBigQueryConnection(ctx context.Context, project string) (*bigqueryapi.C
return client, nil
}
func initDataplexConnection(ctx context.Context) (*dataplex.CatalogClient, 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.NewCatalogClient(ctx, option.WithCredentials(cred))
if err != nil {
return nil, fmt.Errorf("failed to create Dataplex client %w", err)
}
return client, nil
}
func TestDataplexToolEndpoints(t *testing.T) {
sourceConfig := getDataplexVars(t)
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
@@ -81,12 +97,21 @@ func TestDataplexToolEndpoints(t *testing.T) {
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
}
// create table name with UUID
dataplexClient, err := initDataplexConnection(ctx)
if err != nil {
t.Fatalf("unable to create Dataplex connection: %s", err)
}
// create resources with UUID
datasetName := fmt.Sprintf("temp_toolbox_test_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
aspectTypeId := fmt.Sprintf("param-aspect-type-%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
teardownTable1 := setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName)
teardownAspectType1 := setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId)
time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested
defer teardownTable1(t)
defer teardownAspectType1(t)
toolsFile := getDataplexToolsConfig(sourceConfig)
@@ -107,6 +132,7 @@ func TestDataplexToolEndpoints(t *testing.T) {
runDataplexToolGetTest(t)
runDataplexSearchEntriesToolInvokeTest(t, tableName, datasetName)
runDataplexLookupEntryToolInvokeTest(t, tableName, datasetName)
runDataplexSearchAspectTypesToolInvokeTest(t, aspectTypeId)
}
func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) {
@@ -132,8 +158,6 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C
t.Fatalf("Create table job for %s failed: %v", tableName, err)
}
time.Sleep(2 * time.Minute) // wait for table to be ingested
return func(t *testing.T) {
// tear down table
dropSQL := fmt.Sprintf("drop table %s.%s", datasetName, tableName)
@@ -166,6 +190,35 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C
}
}
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{
Parent: parent,
AspectTypeId: aspectTypeId,
AspectType: &dataplexpb.AspectType{
Name: fmt.Sprintf("%s/aspectTypes/%s", parent, aspectTypeId),
MetadataTemplate: &dataplexpb.AspectType_MetadataTemplate{
Name: "UserSchema",
Type: "record",
},
},
}
_, err := client.CreateAspectType(ctx, createAspectTypeReq)
if err != nil {
t.Fatalf("Failed to create aspect type %s: %v", aspectTypeId, err)
}
return func(t *testing.T) {
// tear down aspect type
deleteAspectTypeReq := &dataplexpb.DeleteAspectTypeRequest{
Name: fmt.Sprintf("%s/aspectTypes/%s", parent, aspectTypeId),
}
if _, err := client.DeleteAspectType(ctx, deleteAspectTypeReq); err != nil {
t.Errorf("Failed to delete aspect type %s: %v", aspectTypeId, err)
}
}
}
func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
// Write config into a file and pass it to command
toolsFile := map[string]any{
@@ -182,12 +235,12 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
"my-dataplex-search-entries-tool": map[string]any{
"kind": DataplexSearchEntriesToolKind,
"source": "my-dataplex-instance",
"description": "Simple tool to test end to end functionality.",
"description": "Simple dataplex search entries tool to test end to end functionality.",
},
"my-auth-dataplex-search-entries-tool": map[string]any{
"kind": DataplexSearchEntriesToolKind,
"source": "my-dataplex-instance",
"description": "Simple tool to test end to end functionality.",
"description": "Simple dataplex search entries tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
"my-dataplex-lookup-entry-tool": map[string]any{
@@ -201,6 +254,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
"description": "Simple dataplex lookup entry tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
"my-dataplex-search-aspect-types-tool": map[string]any{
"kind": DataplexSearchAspectTypesToolKind,
"source": "my-dataplex-instance",
"description": "Simple dataplex search aspect types tool to test end to end functionality.",
},
"my-auth-dataplex-search-aspect-types-tool": map[string]any{
"kind": DataplexSearchAspectTypesToolKind,
"source": "my-dataplex-instance",
"description": "Simple dataplex search aspect types tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
},
}
@@ -216,13 +280,18 @@ func runDataplexToolGetTest(t *testing.T) {
{
name: "get my-dataplex-search-entries-tool",
toolName: "my-dataplex-search-entries-tool",
expectedParams: []string{"pageSize", "pageToken", "query", "orderBy", "semanticSearch"},
expectedParams: []string{"pageSize", "query", "orderBy"},
},
{
name: "get my-dataplex-lookup-entry-tool",
toolName: "my-dataplex-lookup-entry-tool",
expectedParams: []string{"name", "view", "aspectTypes", "entry"},
},
{
name: "get my-dataplex-search-aspect-types-tool",
toolName: "my-dataplex-search-aspect-types-tool",
expectedParams: []string{"pageSize", "query", "orderBy"},
},
}
for _, tc := range testCases {
@@ -559,3 +628,119 @@ func runDataplexLookupEntryToolInvokeTest(t *testing.T, tableName string, datase
})
}
}
func runDataplexSearchAspectTypesToolInvokeTest(t *testing.T, aspectTypeId string) {
idToken, err := tests.GetGoogleIdToken(tests.ClientId)
if err != nil {
t.Fatalf("error getting Google ID token: %s", err)
}
testCases := []struct {
name string
api string
requestHeader map[string]string
requestBody io.Reader
wantStatusCode int
expectResult bool
wantContentKey string
}{
{
name: "Success - Aspect Type Found",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-aspect-types-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "metadata_template",
},
{
name: "Success - Aspect Type Found with Authorization",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": idToken},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "metadata_template",
},
{
name: "Failure - Aspect Type Not Found",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-aspect-types-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`"{\"query\":\"name=_aspectType\"}"`)),
wantStatusCode: 400,
expectResult: false,
},
{
name: "Failure - Invalid Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": "invalid_token"},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
wantStatusCode: 401,
expectResult: false,
},
{
name: "Failure - No Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
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)
}
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 {
if result["result"] == nil && !tc.expectResult {
return
}
t.Fatalf("expected 'result' field to be a string, got %T", result["result"])
}
if !tc.expectResult && (resultStr == "" || resultStr == "[]") {
return
}
var entries []interface{}
if err := json.Unmarshal([]byte(resultStr), &entries); err != nil {
t.Fatalf("error unmarshalling result string: %v", err)
}
if tc.expectResult {
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])
}
if _, ok := entry[tc.wantContentKey]; !ok {
t.Fatalf("expected entry to have key '%s', but it was not found in %v", tc.wantContentKey, entry)
}
} else {
if len(entries) != 0 {
t.Fatalf("expected 0 entries, but got %d", len(entries))
}
}
})
}
}