From 69a3cafabec5a40e2776d71de3587c0d16c722a2 Mon Sep 17 00:00:00 2001 From: Pete Hampton Date: Mon, 22 Sep 2025 17:27:06 +0100 Subject: [PATCH] feat(tools/clickhouse-list-tables): add list-tables tool (#1446) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Follows up https://github.com/googleapis/genai-toolbox/pull/1274/ with a list tables tool for ClickHouse ## PR Checklist --- > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes # --------- Co-authored-by: Pete Hampton Co-authored-by: Averi Kitsch --- cmd/root.go | 1 + cmd/root_test.go | 2 +- .../clickhouse/clickhouse-list-tables.md | 60 +++++++ .../prebuiltconfigs/tools/clickhouse.yaml | 6 + .../clickhouselisttables.go | 168 ++++++++++++++++++ .../clickhouselisttables_test.go | 116 ++++++++++++ .../clickhouse/clickhouse_integration_test.go | 149 ++++++++++++++++ 7 files changed, 501 insertions(+), 1 deletion(-) create mode 100644 docs/en/resources/tools/clickhouse/clickhouse-list-tables.md create mode 100644 internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go create mode 100644 internal/tools/clickhouse/clickhouselisttables/clickhouselisttables_test.go diff --git a/cmd/root.go b/cmd/root.go index 21c7d26e52..e38f8432ae 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -66,6 +66,7 @@ import ( _ "github.com/googleapis/genai-toolbox/internal/tools/bigtable" _ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouseexecutesql" _ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselistdatabases" + _ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselisttables" _ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhousesql" _ "github.com/googleapis/genai-toolbox/internal/tools/cloudmonitoring" _ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreatedatabase" diff --git a/cmd/root_test.go b/cmd/root_test.go index 1e1fb7ec70..0ea5cf547a 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -1409,7 +1409,7 @@ func TestPrebuiltTools(t *testing.T) { wantToolset: server.ToolsetConfigs{ "clickhouse_database_tools": tools.ToolsetConfig{ Name: "clickhouse_database_tools", - ToolNames: []string{"execute_sql", "list_databases"}, + ToolNames: []string{"execute_sql", "list_databases", "list_tables"}, }, }, }, diff --git a/docs/en/resources/tools/clickhouse/clickhouse-list-tables.md b/docs/en/resources/tools/clickhouse/clickhouse-list-tables.md new file mode 100644 index 0000000000..c34189e549 --- /dev/null +++ b/docs/en/resources/tools/clickhouse/clickhouse-list-tables.md @@ -0,0 +1,60 @@ +--- +title: "clickhouse-list-tables" +type: docs +weight: 4 +description: > + A "clickhouse-list-tables" tool lists all tables in a specific ClickHouse database. +aliases: +- /resources/tools/clickhouse-list-tables +--- + +## About + +A `clickhouse-list-tables` tool lists all available tables in a specified +ClickHouse database. It's compatible with the [clickhouse](../../sources/clickhouse.md) source. + +This tool executes the `SHOW TABLES FROM ` command and returns a list +of all tables in the specified database that are accessible to the configured +user, making it useful for schema exploration and table discovery tasks. + +## Example + +```yaml +tools: + list_clickhouse_tables: + kind: clickhouse-list-tables + source: my-clickhouse-instance + description: List all tables in a specific ClickHouse database +``` + +## Parameters + +| **parameter** | **type** | **required** | **description** | +|---------------|:--------:|:------------:|---------------------------------------------| +| database | string | true | The database to list tables from. | + +## Return Value + +The tool returns an array of objects, where each object contains: +- `name`: The name of the table +- `database`: The database the table belongs to + +Example response: +```json +[ + {"name": "users", "database": "analytics"}, + {"name": "events", "database": "analytics"}, + {"name": "products", "database": "analytics"}, + {"name": "orders", "database": "analytics"} +] +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|--------------------|:------------------:|:------------:|-----------------------------------------------------------| +| kind | string | true | Must be "clickhouse-list-tables". | +| source | string | true | Name of the ClickHouse source to list tables from. | +| description | string | true | Description of the tool that is passed to the LLM. | +| authRequired | array of string | false | Authentication services required to use this tool. | +| parameters | array of Parameter | false | Parameters for the tool (see Parameters section above). | \ No newline at end of file diff --git a/internal/prebuiltconfigs/tools/clickhouse.yaml b/internal/prebuiltconfigs/tools/clickhouse.yaml index a99c22dab8..165537b5be 100644 --- a/internal/prebuiltconfigs/tools/clickhouse.yaml +++ b/internal/prebuiltconfigs/tools/clickhouse.yaml @@ -32,7 +32,13 @@ tools: source: clickhouse-source description: Use this tool to list all databases in ClickHouse. + list_tables: + kind: clickhouse-list-tables + source: clickhouse-source + description: Use this tool to list all tables in a specific ClickHouse database. + toolsets: clickhouse_database_tools: - execute_sql - list_databases + - list_tables diff --git a/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go new file mode 100644 index 0000000000..37a08e1878 --- /dev/null +++ b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables.go @@ -0,0 +1,168 @@ +// 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 clickhouse + +import ( + "context" + "database/sql" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + "github.com/googleapis/genai-toolbox/internal/tools" +) + +type compatibleSource interface { + ClickHousePool() *sql.DB +} + +var compatibleSources = []string{"clickhouse"} + +const listTablesKind string = "clickhouse-list-tables" +const databaseKey string = "database" + +func init() { + if !tools.Register(listTablesKind, newListTablesConfig) { + panic(fmt.Sprintf("tool kind %q already registered", listTablesKind)) + } +} + +func newListTablesConfig(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 Config struct { + Name string `yaml:"name" validate:"required"` + Kind string `yaml:"kind" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description" validate:"required"` + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` +} + +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return listTablesKind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + s, ok := rawS.(compatibleSource) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", listTablesKind, compatibleSources) + } + + databaseParameter := tools.NewStringParameter(databaseKey, "The database to list tables from.") + parameters := tools.Parameters{databaseParameter} + + allParameters, paramManifest, paramMcpManifest, _ := tools.ProcessParameters(nil, parameters) + + mcpManifest := tools.McpManifest{ + Name: cfg.Name, + Description: cfg.Description, + InputSchema: paramMcpManifest, + } + + t := Tool{ + Name: cfg.Name, + Kind: listTablesKind, + Parameters: parameters, + AllParams: allParameters, + AuthRequired: cfg.AuthRequired, + Pool: s.ClickHousePool(), + manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired}, + mcpManifest: mcpManifest, + } + return t, nil +} + +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + AllParams tools.Parameters `yaml:"allParams"` + + Pool *sql.DB + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, token tools.AccessToken) (any, error) { + mapParams := params.AsMap() + database, ok := mapParams[databaseKey].(string) + if !ok { + return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", databaseKey) + } + + // Query to list all tables in the specified database + query := fmt.Sprintf("SHOW TABLES FROM %s", database) + + results, err := t.Pool.QueryContext(ctx, query) + if err != nil { + return nil, fmt.Errorf("unable to execute query: %w", err) + } + defer results.Close() + + var tables []map[string]any + for results.Next() { + var tableName string + err := results.Scan(&tableName) + if err != nil { + return nil, fmt.Errorf("unable to parse row: %w", err) + } + tables = append(tables, map[string]any{ + "name": tableName, + "database": database, + }) + } + + if err := results.Err(); err != nil { + return nil, fmt.Errorf("errors encountered by results.Scan: %w", err) + } + + return tables, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParseParams(t.AllParams, data, claims) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) McpManifest() tools.McpManifest { + return t.mcpManifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization() bool { + return false +} diff --git a/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables_test.go b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables_test.go new file mode 100644 index 0000000000..d4dcaa8700 --- /dev/null +++ b/internal/tools/clickhouse/clickhouselisttables/clickhouselisttables_test.go @@ -0,0 +1,116 @@ +// 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 clickhouse + +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/sources" + "github.com/googleapis/genai-toolbox/internal/testutils" + "github.com/googleapis/genai-toolbox/internal/tools" +) + +func TestListTablesConfigToolConfigKind(t *testing.T) { + cfg := Config{} + if cfg.ToolConfigKind() != listTablesKind { + t.Errorf("expected %q, got %q", listTablesKind, cfg.ToolConfigKind()) + } +} + +func TestListTablesConfigInitializeMissingSource(t *testing.T) { + cfg := Config{ + Name: "test-list-tables", + Kind: listTablesKind, + Source: "missing-source", + Description: "Test list tables tool", + } + + srcs := map[string]sources.Source{} + _, err := cfg.Initialize(srcs) + if err == nil { + t.Error("expected error for missing source") + } +} + +func TestParseFromYamlClickHouseListTables(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: clickhouse-list-tables + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": Config{ + Name: "example_tool", + Kind: "clickhouse-list-tables", + Source: "my-instance", + Description: "some description", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + 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) + } + }) + } +} + +func TestListTablesToolParseParams(t *testing.T) { + databaseParam := tools.NewStringParameter("database", "The database to list tables from.") + tool := Tool{ + Parameters: tools.Parameters{databaseParam}, + AllParams: tools.Parameters{databaseParam}, + } + + params, err := tool.ParseParams(map[string]any{"database": "test_db"}, map[string]map[string]any{}) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + + if len(params) != 1 { + t.Errorf("expected 1 parameter, got %d", len(params)) + } + + mapParams := params.AsMap() + if mapParams["database"] != "test_db" { + t.Errorf("expected database parameter to be 'test_db', got %v", mapParams["database"]) + } +} diff --git a/tests/clickhouse/clickhouse_integration_test.go b/tests/clickhouse/clickhouse_integration_test.go index a1fb0b6f37..b39c812cfd 100644 --- a/tests/clickhouse/clickhouse_integration_test.go +++ b/tests/clickhouse/clickhouse_integration_test.go @@ -32,6 +32,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/tools" clickhouseexecutesql "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouseexecutesql" clickhouselistdatabases "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselistdatabases" + clickhouselisttables "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselisttables" clickhousesql "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhousesql" "github.com/googleapis/genai-toolbox/tests" "go.opentelemetry.io/otel/trace/noop" @@ -1113,3 +1114,151 @@ func TestClickHouseListDatabasesTool(t *testing.T) { t.Logf("✅ clickhouse-list-databases tool tests completed successfully") } + +func TestClickHouseListTablesTool(t *testing.T) { + _ = getClickHouseVars(t) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + pool, err := initClickHouseConnectionPool(ClickHouseHost, ClickHousePort, ClickHouseUser, ClickHousePass, ClickHouseDatabase, ClickHouseProtocol) + if err != nil { + t.Fatalf("unable to create ClickHouse connection pool: %s", err) + } + defer pool.Close() + + // Create a test database with tables + testDBName := "test_list_tables_db_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:8] + _, err = pool.ExecContext(ctx, fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", testDBName)) + if err != nil { + t.Fatalf("Failed to create test database: %v", err) + } + defer func() { + _, _ = pool.ExecContext(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", testDBName)) + }() + + // Create test tables in the test database + testTable1 := "test_table_1" + testTable2 := "test_table_2" + _, err = pool.ExecContext(ctx, fmt.Sprintf("CREATE TABLE %s.%s (id UInt32, name String) ENGINE = Memory", testDBName, testTable1)) + if err != nil { + t.Fatalf("Failed to create test table 1: %v", err) + } + _, err = pool.ExecContext(ctx, fmt.Sprintf("CREATE TABLE %s.%s (id UInt32, value Float64) ENGINE = Memory", testDBName, testTable2)) + if err != nil { + t.Fatalf("Failed to create test table 2: %v", err) + } + + t.Run("ListTables", func(t *testing.T) { + toolConfig := clickhouselisttables.Config{ + Name: "test-list-tables", + Kind: "clickhouse-list-tables", + Source: "test-clickhouse", + Description: "Test listing tables", + } + + source := createMockSource(t, pool) + sourcesMap := map[string]sources.Source{ + "test-clickhouse": source, + } + + tool, err := toolConfig.Initialize(sourcesMap) + if err != nil { + t.Fatalf("Failed to initialize tool: %v", err) + } + + params := tools.ParamValues{ + {Name: "database", Value: testDBName}, + } + + result, err := tool.Invoke(ctx, params, "") + if err != nil { + t.Fatalf("Failed to list tables: %v", err) + } + + tables, ok := result.([]map[string]any) + if !ok { + t.Fatalf("Expected result to be []map[string]any, got %T", result) + } + + // Should contain exactly 2 tables that we created + if len(tables) != 2 { + t.Errorf("Expected 2 tables, got %d", len(tables)) + } + + foundTable1 := false + foundTable2 := false + for _, table := range tables { + if name, ok := table["name"].(string); ok { + if name == testTable1 { + foundTable1 = true + } + if name == testTable2 { + foundTable2 = true + } + // Verify database field is set correctly + if db, ok := table["database"].(string); ok { + if db != testDBName { + t.Errorf("Expected database to be %s, got %s", testDBName, db) + } + } + } + } + + if !foundTable1 { + t.Errorf("Test table %s not found in list", testTable1) + } + if !foundTable2 { + t.Errorf("Test table %s not found in list", testTable2) + } + + t.Logf("Successfully listed %d tables from database %s", len(tables), testDBName) + }) + + t.Run("ListTablesWithMissingDatabase", func(t *testing.T) { + toolConfig := clickhouselisttables.Config{ + Name: "test-list-tables-missing-db", + Kind: "clickhouse-list-tables", + Source: "test-clickhouse", + Description: "Test listing tables without database parameter", + } + + source := createMockSource(t, pool) + sourcesMap := map[string]sources.Source{ + "test-clickhouse": source, + } + + tool, err := toolConfig.Initialize(sourcesMap) + if err != nil { + t.Fatalf("Failed to initialize tool: %v", err) + } + + params := tools.ParamValues{} + + _, err = tool.Invoke(ctx, params, "") + if err == nil { + t.Error("Expected error for missing database parameter, got nil") + } else { + t.Logf("Got expected error for missing database: %v", err) + } + }) + + t.Run("ListTablesWithInvalidSource", func(t *testing.T) { + toolConfig := clickhouselisttables.Config{ + Name: "test-invalid-source", + Kind: "clickhouse-list-tables", + Source: "non-existent-source", + Description: "Test with invalid source", + } + + sourcesMap := map[string]sources.Source{} + + _, err := toolConfig.Initialize(sourcesMap) + if err == nil { + t.Error("Expected error for non-existent source, got nil") + } else { + t.Logf("Got expected error for invalid source: %v", err) + } + }) + + t.Logf("✅ clickhouse-list-tables tool tests completed successfully") +}