feat: [Dataplex Tool] Search Data Quality Scans

cmd/root.go

@@ -113,6 +113,7 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/dataform/dataformcompilelocal"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchdqscans"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/elasticsearch/elasticsearchesql"

cmd/root_test.go

@@ -2320,7 +2320,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"dataplex_tools": tools.ToolsetConfig{
 					Name:      "dataplex_tools",
-					ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types"},
+					ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "search_dq_scans"},
 				},
 			},
 		},

docs/en/resources/tools/dataplex/dataplex-search-dq-scans.md

@@ -0,0 +1,64 @@
+---
+title: "dataplex-search-dq-scans"
+type: docs
+weight: 1
+description: >
+  A "dataplex-search-dq-scans" tool allows to search for data quality scans based on the provided parameters.
+aliases:
+- /resources/tools/dataplex-search-dq-scans
+---
+
+## About
+
+A `dataplex-search-dq-scans` tool returns data quality scans that match the given criteria.
+It's compatible with the following sources:
+
+- [dataplex](../../sources/dataplex.md)
+
+`dataplex-search-dq-scans` accepts the following optional parameters:
+
+- `filter` - Filter string to search/filter data quality scans. E.g. "display_name = \"my-scan\"".
+- `data_scan_id` - The ID of the data scan to filter by.
+- `table_name` - The name of the table to filter by.
+- `pageSize` - Number of returned data quality scans in the page. Defaults to `10`.
+- `orderBy` - Specifies the ordering of results.
+
+## Requirements
+
+### IAM Permissions
+
+Dataplex uses [Identity and Access Management (IAM)][iam-overview] to control
+user and group access to Dataplex resources. Toolbox will use your
+[Application Default Credentials (ADC)][adc] to authorize and authenticate when
+interacting with [Dataplex][dataplex-docs].
+
+In addition to [setting the ADC for your server][set-adc], you need to ensure
+the IAM identity has been given the correct IAM permissions for the tasks you
+intend to perform. See [Dataplex Universal Catalog IAM permissions][iam-permissions]
+and [Dataplex Universal Catalog IAM roles][iam-roles] for more information on
+applying IAM permissions and roles to an identity.
+
+[iam-overview]: https://cloud.google.com/dataplex/docs/iam-and-access-control
+[adc]: https://cloud.google.com/docs/authentication#adc
+[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc
+[iam-permissions]: https://cloud.google.com/dataplex/docs/iam-permissions
+[iam-roles]: https://cloud.google.com/dataplex/docs/iam-roles
+[dataplex-docs]: https://cloud.google.com/dataplex
+
+## Example
+
+```yaml
+kind: tools
+name: dataplex-search-dq-scans
+type: dataplex-search-dq-scans
+source: my-dataplex-source
+description: Use this tool to search for data quality scans.
+```
+
+## Reference
+
+| **field**   | **type** | **required** | **description**                                    |
+|-------------|:--------:|:------------:|----------------------------------------------------|
+| type        |  string  |     true     | Must be "dataplex-search-dq-scans".                |
+| 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. |

internal/prebuiltconfigs/tools/dataplex.yaml

@@ -30,9 +30,14 @@ tools:
     kind: dataplex-search-aspect-types
     source: dataplex-source
     description: Use this tool to find aspect types relevant to the query.
+  search_dq_scans:
+    kind: dataplex-search-dq-scans
+    source: dataplex-source
+    description: Use this tool to search for data quality scans in Dataplex.
 
 toolsets:
   dataplex_tools:
     - search_entries
     - lookup_entry
     - search_aspect_types
+    - search_dq_scans

internal/sources/dataplex/dataplex.go

@@ -64,13 +64,14 @@ func (r Config) SourceConfigType() string {
 
 func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
 	// Initializes a Dataplex source
-	client, err := initDataplexConnection(ctx, tracer, r.Name, r.Project)
+	client, dataScanClient, err := initDataplexConnection(ctx, tracer, r.Name, r.Project)
 	if err != nil {
 		return nil, err
 	}
 	s := &Source{
-		Config: r,
-		Client: client,
+		Config:         r,
+		Client:         client,
+		DataScanClient: dataScanClient,
 	}
 
 	return s, nil
@@ -80,7 +81,8 @@ var _ sources.Source = &Source{}
 
 type Source struct {
 	Config
-	Client *dataplexapi.CatalogClient
+	Client         *dataplexapi.CatalogClient
+	DataScanClient *dataplexapi.DataScanClient
 }
 
 func (s *Source) SourceType() string {
@@ -100,30 +102,39 @@ func (s *Source) CatalogClient() *dataplexapi.CatalogClient {
 	return s.Client
 }
 
+func (s *Source) GetDataScanClient() *dataplexapi.DataScanClient {
+	return s.DataScanClient
+}
+
 func initDataplexConnection(
 	ctx context.Context,
 	tracer trace.Tracer,
 	name string,
 	project string,
-) (*dataplexapi.CatalogClient, error) {
+) (*dataplexapi.CatalogClient, *dataplexapi.DataScanClient, error) {
 	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
 	defer span.End()
 
 	cred, err := google.FindDefaultCredentials(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err)
+		return nil, nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err)
 	}
 
 	userAgent, err := util.UserAgentFromContext(ctx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	client, err := dataplexapi.NewCatalogClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred))
 	if err != nil {
-		return nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
+		return nil, nil, fmt.Errorf("failed to create Dataplex client for project %q: %w", project, err)
+	}
+
+	dataScanClient, err := dataplexapi.NewDataScanClient(ctx, option.WithUserAgent(userAgent), option.WithCredentials(cred))
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create Dataplex DataScan client for project %q: %w", project, err)
 	}
