Files
genai-toolbox/internal/tools/serverlessspark/createbatch/config.go
Yuan Teoh 293c1d6889 feat!: update configuration file v2 (#2369)
This PR introduces a significant update to the Toolbox configuration
file format, which is one of the primary **breaking changes** required
for the implementation of the Advanced Control Plane.

# Summary of Changes
The configuration schema has been updated to enforce resource isolation
and facilitate atomic, incremental updates.
* Resource Isolation: Resource definitions are now separated into
individual blocks, using a distinct structure for each resource type
(Source, Tool, Toolset, etc.). This improves readability, management,
and auditing of configuration files.
* Field Name Modification: Internal field names have been modified to
align with declarative methodologies. Specifically, the configuration
now separates kind (general resource type, e.g., Source) from type
(specific implementation, e.g., Postgres).

# User Impact
Existing tools.yaml configuration files are now in an outdated format.
Users must eventually update their files to the new YAML format.

# Mitigation & Compatibility
Backward compatibility is maintained during this transition to ensure no
immediate user action is required for existing files.
* Immediate Backward Compatibility: The source code includes a
pre-processing layer that automatically detects outdated configuration
files (v1 format) and converts them to the new v2 format under the hood.
* [COMING SOON] Migration Support: The new toolbox migrate subcommand
will be introduced to allow users to automatically convert their old
configuration files to the latest format.

# Example
Example for config file v2:
```
kind: sources
name: my-pg-instance
type: cloud-sql-postgres
project: my-project
region: my-region
instance: my-instance
database: my_db
user: my_user
password: my_pass
---
kind: authServices
name: my-google-auth
type: google
clientId: testing-id
---
kind: tools
name: example_tool
type: postgres-sql
source: my-pg-instance
description: some description
statement: SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
  type: string
  description: some description
---
kind: tools
name: example_tool_2
type: postgres-sql
source: my-pg-instance
description: returning the number one
statement: SELECT 1;
---
kind: toolsets
name: example_toolset
tools:
- example_tool
```

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Averi Kitsch <akitsch@google.com>
2026-01-27 16:58:43 -08:00

97 lines
3.4 KiB
Go

// 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 createbatch
import (
"context"
"encoding/json"
"fmt"
dataprocpb "cloud.google.com/go/dataproc/v2/apiv1/dataprocpb"
"github.com/goccy/go-yaml"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
)
// unmarshalProto is a helper function to unmarshal a generic interface{} into a proto.Message.
func unmarshalProto(data any, m proto.Message) error {
jsonData, err := json.Marshal(data)
if err != nil {
return fmt.Errorf("failed to marshal to JSON: %w", err)
}
return protojson.Unmarshal(jsonData, m)
}
type compatibleSource interface {
CreateBatch(context.Context, *dataprocpb.Batch) (map[string]any, error)
}
// Config is a common config that can be used with any type of create batch tool. However, each tool
// will still need its own config type, embedding this Config, so it can provide a type-specific
// Initialize implementation.
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"`
RuntimeConfig *dataprocpb.RuntimeConfig `yaml:"runtimeConfig"`
EnvironmentConfig *dataprocpb.EnvironmentConfig `yaml:"environmentConfig"`
AuthRequired []string `yaml:"authRequired"`
}
func NewConfig(ctx context.Context, name string, decoder *yaml.Decoder) (Config, error) {
// Use a temporary struct to decode the YAML, so that we can handle the proto
// conversion for RuntimeConfig and EnvironmentConfig.
var ymlCfg struct {
Name string `yaml:"name"`
Type string `yaml:"type"`
Source string `yaml:"source"`
Description string `yaml:"description"`
RuntimeConfig any `yaml:"runtimeConfig"`
EnvironmentConfig any `yaml:"environmentConfig"`
AuthRequired []string `yaml:"authRequired"`
}
if err := decoder.DecodeContext(ctx, &ymlCfg); err != nil {
return Config{}, err
}
cfg := Config{
Name: name,
Type: ymlCfg.Type,
Source: ymlCfg.Source,
Description: ymlCfg.Description,
AuthRequired: ymlCfg.AuthRequired,
}
if ymlCfg.RuntimeConfig != nil {
rc := &dataprocpb.RuntimeConfig{}
if err := unmarshalProto(ymlCfg.RuntimeConfig, rc); err != nil {
return Config{}, fmt.Errorf("failed to unmarshal runtimeConfig: %w", err)
}
cfg.RuntimeConfig = rc
}
if ymlCfg.EnvironmentConfig != nil {
ec := &dataprocpb.EnvironmentConfig{}
if err := unmarshalProto(ymlCfg.EnvironmentConfig, ec); err != nil {
return Config{}, fmt.Errorf("failed to unmarshal environmentConfig: %w", err)
}
cfg.EnvironmentConfig = ec
}
return cfg, nil
}