feat(tools/dataplex-lookup-entry): Add support for dataplex-lookup-entry tool (#1009)

Added support for lookup entry tool in Dataplex.
Fixes #997

---------

Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
Anuj Jhunjhunwala
2025-08-01 21:18:56 +02:00
committed by GitHub
parent d9ee17d2c7
commit 5fa1660fc8
10 changed files with 904 additions and 109 deletions

View File

@@ -1,20 +0,0 @@
load("//tools/build_defs/go:go_library.bzl", "go_library")
load("//tools/build_defs/go:go_test.bzl", "go_test")
go_library(
name = "cmd",
srcs = [
"options.go",
"root.go",
],
embedsrcs = ["version.txt"],
)
go_test(
name = "cmd_test",
srcs = [
"options_test.go",
"root_test.go",
],
library = ":cmd",
)

View File

@@ -51,6 +51,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
_ "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/dataplexsearchentries"
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/tools/duckdbsql"

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"},
ToolNames: []string{"dataplex_search_entries", "dataplex_lookup_entry"},
},
},
},

View File

@@ -35,10 +35,213 @@ You can use the following system prompt as "Custom Instructions" in your client
application.
```
Whenever you will receive response from dataplex_search_entries tool decide what do to by following these steps:
# Objective
Your primary objective is to help discover, organize and manage metadata related to data assets.
# Tone and Style
1. Adopt the persona of a senior subject matter expert
2. Your communication style must be:
1. Concise: Always favor brevity.
2. Direct: Avoid greetings (e.g., "Hi there!", "Certainly!"). Get straight to the point.
Example (Incorrect): Hi there! I see that you are looking for...
Example (Correct): This problem likely stems from...
3. Do not reiterate or summarize the question in the answer.
4. Crucially, always convey a tone of uncertainty and caution. Since you are interpreting metadata and have no way to externally verify your answers, never express complete confidence. Frame your responses as interpretations based solely on the provided metadata. Use a suggestive tone, not a prescriptive one:
Example (Correct): "The entry describes..."
Example (Correct): "According to catalog,..."
Example (Correct): "Based on the metadata,..."
Example (Correct): "Based on the search results,..."
5. Do not make assumptions
# Data Model
## Entries
Entry represents a specific data asset. Entry acts as a metadata record for something that is managed by Catalog, such as:
- A BigQuery table or dataset
- A Cloud Storage bucket or folder
- An on-premises SQL table
## Aspects
While the Entry itself is a container, the rich descriptive information about the asset (e.g., schema, data types, business descriptions, classifications) is stored in associated components called Aspects. Aspects are created based on pre-defined blueprints known as Aspect Types.
## Aspect Types
Aspect Type is a reusable template that defines the schema for a set of metadata fields. Think of an Aspect Type as a structure for the kind of metadata that is organized in the catalog within the Entry.
Examples:
- projects/dataplex-types/locations/global/aspectTypes/analytics-hub-exchange
- projects/dataplex-types/locations/global/aspectTypes/analytics-hub
- projects/dataplex-types/locations/global/aspectTypes/analytics-hub-listing
- projects/dataplex-types/locations/global/aspectTypes/bigquery-connection
- projects/dataplex-types/locations/global/aspectTypes/bigquery-data-policy
- projects/dataplex-types/locations/global/aspectTypes/bigquery-dataset
- projects/dataplex-types/locations/global/aspectTypes/bigquery-model
- projects/dataplex-types/locations/global/aspectTypes/bigquery-policy
- projects/dataplex-types/locations/global/aspectTypes/bigquery-routine
- projects/dataplex-types/locations/global/aspectTypes/bigquery-row-access-policy
- projects/dataplex-types/locations/global/aspectTypes/bigquery-table
- projects/dataplex-types/locations/global/aspectTypes/bigquery-view
- projects/dataplex-types/locations/global/aspectTypes/cloud-bigtable-instance
- projects/dataplex-types/locations/global/aspectTypes/cloud-bigtable-table
- projects/dataplex-types/locations/global/aspectTypes/cloud-spanner-database
- projects/dataplex-types/locations/global/aspectTypes/cloud-spanner-instance
- projects/dataplex-types/locations/global/aspectTypes/cloud-spanner-table
- projects/dataplex-types/locations/global/aspectTypes/cloud-spanner-view
- projects/dataplex-types/locations/global/aspectTypes/cloudsql-database
- projects/dataplex-types/locations/global/aspectTypes/cloudsql-instance
- projects/dataplex-types/locations/global/aspectTypes/cloudsql-schema
- projects/dataplex-types/locations/global/aspectTypes/cloudsql-table
- projects/dataplex-types/locations/global/aspectTypes/cloudsql-view
- projects/dataplex-types/locations/global/aspectTypes/contacts
- projects/dataplex-types/locations/global/aspectTypes/dataform-code-asset
- projects/dataplex-types/locations/global/aspectTypes/dataform-repository
- projects/dataplex-types/locations/global/aspectTypes/dataform-workspace
- projects/dataplex-types/locations/global/aspectTypes/dataproc-metastore-database
- projects/dataplex-types/locations/global/aspectTypes/dataproc-metastore-service
- projects/dataplex-types/locations/global/aspectTypes/dataproc-metastore-table
- projects/dataplex-types/locations/global/aspectTypes/data-product
- projects/dataplex-types/locations/global/aspectTypes/data-quality-scorecard
- projects/dataplex-types/locations/global/aspectTypes/external-connection
- projects/dataplex-types/locations/global/aspectTypes/overview
- projects/dataplex-types/locations/global/aspectTypes/pubsub-topic
- projects/dataplex-types/locations/global/aspectTypes/schema
- projects/dataplex-types/locations/global/aspectTypes/sensitive-data-protection-job-result
- projects/dataplex-types/locations/global/aspectTypes/sensitive-data-protection-profile
- projects/dataplex-types/locations/global/aspectTypes/sql-access
- projects/dataplex-types/locations/global/aspectTypes/storage-bucket
- projects/dataplex-types/locations/global/aspectTypes/storage-folder
- projects/dataplex-types/locations/global/aspectTypes/storage
- projects/dataplex-types/locations/global/aspectTypes/usage
## Entry Types
Every Entry must conform to an Entry Type. The Entry Type acts as a template, defining the structure, required aspects, and constraints for Entries of that type.
Examples:
- projects/dataplex-types/locations/global/entryTypes/analytics-hub-exchange
- projects/dataplex-types/locations/global/entryTypes/analytics-hub-listing
- projects/dataplex-types/locations/global/entryTypes/bigquery-connection
- projects/dataplex-types/locations/global/entryTypes/bigquery-data-policy
- projects/dataplex-types/locations/global/entryTypes/bigquery-dataset
- projects/dataplex-types/locations/global/entryTypes/bigquery-model
- projects/dataplex-types/locations/global/entryTypes/bigquery-routine
- projects/dataplex-types/locations/global/entryTypes/bigquery-row-access-policy
- projects/dataplex-types/locations/global/entryTypes/bigquery-table
- projects/dataplex-types/locations/global/entryTypes/bigquery-view
- projects/dataplex-types/locations/global/entryTypes/cloud-bigtable-instance
- projects/dataplex-types/locations/global/entryTypes/cloud-bigtable-table
- projects/dataplex-types/locations/global/entryTypes/cloud-spanner-database
- projects/dataplex-types/locations/global/entryTypes/cloud-spanner-instance
- projects/dataplex-types/locations/global/entryTypes/cloud-spanner-table
- projects/dataplex-types/locations/global/entryTypes/cloud-spanner-view
- projects/dataplex-types/locations/global/entryTypes/cloudsql-mysql-database
- projects/dataplex-types/locations/global/entryTypes/cloudsql-mysql-instance
- projects/dataplex-types/locations/global/entryTypes/cloudsql-mysql-table
- projects/dataplex-types/locations/global/entryTypes/cloudsql-mysql-view
- projects/dataplex-types/locations/global/entryTypes/cloudsql-postgresql-database
- projects/dataplex-types/locations/global/entryTypes/cloudsql-postgresql-instance
- projects/dataplex-types/locations/global/entryTypes/cloudsql-postgresql-schema
- projects/dataplex-types/locations/global/entryTypes/cloudsql-postgresql-table
- projects/dataplex-types/locations/global/entryTypes/cloudsql-postgresql-view
- projects/dataplex-types/locations/global/entryTypes/cloudsql-sqlserver-database
- projects/dataplex-types/locations/global/entryTypes/cloudsql-sqlserver-instance
- projects/dataplex-types/locations/global/entryTypes/cloudsql-sqlserver-schema
- projects/dataplex-types/locations/global/entryTypes/cloudsql-sqlserver-table
- projects/dataplex-types/locations/global/entryTypes/cloudsql-sqlserver-view
- projects/dataplex-types/locations/global/entryTypes/dataform-code-asset
- projects/dataplex-types/locations/global/entryTypes/dataform-repository
- projects/dataplex-types/locations/global/entryTypes/dataform-workspace
- projects/dataplex-types/locations/global/entryTypes/dataproc-metastore-database
- projects/dataplex-types/locations/global/entryTypes/dataproc-metastore-service
- projects/dataplex-types/locations/global/entryTypes/dataproc-metastore-table
- projects/dataplex-types/locations/global/entryTypes/pubsub-topic
- projects/dataplex-types/locations/global/entryTypes/storage-bucket
- projects/dataplex-types/locations/global/entryTypes/storage-folder
- projects/dataplex-types/locations/global/entryTypes/vertexai-dataset
- projects/dataplex-types/locations/global/entryTypes/vertexai-feature-group
- projects/dataplex-types/locations/global/entryTypes/vertexai-feature-online-store
## Entry Groups
Entries are organized within Entry Groups, which are logical groupings of Entries. An Entry Group acts as a namespace for its Entries.
## Entry Links
Entries can be linked together using EntryLinks to represent relationships between data assets (e.g. foreign keys).
# Tool instructions
## Tool: dataplex_search_entries
## General
- Do not try to search within search results on your own.
- Do not fetch multiple pages of results unless explicitly asked.
## Search syntax
### Simple search
In its simplest form, a search query consists of a single predicate. Such a predicate can match several pieces of metadata:
- A substring of a name, display name, or description of a resource
- A substring of the type of a resource
- A substring of a column name (or nested column name) in the schema of a resource
- A substring of a project ID
- A string from an overview description
For example, the predicate foo matches the following resources:
- Resource with the name foo.bar
- Resource with the display name Foo Bar
- Resource with the description This is the foo script
- Resource with the exact type foo
- Column foo_bar in the schema of a resource
- Nested column foo_bar in the schema of a resource
- Project prod-foo-bar
- Resource with an overview containing the word foo
### Qualified predicates
You can qualify a predicate by prefixing it with a key that restricts the matching to a specific piece of metadata:
- An equal sign (=) restricts the search to an exact match.
- A colon (:) after the key matches the predicate to either a substring or a token within the value in the search results.
Tokenization splits the stream of text into a series of tokens, with each token usually corresponding to a single word. For example:
- name:foo selects resources with names that contain the foo substring, like foo1 and barfoo.
- description:foo selects resources with the foo token in the description, like bar and foo.
- location=foo matches resources in a specified location with foo as the location name.
The predicate keys type, system, location, and orgid support only the exact match (=) qualifier, not the substring qualifier (:). For example, type=foo or orgid=number.
Search syntax supports the following qualifiers:
- "name:x" - Matches x as a substring of the resource ID.
- "displayname:x" - Match x as a substring of the resource display name.
- "column:x" - Matches x as a substring of the column name (or nested column name) in the schema of the resource.
- "description:x" - Matches x as a token in the resource description.
- "label:bar" - Matches BigQuery resources that have a label (with some value) and the label key has bar as a substring.
- "label=bar" - Matches BigQuery resources that have a label (with some value) and the label key equals bar as a string.
- "label:bar:x" - Matches x as a substring in the value of a label with a key bar attached to a BigQuery resource.
- "label=foo:bar" - Matches BigQuery resources where the key equals foo and the key value equals bar.
- "label.foo=bar" - Matches BigQuery resources where the key equals foo and the key value equals bar.
- "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.
- "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.
### 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.
### Request
1. Always try to rewrite the prompt using search syntax.
### Response
1. If there are multiple search results found
1.1. Present the list of search results
1.2. Format the output in nested ordered list, for example:
1. Present the list of search results
2. Format the output in nested ordered list, for example:
Given
```
{
@@ -75,14 +278,19 @@ Whenever you will receive response from dataplex_search_entries tool decide what
- location: us-central1
- description: Table contains list of best customers.
```
1.3. Ask to select one of the presented search results
3. Ask to select one of the presented search results
2. If there is only one search result found
2.1. Present the search result immediately.
1. Present the search result immediately.
3. If there are no search result found
3.1. Explain that no search result was found
3.2. Suggest to provide a more specific search query.
1. Explain that no search result was found
2. Suggest to provide a more specific search query.
Do not try to search within search results on your own.
## 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.
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.
```
## Reference
@@ -90,4 +298,4 @@ Do not try to search within search results on your own.
| **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,60 @@
---
title: "dataplex-lookup-entry"
type: docs
weight: 1
description: >
A "dataplex-lookup-entry" tool returns details of a particular entry in Dataplex Catalog.
aliases:
- /resources/tools/dataplex-lookup-entry
---
## About
A `dataplex-lookup-entry` tool returns details of a particular entry in Dataplex Catalog.
It's compatible with the following sources:
- [dataplex](../sources/dataplex.md)
`dataplex-lookup-entry` takes a required `name` parameter which contains the project and location to which the request should be attributed in the following form: projects/{project}/locations/{location} and also a required `entry` parameter which is the resource name of the entry in the following form: projects/{project}/locations/{location}/entryGroups/{entryGroup}/entries/{entry}. It also optionally accepts following parameters:
- `view` - View to control which parts of an entry the service should return. It takes integer values from 1-4 corresponding to type of view - BASIC, FULL, CUSTOM, ALL
- `aspectTypes` - Limits the aspects returned to the provided aspect types in the format `projects/{project}/locations/{location}/aspectTypes/{aspectType}`. It only works for CUSTOM view.
- `paths` - Limits the aspects returned to those associated with the provided paths within the Entry. It only works for CUSTOM view.
## 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
## Example
```yaml
tools:
lookup_entry:
kind: dataplex-lookup-entry
source: my-dataplex-source
description: Use this tool to retrieve a specific entry in Dataplex Catalog.
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "dataplex-lookup-entry". |
| 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

@@ -7,9 +7,13 @@ tools:
dataplex_search_entries:
kind: dataplex-search-entries
source: dataplex-source
description: |
Use this tool to search for entries in Dataplex Catalog that represent data assets (e.g. tables, views, models) based on the provided search query.
description: Use this tool to search for entries in Dataplex Catalog based on the provided search query.
dataplex_lookup_entry:
kind: dataplex-lookup-entry
source: dataplex-source
description: Use this tool to retrieve a specific entry from Dataplex Catalog.
toolsets:
dataplex-tools:
- dataplex_search_entries
- dataplex_search_entries
- dataplex_lookup_entry

View File

@@ -0,0 +1,183 @@
// 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 dataplexlookupentry
import (
"context"
"fmt"
dataplexapi "cloud.google.com/go/dataplex/apiv1"
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
"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-lookup-entry"
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
}
// 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"`
Parameters tools.Parameters `yaml:"parameters"`
}
// 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)
}
viewDesc := `
## Argument: view
**Type:** Integer
**Description:** Specifies the parts of the entry and its aspects to return.
**Possible Values:**
* 1 (BASIC): Returns entry without aspects.
* 2 (FULL): Return all required aspects and the keys of non-required aspects. (Default)
* 3 (CUSTOM): Return the entry and aspects requested in aspect_types field (at most 100 aspects). Always use this view when aspect_types is not empty.
* 4 (ALL): Return the entry and both required and optional aspects (at most 100 aspects)
`
name := tools.NewStringParameter("name", "The project to which the request should be attributed in the following form: projects/{project}/locations/{location}.")
view := tools.NewIntParameterWithDefault("view", 2, viewDesc)
aspectTypes := tools.NewArrayParameterWithDefault("aspectTypes", []any{}, "Limits the aspects returned to the provided aspect types. It only works when used together with CUSTOM view.", tools.NewStringParameter("aspectType", "The types of aspects to be included in the response in the format `projects/{project}/locations/{location}/aspectTypes/{aspectType}`."))
entry := tools.NewStringParameter("entry", "The resource name of the Entry in the following form: projects/{project}/locations/{location}/entryGroups/{entryGroup}/entries/{entry}.")
parameters := tools.Parameters{name, view, aspectTypes, entry}
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(),
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
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) {
paramsMap := params.AsMap()
viewMap := map[int]dataplexpb.EntryView{
1: dataplexpb.EntryView_BASIC,
2: dataplexpb.EntryView_FULL,
3: dataplexpb.EntryView_CUSTOM,
4: dataplexpb.EntryView_ALL,
}
name, _ := paramsMap["name"].(string)
entry, _ := paramsMap["entry"].(string)
view, _ := paramsMap["view"].(int)
aspectTypeSlice, err := tools.ConvertAnySliceToTyped(paramsMap["aspectTypes"].([]any), "string")
if err != nil {
return nil, fmt.Errorf("can't convert aspectTypes to array of strings: %s", err)
}
aspectTypes := aspectTypeSlice.([]string)
req := &dataplexpb.LookupEntryRequest{
Name: name,
View: viewMap[view],
AspectTypes: aspectTypes,
Entry: entry,
}
result, err := t.CatalogClient.LookupEntry(ctx, req)
if err != nil {
return nil, err
}
return result, 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,117 @@
// 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 dataplexlookupentry_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"
"github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
)
func TestParseFromYamlDataplexLookupEntry(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-lookup-entry
source: my-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": dataplexlookupentry.Config{
Name: "example_tool",
Kind: "dataplex-lookup-entry",
Source: "my-instance",
Description: "some description",
AuthRequired: []string{},
},
},
},
{
desc: "advanced example",
in: `
tools:
example_tool:
kind: dataplex-lookup-entry
source: my-instance
description: some description
parameters:
- name: name
type: string
description: some name description
- name: view
type: string
description: some view description
- name: aspectTypes
type: array
description: some aspect types description
default: []
items:
name: aspectType
type: string
description: some aspect type description
- name: entry
type: string
description: some entry description
`,
want: server.ToolConfigs{
"example_tool": dataplexlookupentry.Config{
Name: "example_tool",
Kind: "dataplex-lookup-entry",
Source: "my-instance",
Description: "some description",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("name", "some name description"),
tools.NewStringParameter("view", "some view description"),
tools.NewArrayParameterWithDefault("aspectTypes", []any{}, "some aspect types description", tools.NewStringParameter("aspectType", "some aspect type description")),
tools.NewStringParameter("entry", "some entry description"),
},
},
},
},
}
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