-	return client, nil
+	return client, dataScanClient, nil
 }
 
 func (s *Source) LookupEntry(ctx context.Context, name string, view int, aspectTypes []string, entry string) (*dataplexpb.Entry, error) {
@@ -240,3 +251,33 @@ func (s *Source) SearchEntries(ctx context.Context, query string, pageSize int,
 	}
 	return results, nil
 }
+
+func (s *Source) SearchDataQualityScans(ctx context.Context, query string, pageSize int, orderBy string) ([]*dataplexpb.DataScan, error) {
+	// ListDataScansRequest doesn't support generic specific "query" field like SearchEntries,
+	// but it supports "filter".  We will assume `query` is passed as filter.
+	// If the user wants to search by name/display_name they should provide a filter string.
+	// Example filter: "display_name = \"my-scan\""
+	req := &dataplexpb.ListDataScansRequest{
+		Parent:   fmt.Sprintf("projects/%s/locations/-", s.ProjectID()),
+		Filter:   query,
+		PageSize: int32(pageSize),
+		OrderBy:  orderBy,
+	}
+
+	it := s.GetDataScanClient().ListDataScans(ctx, req)
+	var results []*dataplexpb.DataScan
+	for {
+		scan, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			if st, ok := grpcstatus.FromError(err); ok {
+				return nil, fmt.Errorf("failed to list data scans: code=%s message=%s", st.Code(), st.Message())
+			}
+			return nil, fmt.Errorf("failed to list data scans: %w", err)
+		}
+		results = append(results, scan)
+	}
+	return results, nil
+}

internal/tools/dataplex/dataplexsearchdqscans/dataplexsearchdqscans.go

