github-featured/googleapis--mcp-toolbox
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat(tools/dataplex-update-data-product-aspects): Add dataplex-update-data-product-aspects tool (#3607)

Browse Source 
This PR adds the `dataplex-update-data-product-aspects` tool to allow
agent to update metadata aspects (such as documentation `overview` or
contract `refresh-cadence`) on existing Data Product Entries in Google
Cloud Dataplex (Knowledge Catalog). Additionally, it adds the
`dataplex-get-operation` tool to the `data-products` toolset.

Changes overview:

- **New Tool (`dataplex-update-data-product-aspects`)**:
- Implemented the tool with required parameters `locationId`,
`dataProductId`, and `aspects`.
- Configured detailed schema and payload documentation on the `aspects`
and `aspectSchema` parameters (`projectId`, `locationId`,
`aspectTypeId`, and specific fields under `data` for `overview` and
`refresh-cadence` aspects).
- **Documentation & Prebuilt Configurations**:
- Updated `internal/prebuiltconfigs/tools/dataplex.yaml` and
corresponding doc markdown to include detailed payload documentation for
`update_data_product_aspects` and registered the tool under the
`data-products` toolset.
- Updated `docs/en/integrations/knowledge-catalog/source.md` and related
tool references for `update_data_product_aspects` and `get_operation`.
- **Integration Tests**:
- Added unit and end-to-end integration tests for
`update_data_product_aspects`.

---------

Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
This commit is contained in:
Tejas Singh
2026-07-27 23:13:36 +05:30
committed by GitHub
parent 4f7688402b
commit f94093725c
19 changed files with 833 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/internal/config_test.go
+1 -1
View File
@@ -1984,7 +1984,7 @@ func TestPrebuiltTools(t *testing.T) {
},
"data-products": group.GroupConfig{
Name: "data-products",
ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "list_data_products", "get_data_product", "list_data_assets", "get_data_asset", "create_data_product", "update_data_product", "create_data_asset", "update_data_asset"},
ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "list_data_products", "get_data_product", "list_data_assets", "get_data_asset", "create_data_product", "update_data_product", "create_data_asset", "update_data_asset", "update_data_product_aspects", "get_operation"},
},
"enrich": group.GroupConfig{
Name: "enrich",
cmd/internal/imports.go
+1
View File
@@ -185,6 +185,7 @@ import (
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexsearchentries"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexupdatedataasset"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexupdatedataproduct"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexupdatedataproductaspects"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataproc/dataprocgetcluster"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataproc/dataprocgetjob"
_ "github.com/googleapis/mcp-toolbox/internal/tools/dataproc/dataproclistclusters"
docs/KNOWLEDGE_CATALOG_README.md
+18 -15
View File
@@ -50,26 +50,29 @@ Once configured, the MCP server will automatically provide Knowledge Catalog cap
* "Update the display name of Data Product 'my-product' to 'Updated Product'."
* "Create a new Data Asset under 'my-product' with resource '//bigquery.googleapis.com/projects/my-project/datasets/my-dataset/tables/my-table'."
* "Update the labels of Data Asset 'my-asset' under Data Product 'my-product' to have 'env: prod'."
* "Update aspects of Data Product 'my-product' to set overview content to 'New description'."
## Server Capabilities
The Knowledge Catalog MCP server provides the following tools:
| Tool Name | Description |
|:----------------------|:-----------------------------------------------------------------------------------------------------------------------------|
| `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. |
| `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. |
| `create_data_product` | Create a new Data Product. |
| `update_data_product` | Update an existing Data Product. |
| `create_data_asset` | Create a new Data Asset. |
| `update_data_asset` | Update an existing Data Asset. |
| Tool Name | Description |
|:------------------------------|:-----------------------------------------------------------------------------------------------------------------------------|
| `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. |
| `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. |
| `create_data_product` | Create a new Data Product. |
| `update_data_product` | Update an existing Data Product. |
| `create_data_asset` | Create a new Data Asset. |
| `update_data_asset` | Update an existing Data Asset. |
| `update_data_product_aspects` | Update aspects (like documentation or contracts) on a Data Product Entry. |
| `get_operation` | Retrieve the status of a long-running operation. |
## Custom MCP Server Configuration
docs/en/integrations/knowledge-catalog/prebuilt-configs/knowledge-catalog.md
+2 -1
View File
@@ -29,6 +29,7 @@ aliases:
* `update_data_product`: Updates an existing Data Product.
* `create_data_asset`: Creates a new Data Asset under a Data Product.
* `update_data_asset`: Updates an existing Data Asset under a Data Product.
* `update_data_product_aspects`: Updates aspects (like documentation or contracts) on a Data Product Entry.
* `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.
@@ -41,5 +42,5 @@ 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`, `get_data_product`, `list_data_assets`, `get_data_asset`, `create_data_product`, `update_data_product`, `create_data_asset`, `update_data_asset`).
* `data-products`: Data Products and Data Assets curation and management toolset (`search_entries`, `lookup_entry`, `search_aspect_types`, `lookup_context`, `list_data_products`, `get_data_product`, `list_data_assets`, `get_data_asset`, `create_data_product`, `update_data_product`, `create_data_asset`, `update_data_asset`, `update_data_product_aspects`, `get_operation`).
* `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`).
docs/en/integrations/knowledge-catalog/source.md
+22
View File
@@ -439,4 +439,26 @@ This abbreviated syntax works for the qualified predicates except for `label` in
### Response
1. Present the location ID and operation ID returned immediately by the tool call.
2. Poll the returned operation using the `get_operation` tool until completion at intervals of ~5 seconds.
## Tool: update_data_product_aspects
### Request
1. Use this tool to update metadata aspects (like `overview` or `contacts`) on an existing Data Product Entry.
2. You must provide `locationId`, `dataProductId`, and `aspects`.
3. The `aspects` parameter is an array of aspects. Each aspect contains:
- `projectId` (string, required): The project ID of the aspect type (use `"dataplex-types"` for system aspects).
- `locationId` (string, required): The location ID of the aspect type (use `"global"` for system aspects).
- `aspectTypeId` (string, required): The name of the aspect type (e.g. `"overview"` or `"refresh-cadence"`).
- `data` (object, required): The aspect payload details.
- For `"overview"` (documentation), the `data` object accepts: `content` (string, required, markdown or text), `contentType` (string, optional, `MARKDOWN` or `HTML`), and `links` (array of objects with `url` and `title`).
- For `"refresh-cadence"` (contract), the `data` object accepts: `frequency` (string, required: `Daily`, `Weekly`, `Monthly`, etc.), `refreshTime` (string, optional, e.g. `"09:00 PST"`), `thresholdInMinutes` (int, optional), and `cronSchedule` (string, optional).
### Response
1. Present the updated entry details including `name`, `entrySource`, `entryType`, and the list of updated `aspects` (each containing `projectId`, `locationId`, `aspectTypeId`, and `data`).
## Tool: get_operation
### Request
1. Use this tool to retrieve the status of an asynchronous long-running operation (LRO) like scan creation or data product/asset creation/updation.
2. You must provide `operationName`, a fully-qualified operation name string having format `projects/{projectId}/locations/{locationId}/operations/{operationId}`.
### Response
1. Present whether the operation is done (`done: true` or `done: false`).
2. Once the operation is completed (`done: true`), the response result will contain either an error (if it failed), or the created/updated resource (like `DataScan`, `DataProduct`, or `DataAsset`).
```
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-create-data-asset.md
+1 -1
View File
@@ -1,7 +1,7 @@
---
title: "dataplex-create-data-asset"
type: docs
weight: 2
weight: 1
description: >
A "dataplex-create-data-asset" tool creates a new Data Asset under an existing Data Product in Knowledge Catalog.
---
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-create-data-product.md
+1 -1
View File
@@ -1,7 +1,7 @@
---
title: "dataplex-create-data-product"
type: docs
weight: 2
weight: 1
description: >
A "dataplex-create-data-product" tool allows to create a new Data Product.
---
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-get-operation.md
+9 -8
View File
@@ -10,11 +10,12 @@ aliases:
## About
A `dataplex-get-operation` tool retrieves the status of a Dataplex long-running operation (LRO) like scan creation.
A `dataplex-get-operation` tool retrieves the status of a Dataplex long-running operation (LRO) (such as scan creation, data product creation/update, or data asset creation/update).
Poll this tool until the `done` field from the response is `true`. Once completed, the `response` field will contain the created DataScan resource, from which you can extract the `scanId` (the last part of the `name` field, e.g. `nq-doc-1234`) to pass to `get_run_status` and get results.
WARNING: This only tracks the creation of the scan template, NOT the actual background execution.
Poll this tool until the `done` field from the response is `true`. Once completed, the `response` field will contain the details of the created or updated resource (such as `DataScan`, `DataProduct`, or `DataAsset`).
For scan template creation, you can extract the `scanId` (the last part of the `name` field, e.g. `nq-doc-1234`) to pass to `get_run_status` and retrieve execution results.
WARNING: For scans, this only tracks the creation of the scan template, NOT the actual background execution.
## Compatible Sources
@@ -45,9 +46,9 @@ applying IAM permissions and roles to an identity.
The `dataplex-get-operation` tool accepts the following parameters:
| **field** | **type** | **required** | **description** |
| --------- | :------: | :----------: | --------------- |
| operationName | string | true | The full operation resource name (format: `projects/{project}/locations/{location}/operations/{operation_id}`). |
| **field** | **type** | **required** | **description** |
| ------------- | :------: | :----------: | ------------------------------------------------------------------------------------------------------------------ |
| operationName | string | true | The full operation resource name (format: `projects/{projectId}/locations/{locationId}/operations/{operationId}`). |
## Example
@@ -62,7 +63,7 @@ description: Check the status of a long-running scan template creation.
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:--------:|:------------:|----------------------------------------------------|
| type | string | true | Must be "dataplex-get-operation". |
| ----------- | :------: | :----------: | -------------------------------------------------- |
| type | string | true | Must be "dataplex-get-operation". |
| 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. |
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-update-data-asset.md
+1 -1
View File
@@ -1,7 +1,7 @@
---
title: "dataplex-update-data-asset"
type: docs
weight: 2
weight: 1
description: >
A "dataplex-update-data-asset" tool updates an existing Data Asset under a Data Product in Knowledge Catalog.
---
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-update-data-product-aspects.md
+114
View File
@@ -0,0 +1,114 @@
---
title: "dataplex-update-data-product-aspects"
type: docs
weight: 1
description: >
A "dataplex-update-data-product-aspects" tool updates aspects for an existing Data Product Entry in Knowledge Catalog.
---
## About
A `dataplex-update-data-product-aspects` tool updates aspects on an existing Data Product Entry in Knowledge Catalog (formerly known as Dataplex). This tool operates on the catalog entry associated with the Data Product, allowing you to add or modify metadata aspects in a single request.
View the [Data Products guide][guide] for more information.
[guide]: https://docs.cloud.google.com/dataplex/docs/data-products-overview
## 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
## Parameters
The `dataplex-update-data-product-aspects` tool accepts the following parameters:
| **field** | **type** | **required** | **description** |
| ------------- | ---------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| locationId | string | true | The location ID (e.g. `us`, `us-central1`) of the Data Product. |
| dataProductId | string | true | The unique ID of the Data Product. |
| aspects | array of objects | true | The list of aspects to add or update on the Data Product Entry. Each object contains `projectId` (string, required), `locationId` (string, required), `aspectTypeId` (string, required), and `data` (object, optional, JSON map of the aspect details). |
### Aspect Object Fields
Each object in the `aspects` array must have the following fields:
* **`projectId`**: The project ID of the aspect type. The value is `"dataplex-types"` for system aspects (like `overview` and `refresh-cadence`).
* **`locationId`**: The location of the aspect type. The value is `"global"` for system aspects (like `overview` and `refresh-cadence`).
* **`aspectTypeId`**: The unique name of the aspect type (e.g. `"overview"` or `"refresh-cadence"`).
* **`data`**: The JSON payload conforming to the aspect type's schema:
* **System `overview` (Documentation) Schema**:
* `content` (string, required): The detailed documentation content (markdown or plain text).
* `contentType` (string, optional): The content type format. Values: `MARKDOWN`, `HTML`.
* `links` (array of objects, optional): List of relevant URL links. Each link object contains `url` (string, required) and `title` (string, optional).
* **System `refresh-cadence` (Contract) Schema**:
* `frequency` (string, required): How often the data is updated. Values: `Daily`, `Weekly`, `Monthly`, `Quarterly`, `Half-Yearly`, `Yearly`.
* `refreshTime` (string, optional): Time of day when the data is updated (e.g. `"09:00 PST"`).
* `thresholdInMinutes` (int, optional): Delinquency threshold in minutes (e.g. `15`).
* `cronSchedule` (string, optional): Optional cron schedule expression (e.g. `"0 * * * *"`).
## Example
```yaml
kind: tool
name: update_data_product_aspects
type: dataplex-update-data-product-aspects
source: my-dataplex-source
description: Use this tool to update aspects (like overview or contacts) on a Data Product Entry.
```
## Output Format
The tool returns the updated catalog entry details and aspects:
```json
{
"locationId": "us",
"dataProductId": "my-data-product",
"entrySource": {
"resource": "projects/...",
"displayName": "My Data Product Resource",
"description": "Resource details"
},
"entryType": "projects/dataplex-types/locations/global/entryTypes/data-product",
"aspects": [
{
"projectId": "dataplex-types",
"locationId": "global",
"aspectTypeId": "overview",
"data": {
"content": "Updated description content"
}
}
]
}
```
## Reference
| **field** | **type** | **required** | **description** |
| ----------- | -------- | ------------ | ---------------------------------------------------- |
| type | string | true | Must be "dataplex-update-data-product-aspects". |
| 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. |
docs/en/integrations/knowledge-catalog/tools/knowledge-catalog-update-data-product.md
+1 -1
View File
@@ -1,7 +1,7 @@
---
title: "dataplex-update-data-product"
type: docs
weight: 2
weight: 1
description: >
A "dataplex-update-data-product" tool updates metadata for an existing Data Product in Knowledge Catalog.
---
go.mod
+1
View File
@@ -18,6 +18,7 @@ require (
cloud.google.com/go/geminidataanalytics v1.2.0
cloud.google.com/go/logging v1.18.0
cloud.google.com/go/longrunning v1.0.0
cloud.google.com/go/resourcemanager v1.15.0
cloud.google.com/go/spanner v1.92.0
cloud.google.com/go/storage v1.62.3
github.com/ClickHouse/clickhouse-go/v2 v2.46.0
go.sum
+2
View File
@@ -37,6 +37,8 @@ cloud.google.com/go/longrunning v1.0.0 h1:lwzWEYD8+NkYV7dhexOz6kmlvajZA70+bW/xMh
cloud.google.com/go/longrunning v1.0.0/go.mod h1:8nqFBPOO1U/XkhWl0I19AMZEphrHi73VNABIpKYaTwM=
cloud.google.com/go/monitoring v1.29.0 h1:AHhDsFaSax1/4k+qlIDX/SDGe6hggnfXJ9dkgD9qBPY=
cloud.google.com/go/monitoring v1.29.0/go.mod h1:72NOVjJXHY/HBfoLT0+qlCZBT059+9VXLeAnL2PeeVM=
cloud.google.com/go/resourcemanager v1.15.0 h1:OwcTLrKaly0SMPoYHssPG4FBzRF0tyimeySOFD/YPJ0=
cloud.google.com/go/resourcemanager v1.15.0/go.mod h1:ve0VNxPoDU6XxDuEMCjkineb0YzXQXx3mOWwnNckGDE=
cloud.google.com/go/spanner v1.92.0 h1:cfeMNmtFjz+OYzQVCIuGBw4Cik4CbF2ptXMuRQcUar0=
cloud.google.com/go/spanner v1.92.0/go.mod h1:rCDPfWXNX0h+t484r+crCEaaMKbJfoWkHRDKU3H3+oY=
cloud.google.com/go/storage v1.62.3 h1:SZq1t23NCI+e96dH77Dg3PEfsNNEjqO8zE5AnD8gVD0=
internal/prebuiltconfigs/tools/dataplex.yaml
+10 -6
View File
@@ -96,6 +96,12 @@ source: dataplex-source
description: Updates an existing Data Asset under a Data Product.
---
kind: tool
name: update_data_product_aspects
type: dataplex-update-data-product-aspects
source: dataplex-source
description: Updates metadata aspects (like overview or contacts) on an existing Data Product Entry.
---
kind: tool
name: generate_data_insights
type: dataplex-generate-data-insights
source: dataplex-source
@@ -134,12 +140,8 @@ name: get_operation
type: dataplex-get-operation
source: dataplex-source
description: >-
Retrieves the status of a Dataplex long-running operation (LRO) like scan creation.
Poll this tool until the 'done' field from the tool's response is true. Once completed, the 'response' field will
contain the created DataScan resource, from which you can extract the 'scanId'
(the last part of the 'name' field, e.g. 'nq-doc-1234') to pass to get_run_status
and get_data_insights. WARNING: This only tracks the creation of the scan, NOT
its execution.
Retrieves the status of a Dataplex long-running operation (LRO) (such as scan creation, data product creation/update, or data asset creation/update).
Poll this tool until the 'done' field from the tool's response is true. Once completed, the 'response' field contains the details of the created or updated resource (such as DataScan, DataProduct, or DataAsset). For scan template creation, you can extract the 'scanId' from the returned DataScan resource to pass to get_run_status. WARNING: For scans, this only tracks the creation of the scan template, NOT its background execution.
---
kind: tool
name: get_run_status
@@ -294,6 +296,8 @@ tools:
- update_data_product
- create_data_asset
- update_data_asset
- update_data_product_aspects
- get_operation
---
kind: toolset
name: enrich
internal/sources/dataplex/dataplex.go
+74 -9
View File
@@ -19,11 +19,14 @@ import (
"encoding/json"
"fmt"
"regexp"
"strconv"
"strings"
dataplexapi "cloud.google.com/go/dataplex/apiv1"
"cloud.google.com/go/dataplex/apiv1/dataplexpb"
"cloud.google.com/go/longrunning/autogen/longrunningpb"
resourcemanager "cloud.google.com/go/resourcemanager/apiv3"
resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb"
"github.com/cenkalti/backoff/v6"
"github.com/goccy/go-yaml"
"github.com/google/uuid"
@@ -79,15 +82,47 @@ func (r Config) SourceConfigType() string {
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
// Initializes a Dataplex source
client, dataScanClient, dataProductClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project, r.ImpersonateServiceAccount, r.Scopes)
client, dataScanClient, dataProductClient, projectsClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project, r.ImpersonateServiceAccount, r.Scopes)
if err != nil {
return nil, err
}
// Resolve project number
proj, err := projectsClient.GetProject(ctx, &resourcemanagerpb.GetProjectRequest{
Name: "projects/" + r.Project,
})
if err != nil {
client.Close()
dataScanClient.Close()
dataProductClient.Close()
projectsClient.Close()
return nil, fmt.Errorf("failed to get project details for project %q: %w", r.Project, err)
}
parts := strings.Split(proj.Name, "/")
if len(parts) < 2 {
client.Close()
dataScanClient.Close()
dataProductClient.Close()
projectsClient.Close()
return nil, fmt.Errorf("unexpected project resource name format: %q", proj.Name)
}
projectNumberStr := parts[1]
projectNumber, err := strconv.ParseInt(projectNumberStr, 10, 64)
if err != nil {
client.Close()
dataScanClient.Close()
dataProductClient.Close()
projectsClient.Close()
return nil, fmt.Errorf("failed to parse project number %q as int64: %w", projectNumberStr, err)
}
s := &Source{
Config: r,
Client: client,
DataScanClient: dataScanClient,
dataProductClient: dataProductClient,
projectsClient: projectsClient,
projectNumber: projectNumber,
}
return s, nil
@@ -100,6 +135,8 @@ type Source struct {
Client *dataplexapi.CatalogClient
DataScanClient *dataplexapi.DataScanClient
dataProductClient *dataplexapi.DataProductClient
projectsClient *resourcemanager.ProjectsClient
projectNumber int64
}
func (s *Source) SourceType() string {
@@ -115,6 +152,14 @@ func (s *Source) ProjectID() string {
return s.Project
}
func (s *Source) ProjectNumber() int64 {
return s.projectNumber
}
func (s *Source) ProjectsClient() *resourcemanager.ProjectsClient {
return s.projectsClient
}
func (s *Source) CatalogClient() *dataplexapi.CatalogClient {
return s.Client
}
@@ -134,13 +179,13 @@ func initDataplexConnection(
project string,
impersonateServiceAccount string,
scopes []string,
) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, *dataplexapi.DataProductClient, error) {
) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, *dataplexapi.DataProductClient, *resourcemanager.ProjectsClient, error) {
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)
if err != nil {
return nil, nil, nil, err
return nil, nil, nil, nil, err
}
var opts []option.ClientOption
@@ -157,7 +202,7 @@ func initDataplexConnection(
Scopes: credScopes,
})
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to create impersonated credentials for %q for project %q: %w", impersonateServiceAccount, project, err)
return nil, nil, nil, nil, fmt.Errorf("failed to create impersonated credentials for %q for project %q: %w", impersonateServiceAccount, project, err)
}
opts = []option.ClientOption{
option.WithUserAgent(userAgent),
@@ -167,7 +212,7 @@ func initDataplexConnection(
// Use default credentials
cred, err := google.FindDefaultCredentials(ctx, credScopes...)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials for project %q: %w", project, err)
return nil, nil, nil, nil, fmt.Errorf("failed to find default Google Cloud credentials for project %q: %w", project, err)
}
opts = []option.ClientOption{
option.WithUserAgent(userAgent),
@@ -177,19 +222,31 @@ func initDataplexConnection(
client, err := dataplexapi.NewCatalogClient(ctx, opts...)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
return nil, nil, nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
}
dataScanClient, err := dataplexapi.NewDataScanClient(ctx, opts...)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to create Dataplex DataScan client for project %q: %w", project, err)
client.Close()
return nil, nil, nil, nil, fmt.Errorf("failed to create Dataplex DataScan client for project %q: %w", project, err)
}
dataProductClient, err := dataplexapi.NewDataProductClient(ctx, opts...)
if err != nil {
return nil, nil, nil, fmt.Errorf("failed to create Dataplex DataProduct client for project %q: %w", project, err)
client.Close()
dataScanClient.Close()
return nil, nil, nil, nil, fmt.Errorf("failed to create Dataplex DataProduct client for project %q: %w", project, err)
}
return client, dataScanClient, dataProductClient, nil
projectsClient, err := resourcemanager.NewProjectsClient(ctx, opts...)
if err != nil {
client.Close()
dataScanClient.Close()
dataProductClient.Close()
return nil, nil, nil, nil, fmt.Errorf("failed to create ResourceManager projects client for project %q: %w", project, err)
}
return client, dataScanClient, dataProductClient, projectsClient, nil
}
func (s *Source) LookupEntry(ctx context.Context, name string, view int, aspectTypes []string, entry string) (*dataplexpb.Entry, error) {
@@ -787,6 +844,14 @@ func (s *Source) UpdateDataAsset(
}, nil
}
func (s *Source) UpdateEntry(ctx context.Context, entry *dataplexpb.Entry, updateMask *fieldmaskpb.FieldMask) (*dataplexpb.Entry, error) {
req := &dataplexpb.UpdateEntryRequest{
Entry: entry,
UpdateMask: updateMask,
}
return s.CatalogClient().UpdateEntry(ctx, req)
}
func (s *Source) GenerateDataInsights(ctx context.Context, location, resourcePath string, publish bool) (string, error) {
parent := fmt.Sprintf("projects/%s/locations/%s", s.ProjectID(), location)
dataScanID := fmt.Sprintf("nq-doc-%s", uuid.New().String())
internal/tools/dataplex/dataplexgetoperation/dataplexgetoperation.go
+1 -1
View File
@@ -61,7 +61,7 @@ func (cfg Config) ToolConfigType() string {
}
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
operationName := parameters.NewStringParameter("operationName", "The fully-qualified resource name of the operation returned by generate_data_insights. Format: projects/{project}/locations/{location}/operations/{operation_id}.")
operationName := parameters.NewStringParameter("operationName", "The fully-qualified resource name of the operation returned by generate_data_insights. Format: projects/{projectId}/locations/{locationId}/operations/{operationId}.")
allParameters := parameters.Parameters{operationName}
internal/tools/dataplex/dataplexupdatedataproductaspects/dataplexupdatedataproductaspects.go
+271
View File
@@ -0,0 +1,271 @@
// 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 dataplexupdatedataproductaspects
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
"github.com/goccy/go-yaml"
"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"
"google.golang.org/protobuf/types/known/fieldmaskpb"
"google.golang.org/protobuf/types/known/structpb"
)
const resourceType string = "dataplex-update-data-product-aspects"
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{ConfigBase: tools.ConfigBase{Name: name}}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type compatibleSource interface {
ProjectID() string
ProjectNumber() int64
UpdateEntry(ctx context.Context, entry *dataplexpb.Entry, updateMask *fieldmaskpb.FieldMask) (*dataplexpb.Entry, error)
}
type Config struct {
tools.ConfigBase `yaml:",inline"`
Type string `yaml:"type" validate:"required"`
Source string `yaml:"source" validate:"required"`
Parameters parameters.Parameters `yaml:"parameters"`
Annotations *tools.ToolAnnotations `yaml:"annotations,omitempty"`
}
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigType() string {
return resourceType
}
type Aspect struct {
ProjectID string `json:"projectId" validate:"required"`
LocationID string `json:"locationId" validate:"required"`
AspectTypeId string `json:"aspectTypeId" validate:"required"`
Data map[string]any `json:"data,omitempty"`
}
type EntrySource struct {
Resource string `json:"resource,omitempty"`
DisplayName string `json:"displayName,omitempty"`
Description string `json:"description,omitempty"`
}
type UpdateDataProductAspectsResponse struct {
Name string `json:"name"`
EntrySource *EntrySource `json:"entrySource,omitempty"`
EntryType string `json:"entryType,omitempty"`
Aspects []Aspect `json:"aspects"`
}
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
locationID := parameters.NewStringParameter("locationId", "The location ID (e.g. 'us', 'us-central1') of the Data Product.")
dataProductID := parameters.NewStringParameter("dataProductId", "The unique ID of the Data Product.")
aspectSchema := parameters.NewMapParameter(
"aspect",
"Aspect details containing: projectId (string, required, 'dataplex-types' for system aspects), locationId (string, required, 'global' for system aspects), aspectTypeId (string, required, e.g. 'overview' or 'refresh-cadence'), and data (object, required, the aspect payload details. For 'overview' (documentation), data accepts: content (string, required, markdown or text), contentType (string, optional, MARKDOWN or HTML), and links (array of objects with url and title). For 'refresh-cadence' (contract), data accepts: frequency (string, required: Daily, Weekly, Monthly, etc.), refreshTime (string, optional, e.g. '09:00 PST'), thresholdInMinutes (int, optional), and cronSchedule (string, optional)).",
"",
)
aspects := parameters.NewArrayParameter(
"aspects",
"The list of aspects to update on the Data Product Entry.",
aspectSchema,
)
allParameters := parameters.Parameters{locationID, dataProductID, aspects}
return Tool{
BaseTool: tools.NewBaseTool(
cfg,
tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewDestructiveAnnotations),
tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
allParameters,
),
}, nil
}
// validate interface
var _ tools.Tool = Tool{}
type Tool struct {
tools.BaseTool[Config]
}
func (t Tool) GetSourceName() string {
return t.Cfg.Source
}
func (t Tool) ToConfig() tools.ToolConfig {
return t.Cfg
}
func (t Tool) ValidateSource(source sources.Source) error {
_, ok := source.(compatibleSource)
if !ok {
return fmt.Errorf("invalid source for %q tool: source %q is not a compatible type", t.Cfg.Type, t.Cfg.Source)
}
return nil
}
func (t Tool) Invoke(ctx context.Context, s sources.Source, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, ok := s.(compatibleSource)
if !ok {
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, nil)
}
paramsMap := params.AsMap()
locationID, _ := paramsMap["locationId"].(string)
if locationID == "" {
return nil, util.NewAgentError("locationId parameter is required and must be a non-empty string", nil)
}
dataProductID, _ := paramsMap["dataProductId"].(string)
if dataProductID == "" {
return nil, util.NewAgentError("dataProductId parameter is required and must be a non-empty string", nil)
}
rawAspects, ok := paramsMap["aspects"].([]any)
if !ok {
return nil, util.NewAgentError("aspects parameter is required and must be an array", nil)
}
rawAspectsBytes, err := json.Marshal(rawAspects)
if err != nil {
return nil, util.NewAgentError("failed to marshal aspects parameter", err)
}
var parsedAspects []Aspect
if err := json.Unmarshal(rawAspectsBytes, &parsedAspects); err != nil {
return nil, util.NewAgentError("failed to unmarshal aspects parameter into required format", err)
}
projectID := source.ProjectID()
// Convert input array of aspects to aspects map
aspectsMap := make(map[string]*dataplexpb.Aspect)
for i, aspect := range parsedAspects {
if aspect.AspectTypeId == "" {
return nil, util.NewAgentError(fmt.Sprintf("aspectTypeId is required for aspect at index %d", i), nil)
}
if aspect.Data == nil {
return nil, util.NewAgentError(fmt.Sprintf("data is required for aspect at index %d", i), nil)
}
aspectProjID := aspect.ProjectID
if aspectProjID == "" {
if aspect.AspectTypeId == "overview" || aspect.AspectTypeId == "refresh-cadence" {
aspectProjID = "dataplex-types"
} else {
aspectProjID = projectID
}
}
aspectLocID := aspect.LocationID
if aspectLocID == "" {
if aspect.AspectTypeId == "overview" || aspect.AspectTypeId == "refresh-cadence" {
aspectLocID = "global"
} else {
aspectLocID = locationID
}
}
aspectType := fmt.Sprintf("projects/%s/locations/%s/aspectTypes/%s", aspectProjID, aspectLocID, aspect.AspectTypeId)
aspectKey := fmt.Sprintf("%s.%s.%s", aspectProjID, aspectLocID, aspect.AspectTypeId)
structData, err := structpb.NewStruct(aspect.Data)
if err != nil {
return nil, util.NewAgentError(fmt.Sprintf("failed to serialize data for aspect %q: %s", aspect.AspectTypeId, err), err)
}
aspectsMap[aspectKey] = &dataplexpb.Aspect{
AspectType: aspectType,
Data: structData,
}
}
entryName := fmt.Sprintf(
"projects/%s/locations/%s/entryGroups/@dataplex/entries/projects/%d/locations/%s/dataProducts/%s",
projectID, locationID, source.ProjectNumber(), locationID, dataProductID,
)
entry := &dataplexpb.Entry{
Name: entryName,
Aspects: aspectsMap,
}
updateMask, _ := fieldmaskpb.New(entry, "aspects")
returnedEntry, err := source.UpdateEntry(ctx, entry, updateMask)
if err != nil {
return nil, util.ProcessGcpError(err)
}
// Format returned entry aspects back to the expected output Aspects format
var returnedAspects []Aspect
for _, aspectProto := range returnedEntry.Aspects {
parts := strings.Split(aspectProto.AspectType, "/")
if len(parts) < 6 {
continue
}
aspectProjID := parts[1]
aspectLocID := parts[3]
aspectTypeName := parts[5]
data := aspectProto.Data.AsMap()
returnedAspects = append(returnedAspects, Aspect{
AspectTypeId: aspectTypeName,
Data: data,
ProjectID: aspectProjID,
LocationID: aspectLocID,
})
}
var entrySource *EntrySource
if returnedEntry.GetEntrySource() != nil {
entrySource = &EntrySource{
Resource: returnedEntry.GetEntrySource().GetResource(),
DisplayName: returnedEntry.GetEntrySource().GetDisplayName(),
Description: returnedEntry.GetEntrySource().GetDescription(),
}
}
return UpdateDataProductAspectsResponse{
Name: returnedEntry.GetName(),
EntrySource: entrySource,
EntryType: returnedEntry.GetEntryType(),
Aspects: returnedAspects,
}, nil
}
internal/tools/dataplex/dataplexupdatedataproductaspects/dataplexupdatedataproductaspects_test.go
+71
View File
@@ -0,0 +1,71 @@
// 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 dataplexupdatedataproductaspects_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"
"github.com/googleapis/mcp-toolbox/internal/tools/dataplex/dataplexupdatedataproductaspects"
)
func TestParseFromYamlDataplexUpdateDataProductAspects(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-update-data-product-aspects
source: my-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": dataplexupdatedataproductaspects.Config{
ConfigBase: tools.ConfigBase{
Name: "example_tool",
Description: "some description",
AuthRequired: []string{},
},
Type: "dataplex-update-data-product-aspects",
Source: "my-instance",
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Parse contents
_, _, _, got, _, _, err := server.UnmarshalPrimitiveConfig(ctx, testutils.FormatYaml(tc.in))
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Errorf("unexpected diff (-want +got):\n%s", diff)
}
})
}
}
tests/dataplex/dataplex_integration_test.go
+232 -198
View File
@@ -48,31 +48,32 @@ import (
)
var (
DataplexSourceType = "dataplex"
DataplexLookupContextToolType = "dataplex-lookup-context"
DataplexSearchEntriesToolType = "dataplex-search-entries"
DataplexLookupEntryToolType = "dataplex-lookup-entry"
DataplexSearchAspectTypesToolType = "dataplex-search-aspect-types"
DataplexSearchDataQualityScansToolType = "dataplex-search-dq-scans"
DataplexListDataProductsToolType = "dataplex-list-data-products"
DataplexGetDataProductToolType = "dataplex-get-data-product"
DataplexListDataAssetsToolType = "dataplex-list-data-assets"
DataplexGetDataAssetToolType = "dataplex-get-data-asset"
DataplexCreateDataProductToolType = "dataplex-create-data-product"
DataplexCreateDataAssetToolType = "dataplex-create-data-asset"
DataplexUpdateDataAssetToolType = "dataplex-update-data-asset"
DataplexUpdateDataProductToolType = "dataplex-update-data-product"
DataplexGenerateDataProfileToolType = "dataplex-generate-data-profile"
DataplexGetDataProfileToolType = "dataplex-get-data-profile"
DataplexGetOperationToolType = "dataplex-get-operation"
DataplexGetRunStatusToolType = "dataplex-get-run-status"
DataplexGenerateDataInsightsToolType = "dataplex-generate-data-insights"
DataplexGetDataInsightsToolType = "dataplex-get-data-insights"
DataplexDiscoverMetadataToolType = "dataplex-discover-metadata"
DataplexGetDiscoveryResultsToolType = "dataplex-get-discovery-results"
DataplexCheckDataQualityToolType = "dataplex-check-data-quality"
DataplexGetDataQualityResultsToolType = "dataplex-get-data-quality-results"
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
DataplexSourceType = "dataplex"
DataplexLookupContextToolType = "dataplex-lookup-context"
DataplexSearchEntriesToolType = "dataplex-search-entries"
DataplexLookupEntryToolType = "dataplex-lookup-entry"
DataplexSearchAspectTypesToolType = "dataplex-search-aspect-types"
DataplexSearchDataQualityScansToolType = "dataplex-search-dq-scans"
DataplexListDataProductsToolType = "dataplex-list-data-products"
DataplexGetDataProductToolType = "dataplex-get-data-product"
DataplexListDataAssetsToolType = "dataplex-list-data-assets"
DataplexGetDataAssetToolType = "dataplex-get-data-asset"
DataplexCreateDataProductToolType = "dataplex-create-data-product"
DataplexUpdateDataProductToolType = "dataplex-update-data-product"
DataplexCreateDataAssetToolType = "dataplex-create-data-asset"
DataplexUpdateDataAssetToolType = "dataplex-update-data-asset"
DataplexUpdateDataProductAspectsToolType = "dataplex-update-data-product-aspects"
DataplexGenerateDataProfileToolType = "dataplex-generate-data-profile"
DataplexGetDataProfileToolType = "dataplex-get-data-profile"
DataplexGetOperationToolType = "dataplex-get-operation"
DataplexGetRunStatusToolType = "dataplex-get-run-status"
DataplexGenerateDataInsightsToolType = "dataplex-generate-data-insights"
DataplexGetDataInsightsToolType = "dataplex-get-data-insights"
DataplexDiscoverMetadataToolType = "dataplex-discover-metadata"
DataplexGetDiscoveryResultsToolType = "dataplex-get-discovery-results"
DataplexCheckDataQualityToolType = "dataplex-check-data-quality"
DataplexGetDataQualityResultsToolType = "dataplex-get-data-quality-results"
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
)
func getDataplexVars(t *testing.T) map[string]any {
@@ -484,6 +485,7 @@ func TestDataplexToolEndpoints(t *testing.T) {
runDataplexGetDataAssetToolInvokeTest(t, dataProductId1, dataAssetId1)
runDataplexCreateAndUpdateDataProductToolsInvokeTest(t, dataplexDataProductClient, dataProductId3, dataProductId4)
runDataplexCreateAndUpdateDataAssetToolsInvokeTest(t, dataplexDataProductClient, dataProductId2, dataAssetId2, dataAssetId3, datasetName1, tableName1, datasetName2, tableName2)
runDataplexUpdateDataProductAspectsToolInvokeTest(t, dataProductId1, aspectTypeId)
runDataplexEnrichmentToolInvokeTest(t, tableName1, datasetName1, bucketName, dataplexDataScanClient)
}
@@ -710,9 +712,12 @@ func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client
},
},
}
_, err := client.CreateAspectType(ctx, createAspectTypeReq)
op, err := client.CreateAspectType(ctx, createAspectTypeReq)
if err != nil {
t.Fatalf("Failed to create aspect type %s: %v", aspectTypeId, err)
t.Fatalf("Failed to initiate CreateAspectType %s: %v", aspectTypeId, err)
}
if _, err := op.Wait(ctx); err != nil {
t.Fatalf("Failed to wait for CreateAspectType %s: %v", aspectTypeId, err)
}
return func(t *testing.T) {
@@ -885,6 +890,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
"description": "Simple dataplex update data asset tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
"my-dataplex-update-data-product-aspects-tool": map[string]any{
"type": DataplexUpdateDataProductAspectsToolType,
"source": "my-dataplex-instance",
"description": "Simple dataplex update data product aspects tool to test end to end functionality.",
},
"my-auth-dataplex-update-data-product-aspects-tool": map[string]any{
"type": DataplexUpdateDataProductAspectsToolType,
"source": "my-dataplex-instance",
"description": "Simple dataplex update data product aspects tool to test end to end functionality.",
"authRequired": []string{"my-google-auth"},
},
"my-dataplex-generate-data-profile-tool": map[string]any{
"type": DataplexGenerateDataProfileToolType,
"source": "my-dataplex-instance",
@@ -1007,6 +1023,11 @@ func runDataplexToolGetTest(t *testing.T) {
toolName: "my-dataplex-update-data-asset-tool",
expectedParams: []string{"locationId", "dataProductId", "dataAssetId", "labels", "accessGroupConfigs", "updateMask"},
},
{
name: "get my-dataplex-update-data-product-aspects-tool",
toolName: "my-dataplex-update-data-product-aspects-tool",
expectedParams: []string{"locationId", "dataProductId", "aspects"},
},
{
name: "get my-dataplex-generate-data-profile-tool",
toolName: "my-dataplex-generate-data-profile-tool",
@@ -1197,26 +1218,12 @@ func runDataplexSearchEntriesToolInvokeTest(t *testing.T, tableName string, data
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != 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))
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 {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -1363,27 +1370,13 @@ func runDataplexLookupEntryToolInvokeTest(t *testing.T, tableName string, datase
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
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 {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("Error parsing response body: %v", err)
}
@@ -1492,24 +1485,12 @@ func runDataplexSearchAspectTypesToolInvokeTest(t *testing.T, aspectTypeId strin
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
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 {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -1614,23 +1595,7 @@ func runDataplexLookupContextToolInvokeTest(t *testing.T, tableName string, data
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()
bodyBytes, _ := io.ReadAll(resp.Body)
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("Response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
@@ -1748,24 +1713,12 @@ func runDataplexSearchDataQualityScansToolInvokeTest(t *testing.T, dataScanId st
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
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 {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -1859,27 +1812,15 @@ func runDataplexListDataProductsToolInvokeTest(t *testing.T, dataProductId1 stri
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -1982,27 +1923,15 @@ func runDataplexGetDataProductToolInvokeTest(t *testing.T, dataProductId string)
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -2098,27 +2027,15 @@ func runDataplexListDataAssetsToolInvokeTest(t *testing.T, dataProductId string,
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -2232,27 +2149,16 @@ func runDataplexGetDataAssetToolInvokeTest(t *testing.T, dataProductId string, d
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -2426,28 +2332,15 @@ func runDataplexCreateAndUpdateDataProductToolsInvokeTest(t *testing.T, client *
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -2638,28 +2531,16 @@ func runDataplexCreateAndUpdateDataAssetToolsInvokeTest(
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()
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(bodyBytes, &result); err != nil {
t.Fatalf("error parsing response body: %s", err)
}
resultStr, ok := result["result"].(string)
@@ -2953,3 +2834,156 @@ func runDataplexEnrichmentToolInvokeTest(t *testing.T, tableName string, dataset
})
}
}
func runDataplexUpdateDataProductAspectsToolInvokeTest(t *testing.T, dataProductId string, aspectTypeId string) {
idToken, err := tests.GetGoogleIdToken(t)
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
expectedAspectTypeId string
expectedAspectValue string
}{
{
name: "Success - Update Data Product Aspects (Authorized)",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-update-data-product-aspects-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": idToken},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(
`{"locationId":"us","dataProductId":"%s","aspects":[{"projectId":"dataplex-types","locationId":"global","aspectTypeId":"overview","data":{"content":"Auth updated description"}}]}`,
dataProductId,
))),
wantStatusCode: 200,
expectResult: true,
expectedAspectTypeId: "overview",
expectedAspectValue: "Auth updated description",
},
{
name: "Success - Update Data Product Aspects (Un-authorized)",
api: "http://127.0.0.1:5000/api/tool/my-dataplex-update-data-product-aspects-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(
`{"locationId":"us","dataProductId":"%s","aspects":[{"projectId":"%s","locationId":"us","aspectTypeId":"%s","data":{}}]}`,
dataProductId, DataplexProject, aspectTypeId,
))),
wantStatusCode: 200,
expectResult: true,
expectedAspectTypeId: aspectTypeId,
expectedAspectValue: "",
},
{
name: "Failure - Without Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-update-data-product-aspects-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(
`{"locationId":"us","dataProductId":"%s","aspects":[{"projectId":"dataplex-types","locationId":"global","aspectTypeId":"overview","data":{"content":"Unauth"}}]}`,
dataProductId,
))),
wantStatusCode: 401,
expectResult: false,
},
{
name: "Failure - Invalid Authorization Token",
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-update-data-product-aspects-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": "invalid_token"},
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(
`{"locationId":"us","dataProductId":"%s","aspects":[{"projectId":"dataplex-types","locationId":"global","aspectTypeId":"overview","data":{"content":"Invalid"}}]}`,
dataProductId,
))),
wantStatusCode: 401,
expectResult: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
resp, bodyBytes := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("response status code got %d, want %d\nResponse body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
}
if !tc.expectResult {
return
}
var result map[string]interface{}
if err := json.Unmarshal(bodyBytes, &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 invokeResp map[string]interface{}
if err := json.Unmarshal([]byte(resultStr), &invokeResp); err != nil {
t.Fatalf("error unmarshalling result string: %v", err)
}
aspectsList, ok := invokeResp["aspects"].([]interface{})
if !ok {
t.Fatalf("expected 'aspects' field to be an array, got %T. Full response: %s", invokeResp["aspects"], resultStr)
}
var foundAspect bool
for _, rawAspect := range aspectsList {
aspect, ok := rawAspect.(map[string]interface{})
if !ok {
continue
}
typeId, _ := aspect["aspectTypeId"].(string)
if typeId == tc.expectedAspectTypeId {
projId, _ := aspect["projectId"].(string)
locId, _ := aspect["locationId"].(string)
data, _ := aspect["data"].(map[string]interface{})
if typeId == "overview" {
if data != nil {
desc, _ := data["content"].(string)
if desc == tc.expectedAspectValue {
if projId != "dataplex-types" {
t.Errorf("expected aspect projectId to be 'dataplex-types', got %q", projId)
}
if locId != "global" {
t.Errorf("expected aspect locationId to be 'global', got %q", locId)
}
foundAspect = true
break
}
}
} else {
// Custom aspect type validations
if projId != DataplexProject {
t.Errorf("expected aspect projectId to be %q, got %q", DataplexProject, projId)
}
if locId != "us" {
t.Errorf("expected aspect locationId to be 'us', got %q", locId)
}
foundAspect = true
break
}
}
}
if !foundAspect {
t.Fatalf("expected aspect %q with value %q in returned aspects, but not found. Response: %v", tc.expectedAspectTypeId, tc.expectedAspectValue, invokeResp)
}
entryType, _ := invokeResp["entryType"].(string)
if entryType == "" {
t.Errorf("expected non-empty 'entryType' in response. Response: %v", invokeResp)
}
if !strings.Contains(entryType, "/entryTypes/") {
t.Errorf("expected 'entryType' to contain '/entryTypes/', got %q", entryType)
}
if _, ok := invokeResp["entrySource"]; !ok {
t.Errorf("expected 'entrySource' key in response. Response: %v", invokeResp)
}
})
}
}