@@ -80,12 +80,11 @@ 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.")
name := tools.NewStringParameterWithDefault("name", fmt.Sprintf("projects/%s/locations/global", s.ProjectID()), "The project to which the request should be attributed in the following form: projects/{project}/locations/global")
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, name, pageSize, pageToken, orderBy, semanticSearch}
parameters := tools.Parameters{query, pageSize, pageToken, orderBy, semanticSearch}
mcpManifest := tools.McpManifest{
Name: cfg.Name,
@@ -93,7 +92,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
InputSchema: parameters.McpManifest(),
}
t := &SearchTool{
t := &Tool{
Name: cfg.Name,
Kind: kind,
Parameters: parameters,
@@ -110,7 +109,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
return t, nil
}
type SearchTool struct {
type Tool struct {
Name string
Kind string
Parameters tools.Parameters
@@ -121,14 +120,13 @@ type SearchTool struct {
mcpManifest tools.McpManifest
}
func (t *SearchTool) Authorized(verifiedAuthServices []string) bool {
func (t *Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t *SearchTool) 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)
name, _ := paramsMap["name"].(string)
pageSize, _ := paramsMap["pageSize"].(int32)
pageToken, _ := paramsMap["pageToken"].(string)
orderBy, _ := paramsMap["orderBy"].(string)
@@ -136,7 +134,7 @@ func (t *SearchTool) Invoke(ctx context.Context, params tools.ParamValues) (any,
req := &dataplexpb.SearchEntriesRequest{
Query: query,
Name: name,
Name: fmt.Sprintf("projects/%s/locations/global", t.ProjectID),
PageSize: pageSize,
PageToken: pageToken,
OrderBy: orderBy,
@@ -159,17 +157,17 @@ func (t *SearchTool) Invoke(ctx context.Context, params tools.ParamValues) (any,
return results, nil
}
func (t *SearchTool) 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 *SearchTool) Manifest() tools.Manifest {
func (t *Tool) Manifest() tools.Manifest {
// Returns the tool manifest
return t.manifest
}
func (t *SearchTool) McpManifest() tools.McpManifest {
func (t *Tool) McpManifest() tools.McpManifest {
// Returns the tool MCP manifest
return t.mcpManifest
}

View File

@@ -19,6 +19,7 @@ import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"regexp"
@@ -39,6 +40,7 @@ import (
var (
DataplexSourceKind = "dataplex"
DataplexSearchEntriesToolKind = "dataplex-search-entries"
DataplexLookupEntryToolKind = "dataplex-lookup-entry"
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
)
@@ -69,7 +71,7 @@ func initBigQueryConnection(ctx context.Context, project string) (*bigqueryapi.C
func TestDataplexToolEndpoints(t *testing.T) {
sourceConfig := getDataplexVars(t)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
defer cancel()
var args []string
@@ -94,7 +96,7 @@ func TestDataplexToolEndpoints(t *testing.T) {
}
defer cleanup()
waitCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
waitCtx, cancel := context.WithTimeout(ctx, 3*time.Minute)
defer cancel()
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
@@ -102,8 +104,9 @@ func TestDataplexToolEndpoints(t *testing.T) {
t.Fatalf("toolbox didn't start successfully: %s", err)
}
runDataplexSearchEntriesToolGetTest(t)
runDataplexToolGetTest(t)
runDataplexSearchEntriesToolInvokeTest(t, tableName, datasetName)
runDataplexLookupEntryToolInvokeTest(t, tableName, datasetName)
}
func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) {
@@ -169,92 +172,169 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
"sources": map[string]any{
"my-dataplex-instance": sourceConfig,
},
"authServices": map[string]any{
"my-google-auth": map[string]any{
"kind": "google",
"clientId": tests.ClientId,
},
},
"tools": map[string]any{
"my-search-entries-tool": 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.",
},
"my-auth-dataplex-search-entries-tool": map[string]any{
"kind": DataplexSearchEntriesToolKind,
"source": "my-dataplex-instance",
"description": "Simple tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
"my-dataplex-lookup-entry-tool": map[string]any{
"kind": DataplexLookupEntryToolKind,
"source": "my-dataplex-instance",
"description": "Simple dataplex lookup entry tool to test end to end functionality.",
},
"my-auth-dataplex-lookup-entry-tool": map[string]any{
"kind": DataplexLookupEntryToolKind,
"source": "my-dataplex-instance",
"description": "Simple dataplex lookup entry tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
},
}
return toolsFile
}
func runDataplexSearchEntriesToolGetTest(t *testing.T) {
resp, err := http.Get("http://127.0.0.1:5000/api/tool/my-search-entries-tool/")
if err != nil {
t.Fatalf("error making GET request: %s", err)
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
t.Fatalf("expected status code 200, got %d", resp.StatusCode)
}
var body map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
t.Fatalf("error decoding response body: %s", err)
}
got, ok := body["tools"]
if !ok {
t.Fatalf("unable to find 'tools' key in response body")
func runDataplexToolGetTest(t *testing.T) {
testCases := []struct {
name string
toolName string
expectedParams []string
}{
{
name: "get my-dataplex-search-entries-tool",
toolName: "my-dataplex-search-entries-tool",
expectedParams: []string{"pageSize", "pageToken", "query", "orderBy", "semanticSearch"},
},
{
name: "get my-dataplex-lookup-entry-tool",
toolName: "my-dataplex-lookup-entry-tool",
expectedParams: []string{"name", "view", "aspectTypes", "entry"},
},
}
toolsMap, ok := got.(map[string]interface{})
if !ok {
t.Fatalf("tools is not a map")
}
tool, ok := toolsMap["my-search-entries-tool"].(map[string]interface{})
if !ok {
t.Fatalf("tool not found in manifest")
}
params, ok := tool["parameters"].([]interface{})
if !ok {
t.Fatalf("parameters not found")
}
paramNames := []string{}
for _, param := range params {
paramMap, ok := param.(map[string]interface{})
if ok {
paramNames = append(paramNames, paramMap["name"].(string))
}
}
expected := []string{"name", "pageSize", "pageToken", "orderBy", "query"}
for _, want := range expected {
found := false
for _, got := range paramNames {
if got == want {
found = true
break
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
resp, err := http.Get(fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/", tc.toolName))
if err != nil {
t.Fatalf("error when sending a request: %s", err)
}
}
if !found {
t.Fatalf("expected parameter %q not found in tool parameters", want)
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
t.Fatalf("response status code is not 200")
}
var body map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&body)
if err != nil {
t.Fatalf("error parsing response body")
}
got, ok := body["tools"]
if !ok {
t.Fatalf("unable to find tools in response body")
}
toolsMap, ok := got.(map[string]interface{})
if !ok {
t.Fatalf("expected 'tools' to be a map, got %T", got)
}
tool, ok := toolsMap[tc.toolName].(map[string]interface{})
if !ok {
t.Fatalf("expected tool %q to be a map, got %T", tc.toolName, toolsMap[tc.toolName])
}
params, ok := tool["parameters"].([]interface{})
if !ok {
t.Fatalf("expected 'parameters' to be a slice, got %T", tool["parameters"])
}
paramSet := make(map[string]struct{})
for _, param := range params {
paramMap, ok := param.(map[string]interface{})
if ok {
if name, ok := paramMap["name"].(string); ok {
paramSet[name] = struct{}{}
}
}
}
var missing []string
for _, want := range tc.expectedParams {
if _, found := paramSet[want]; !found {
missing = append(missing, want)
}
}
if len(missing) > 0 {
t.Fatalf("missing parameters for tool %q: %v", tc.toolName, missing)
}
})
}
}
func runDataplexSearchEntriesToolInvokeTest(t *testing.T, tableName string, datasetName string) {
idToken, err := tests.GetGoogleIdToken(tests.ClientId)
if err != nil {
t.Fatalf("error getting Google ID token: %s", err)
}
testCases := []struct {
name string
tableName string
datasetName string
api string
requestHeader map[string]string
requestBody io.Reader
wantStatusCode int
expectResult bool
wantContentKey string
}{
{
name: "Success - Entry Found",
tableName: tableName,
datasetName: datasetName,
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-entries-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"displayname=%s system=bigquery parent=%s\"}", tableName, datasetName))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "dataplex_entry",
},
{
name: "Success with Authorization - Entry Found",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-entries-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": idToken},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"displayname=%s system=bigquery parent=%s\"}", tableName, datasetName))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "dataplex_entry",
},
{
name: "Failure - Invalid Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-entries-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": "invalid_token"},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"displayname=%s system=bigquery parent=%s\"}", tableName, datasetName))),
wantStatusCode: 401,
expectResult: false,
wantContentKey: "dataplex_entry",
},
{
name: "Failure - Without Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-entries-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"displayname=%s system=bigquery parent=%s\"}", tableName, datasetName))),
wantStatusCode: 401,
expectResult: false,
wantContentKey: "dataplex_entry",
},
{
name: "Failure - Entry Not Found",
tableName: "",
datasetName: "",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-entries-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"query":"displayname=\"\" system=bigquery parent=\"\""}`)),
wantStatusCode: 200,
expectResult: false,
wantContentKey: "",
@@ -263,19 +343,23 @@ func runDataplexSearchEntriesToolInvokeTest(t *testing.T, tableName string, data
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
query := fmt.Sprintf("displayname=\"%s\" system=bigquery parent:\"%s\"", tc.tableName, tc.datasetName)
reqBodyMap := map[string]string{"query": query}
reqBodyBytes, err := json.Marshal(reqBodyMap)
req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
if err != nil {
t.Fatalf("error marshalling request body: %s", err)
t.Fatalf("unable to create request: %s", err)
}
resp, err := http.Post("http://127.0.0.1:5000/api/tool/my-search-entries-tool/invoke", "application/json", bytes.NewBuffer(reqBodyBytes))
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 making POST request: %s", err)
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.", tc.wantStatusCode)
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
bodyBytes, _ := io.ReadAll(resp.Body)
t.Fatalf("Response body: %s", string(bodyBytes))
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
@@ -297,8 +381,8 @@ func runDataplexSearchEntriesToolInvokeTest(t *testing.T, tableName string, data
}
if tc.expectResult {
if len(entries) == 0 {
t.Fatal("expected at least one entry, but got 0")
if len(entries) != 1 {
t.Fatalf("expected exactly one entry, but got %d", len(entries))
}
entry, ok := entries[0].(map[string]interface{})
if !ok {
@@ -315,3 +399,163 @@ func runDataplexSearchEntriesToolInvokeTest(t *testing.T, tableName string, data
})
}
}
func runDataplexLookupEntryToolInvokeTest(t *testing.T, tableName string, datasetName string) {
idToken, err := tests.GetGoogleIdToken(tests.ClientId)
if err != nil {
t.Fatalf("error getting Google ID token: %s", err)
}
testCases := []struct {
name string
wantStatusCode int
api string
requestHeader map[string]string
requestBody io.Reader
expectResult bool
wantContentKey string
dontWantContentKey string
aspectCheck bool
reqBodyMap map[string]any
}{
{
name: "Success - Entry Found",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s\"}", DataplexProject, DataplexProject, DataplexProject, datasetName))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "name",
},
{
name: "Success - Entry Found with Authorization",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": idToken},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s\"}", DataplexProject, DataplexProject, DataplexProject, datasetName))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "name",
},
{
name: "Failure - Invalid Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": "invalid_token"},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s\"}", DataplexProject, DataplexProject, DataplexProject, datasetName))),
wantStatusCode: 401,
expectResult: false,
wantContentKey: "name",
},
{
name: "Failure - Without Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s\"}", DataplexProject, DataplexProject, DataplexProject, datasetName))),
wantStatusCode: 401,
expectResult: false,
wantContentKey: "name",
},
{
name: "Failure - Entry Not Found or Permission Denied",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s\"}", DataplexProject, DataplexProject, DataplexProject, "non-existent-dataset"))),
wantStatusCode: 400,
expectResult: false,
},
{
name: "Success - Entry Found with Basic View",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s\", \"view\": %d}", DataplexProject, DataplexProject, DataplexProject, datasetName, tableName, 1))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "name",
dontWantContentKey: "aspects",
},
{
name: "Failure - Entry with Custom View without Aspect Types",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s\", \"view\": %d}", DataplexProject, DataplexProject, DataplexProject, datasetName, tableName, 3))),
wantStatusCode: 400,
expectResult: false,
},
{
name: "Success - Entry Found with only Schema Aspect",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-lookup-entry-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"name\":\"projects/%s/locations/us\", \"entry\":\"projects/%s/locations/us/entryGroups/@bigquery/entries/bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s\", \"aspectTypes\":[\"projects/dataplex-types/locations/global/aspectTypes/schema\"], \"view\": %d}", DataplexProject, DataplexProject, DataplexProject, datasetName, tableName, 3))),
wantStatusCode: 200,
expectResult: true,
wantContentKey: "aspects",
aspectCheck: true,
},
}
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 {
bodyBytes, _ := io.ReadAll(resp.Body)
t.Fatalf("Response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
t.Fatalf("Error parsing response body: %v", err)
}
if tc.expectResult {
resultStr, ok := result["result"].(string)
if !ok {
t.Fatalf("Expected 'result' field to be a string on success, got %T", result["result"])
}
if resultStr == "" || resultStr == "{}" || resultStr == "null" {
t.Fatal("Expected an entry, but got empty result")
}
var entry map[string]interface{}
if err := json.Unmarshal([]byte(resultStr), &entry); err != nil {
t.Fatalf("Error unmarshalling result string into entry map: %v", err)
}
if _, ok := entry[tc.wantContentKey]; !ok {
t.Fatalf("Expected entry to have key '%s', but it was not found in %v", tc.wantContentKey, entry)
}
if _, ok := entry[tc.dontWantContentKey]; ok {
t.Fatalf("Expected entry to not have key '%s', but it was found in %v", tc.dontWantContentKey, entry)
}
if tc.aspectCheck {
// Check length of aspects
aspects, ok := entry["aspects"].(map[string]interface{})
if !ok {
t.Fatalf("Expected 'aspects' to be a map, got %T", aspects)
}
if len(aspects) != 1 {
t.Fatalf("Expected exactly one aspect, but got %d", len(aspects))
}
}
} else { // Handle expected error response
_, ok := result["error"]
if !ok {
t.Fatalf("Expected 'error' field in response, got %v", result)
}
}
})
}
}