@@ -0,0 +1,165 @@
+// 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 dataplexsearchdqscans
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"cloud.google.com/go/dataplex/apiv1/dataplexpb"
+	"github.com/goccy/go-yaml"
+	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
+	"github.com/googleapis/genai-toolbox/internal/sources"
+	"github.com/googleapis/genai-toolbox/internal/tools"
+	"github.com/googleapis/genai-toolbox/internal/util/parameters"
+)
+
+const resourceType string = "dataplex-search-dq-scans"
+
+func init() {
+	if !tools.Register(resourceType, newConfig) {
+		panic(fmt.Sprintf("tool type %q already registered", resourceType))
+	}
+}
+
+func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
+	actual := Config{Name: name}
+	if err := decoder.DecodeContext(ctx, &actual); err != nil {
+		return nil, err
+	}
+	return actual, nil
+}
+
+type compatibleSource interface {
+	SearchDataQualityScans(context.Context, string, int, string) ([]*dataplexpb.DataScan, error)
+}
+
+type Config struct {
+	Name         string   `yaml:"name" validate:"required"`
+	Type         string   `yaml:"type" validate:"required"`
+	Source       string   `yaml:"source" validate:"required"`
+	Description  string   `yaml:"description"`
+	AuthRequired []string `yaml:"authRequired"`
+}
+
+// validate interface
+var _ tools.ToolConfig = Config{}
+
+func (cfg Config) ToolConfigType() string {
+	return resourceType
+}
+
+func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
+	filter := parameters.NewStringParameterWithDefault("filter", "", "Optional. Filter string to search/filter data quality scans. E.g. \"display_name = \\\"my-scan\\\"\"")
+	dataScanID := parameters.NewStringParameterWithDefault("data_scan_id", "", "Optional. The ID of the data scan to filter by.")
+	tableName := parameters.NewStringParameterWithDefault("table_name", "", "Optional. The name of the table to filter by.")
+	pageSize := parameters.NewIntParameterWithDefault("pageSize", 10, "Number of returned data quality scans in the page.")
+	orderBy := parameters.NewStringParameterWithDefault("orderBy", "", "Specifies the ordering of results.")
+	params := parameters.Parameters{filter, dataScanID, tableName, pageSize, orderBy}
+
+	mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, params, nil)
+
+	t := Tool{
+		Config:     cfg,
+		Parameters: params,
+		manifest: tools.Manifest{
+			Description:  cfg.Description,
+			Parameters:   params.Manifest(),
+			AuthRequired: cfg.AuthRequired,
+		},
+		mcpManifest: mcpManifest,
+	}
+	return t, nil
+}
+
+type Tool struct {
+	Config
+	Parameters  parameters.Parameters
+	manifest    tools.Manifest
+	mcpManifest tools.McpManifest
+}
+
+func (t Tool) ToConfig() tools.ToolConfig {
+	return t.Config
+}
+
+func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
+	source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
+	if err != nil {
+		return nil, err
+	}
+	paramsMap := params.AsMap()
+	filter, _ := paramsMap["filter"].(string)
+	dataScanID, _ := paramsMap["data_scan_id"].(string)
+	tableName, _ := paramsMap["table_name"].(string)
+	pageSize, _ := paramsMap["pageSize"].(int)
+	orderBy, _ := paramsMap["orderBy"].(string)
+
+	var filters []string
+	if filter != "" {
+		filters = append(filters, filter)
+	}
+	if dataScanID != "" {
+		// assuming data_scan_id usually maps to name or id filter if supported, or display_name
+		// referencing user request "parameters could be datascan name"
+		// The generic filter "resource.name" or "display_name" is often used.
+		// Let's assume display_name for user convenience or name if full resource name.
+		// If it's just ID, we might need wildcard?
+		// Actually, `id` might be part of the resource name.
+		// Let's use `display_name` as it's more likely what user means by "name".
+		// Or if they mean ID, it might be `resource.name : id`.
+		// Let's try `display_name = "ID"` first as safe bet or just append to filter.
+		filters = append(filters, fmt.Sprintf("display_name = %q", dataScanID))
+	}
+	if tableName != "" {
+		// "data.entity" is typically used for table in DataScan filters
+		filters = append(filters, fmt.Sprintf("data.entity = %q", tableName))
+	}
+
+	finalFilter := strings.Join(filters, " AND ")
+
+	return source.SearchDataQualityScans(ctx, finalFilter, pageSize, orderBy)
+}
+
+func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
+	return parameters.EmbedParams(ctx, t.Parameters, paramValues, embeddingModelsMap, nil)
+}
+
+func (t Tool) Manifest() tools.Manifest {
+	// Returns the tool manifest
+	return t.manifest
+}
+
+func (t Tool) McpManifest() tools.McpManifest {
+	// Returns the tool MCP manifest
+	return t.mcpManifest
+}
+
+func (t Tool) Authorized(verifiedAuthServices []string) bool {
+	return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
+}
+
+func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
+	return false, nil
+}
+
+func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
+	return "Authorization", nil
+}
+
+func (t Tool) GetParameters() parameters.Parameters {
+	return t.Parameters
+}