diff --git a/cmd/root.go b/cmd/root.go index 4231a69dd3..5ddc1b88e4 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -102,6 +102,11 @@ import ( _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerconversationalanalytics" _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerdeleteprojectfile" _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerdevmode" + _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiondatabases" + _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnections" + _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectionschemas" + _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontables" + _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontablecolumns" _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdashboards" _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdimensions" _ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetexplores" diff --git a/cmd/root_test.go b/cmd/root_test.go index 472bf4309c..1c3d9c1be7 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -1507,7 +1507,7 @@ func TestPrebuiltTools(t *testing.T) { wantToolset: server.ToolsetConfigs{ "looker_tools": tools.ToolsetConfig{ Name: "looker_tools", - ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "make_dashboard", "add_dashboard_element", "health_pulse", "health_analyze", "health_vacuum", "dev_mode", "get_projects", "get_project_files", "get_project_file", "create_project_file", "update_project_file", "delete_project_file"}, + ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "make_dashboard", "add_dashboard_element", "health_pulse", "health_analyze", "health_vacuum", "dev_mode", "get_projects", "get_project_files", "get_project_file", "create_project_file", "update_project_file", "delete_project_file", "get_connections", "get_connection_schemas", "get_connection_databases", "get_connection_tables", "get_connection_table_columns"}, }, }, }, diff --git a/docs/en/resources/tools/looker/looker-get-connection-databases.md b/docs/en/resources/tools/looker/looker-get-connection-databases.md new file mode 100644 index 0000000000..bfeab7a0f4 --- /dev/null +++ b/docs/en/resources/tools/looker/looker-get-connection-databases.md @@ -0,0 +1,41 @@ +--- +title: "looker-get-connection-databases" +type: docs +weight: 1 +description: > + A "looker-get-connection-databases" tool returns all the databases in a connection. +aliases: +- /resources/tools/looker-get-connection-databases +--- + +## About + +A `looker-get-connection-databases` tool returns all the databases in a connection. + +It's compatible with the following sources: + +- [looker](../../sources/looker.md) + +`looker-get-connection-databases` accepts a `conn` parameter. + +## Example + +```yaml +tools: + get_connection_databases: + kind: looker-get-connection-databases + source: looker-source + description: | + get_connection_databases Tool + + This tool will list the databases available from a connection if the connection + supports multiple databases. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| kind | string | true | Must be "looker-get-connection-databases". | +| source | string | true | Name of the source Looker instance. | +| description | string | true | Description of the tool that is passed to the LLM. | \ No newline at end of file diff --git a/docs/en/resources/tools/looker/looker-get-connection-schemas.md b/docs/en/resources/tools/looker/looker-get-connection-schemas.md new file mode 100644 index 0000000000..b97e9c9fa6 --- /dev/null +++ b/docs/en/resources/tools/looker/looker-get-connection-schemas.md @@ -0,0 +1,41 @@ +--- +title: "looker-get-connection-schemas" +type: docs +weight: 1 +description: > + A "looker-get-connection-schemas" tool returns all the schemas in a connection. +aliases: +- /resources/tools/looker-get-connection-schemas +--- + +## About + +A `looker-get-connection-schemas` tool returns all the schemas in a connection. + +It's compatible with the following sources: + +- [looker](../../sources/looker.md) + +`looker-get-connection-schemas` accepts a `conn` parameter and an optional `db` parameter. + +## Example + +```yaml +tools: + get_connection_schemas: + kind: looker-get-connection-schemas + source: looker-source + description: | + get_connection_schemas Tool + + This tool will list the schemas available from a connection, filtered by + an optional database name. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| kind | string | true | Must be "looker-get-connection-schemas". | +| source | string | true | Name of the source Looker instance. | +| description | string | true | Description of the tool that is passed to the LLM. | \ No newline at end of file diff --git a/docs/en/resources/tools/looker/looker-get-connection-table-columns.md b/docs/en/resources/tools/looker/looker-get-connection-table-columns.md new file mode 100644 index 0000000000..3dfba7ffea --- /dev/null +++ b/docs/en/resources/tools/looker/looker-get-connection-table-columns.md @@ -0,0 +1,43 @@ +--- +title: "looker-get-connection-table-columns" +type: docs +weight: 1 +description: > + A "looker-get-connection-table-columns" tool returns all the columns for each table specified. +aliases: +- /resources/tools/looker-get-connection-table-columns +--- + +## About + +A `looker-get-connection-table-columns` tool returns all the columnes for each table specified. + + +It's compatible with the following sources: + +- [looker](../../sources/looker.md) + +`looker-get-connection-table-columns` accepts a `conn` parameter, a `schema` parameter, a `tables` parameter with a comma separated list of tables, and an optional `db` parameter. + +## Example + +```yaml +tools: + get_connection_table_columns: + kind: looker-get-connection-table-columns + source: looker-source + description: | + get_connection_table_columns Tool + + This tool will list the columns available from a connection, for all the tables + given in a comma separated list of table names, filtered by the + schema name and optional database name. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| kind | string | true | Must be "looker-get-connection-table-columns". | +| source | string | true | Name of the source Looker instance. | +| description | string | true | Description of the tool that is passed to the LLM. | \ No newline at end of file diff --git a/docs/en/resources/tools/looker/looker-get-connection-tables.md b/docs/en/resources/tools/looker/looker-get-connection-tables.md new file mode 100644 index 0000000000..7b66ef2a76 --- /dev/null +++ b/docs/en/resources/tools/looker/looker-get-connection-tables.md @@ -0,0 +1,41 @@ +--- +title: "looker-get-connection-tables" +type: docs +weight: 1 +description: > + A "looker-get-connection-tables" tool returns all the tables in a connection. +aliases: +- /resources/tools/looker-get-connection-tables +--- + +## About + +A `looker-get-connection-tables` tool returns all the tables in a connection. + +It's compatible with the following sources: + +- [looker](../../sources/looker.md) + +`looker-get-connection-tables` accepts a `conn` parameter, a `schema` parameter, and an optional `db` parameter. + +## Example + +```yaml +tools: + get_connection_tables: + kind: looker-get-connection-tables + source: looker-source + description: | + get_connection_tables Tool + + This tool will list the tables available from a connection, filtered by the + schema name and optional database name. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| kind | string | true | Must be "looker-get-connection-tables". | +| source | string | true | Name of the source Looker instance. | +| description | string | true | Description of the tool that is passed to the LLM. | \ No newline at end of file diff --git a/docs/en/resources/tools/looker/looker-get-connections.md b/docs/en/resources/tools/looker/looker-get-connections.md new file mode 100644 index 0000000000..79cbac30e3 --- /dev/null +++ b/docs/en/resources/tools/looker/looker-get-connections.md @@ -0,0 +1,42 @@ +--- +title: "looker-get-connections" +type: docs +weight: 1 +description: > + A "looker-get-connections" tool returns all the connections in the source. +aliases: +- /resources/tools/looker-get-connections +--- + +## About + +A `looker-get-connections` tool returns all the connections in the source. + +It's compatible with the following sources: + +- [looker](../../sources/looker.md) + +`looker-get-connections` accepts no parameters. + +## Example + +```yaml +tools: + get_connections: + kind: looker-get-connections + source: looker-source + description: | + get_connections Tool + + This tool will list all the connections available in the Looker system, as + well as the dialect name, the default schema, the database if applicable, + and whether the connection supports multiple databases. +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|----------------------------------------------------| +| kind | string | true | Must be "looker-get-connections". | +| source | string | true | Name of the source Looker instance. | +| description | string | true | Description of the tool that is passed to the LLM. | \ No newline at end of file diff --git a/internal/prebuiltconfigs/tools/looker.yaml b/internal/prebuiltconfigs/tools/looker.yaml index 30ebe38bd8..3e049a58e8 100644 --- a/internal/prebuiltconfigs/tools/looker.yaml +++ b/internal/prebuiltconfigs/tools/looker.yaml @@ -823,6 +823,54 @@ tools: This tool must be called after the dev_mode tool has changed the session to dev mode. + get_connections: + kind: looker-get-connections + source: looker-source + description: | + get_connections Tool + + This tool will list all the connections available in the Looker system, as + well as the dialect name, the default schema, the database if applicable, + and whether the connection supports multiple databases. + + get_connection_schemas: + kind: looker-get-connection-schemas + source: looker-source + description: | + get_connection_schemas Tool + + This tool will list the schemas available from a connection, filtered by + an optional database name. + + get_connection_databases: + kind: looker-get-connection-databases + source: looker-source + description: | + get_connection_databases Tool + + This tool will list the databases available from a connection if the connection + supports multiple databases. + + get_connection_tables: + kind: looker-get-connection-tables + source: looker-source + description: | + get_connection_tables Tool + + This tool will list the tables available from a connection, filtered by the + schema name and optional database name. + + get_connection_table_columns: + kind: looker-get-connection-table-columns + source: looker-source + description: | + get_connection_table_columns Tool + + This tool will list the columns available from a connection, for all the tables + given in a comma separated list of table names, filtered by the + schema name and optional database name. + + toolsets: looker_tools: - get_models @@ -850,3 +898,8 @@ toolsets: - create_project_file - update_project_file - delete_project_file + - get_connections + - get_connection_schemas + - get_connection_databases + - get_connection_tables + - get_connection_table_columns diff --git a/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases.go b/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases.go new file mode 100644 index 0000000000..ac758fd65c --- /dev/null +++ b/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases.go @@ -0,0 +1,150 @@ +// 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 lookergetconnectiondatabases + +import ( + "context" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon" + + "github.com/looker-open-source/sdk-codegen/go/rtl" + v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4" +) + +const kind string = "looker-get-connection-databases" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type 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"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(*lookersrc.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind) + } + + connParameter := tools.NewStringParameter("conn", "The connection containing the databases.") + parameters := tools.Parameters{connParameter} + + mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) + + // finish tool setup + return Tool{ + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientOAuth, + Client: s.Client, + ApiSettings: s.ApiSettings, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: parameters.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + mcpManifest: mcpManifest, + }, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + UseClientOAuth bool + Client *v4.LookerSDK + ApiSettings *rtl.ApiSettings + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { + mapParams := params.AsMap() + conn, ok := mapParams["conn"].(string) + if !ok { + return nil, fmt.Errorf("'conn' must be a string, got %T", mapParams["conn"]) + } + + sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken) + if err != nil { + return nil, fmt.Errorf("error getting sdk: %w", err) + } + resp, err := sdk.ConnectionDatabases(conn, t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connection_databases request: %s", err) + } + //logger.DebugContext(ctx, "Got response of %v\n", resp) + + return resp, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParseParams(t.Parameters, 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 t.UseClientOAuth +} \ No newline at end of file diff --git a/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases_test.go b/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases_test.go new file mode 100644 index 0000000000..284d4102bb --- /dev/null +++ b/internal/tools/looker/lookergetconnectiondatabases/lookergetconnectiondatabases_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 lookergetconnectiondatabases_test + +import ( + "strings" + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + lkr "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiondatabases" +) + +func TestParseFromYamlLookerGetConnectionDatabases(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: looker-get-connection-databases + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": lkr.Config{ + Name: "example_tool", + Kind: "looker-get-connection-databases", + 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"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } + +} + +func TestFailParseFromYamlLookerGetConnectionDatabases(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "Invalid method", + in: ` + tools: + example_tool: + kind: looker-get-connection-databases + source: my-instance + method: GOT + description: some description + `, + err: "unable to parse tool \"example_tool\" as kind \"looker-get-connection-databases\": [4:1] unknown field \"method\"\n 1 | authRequired: []\n 2 | description: some description\n 3 | kind: looker-get-connection-databases\n> 4 | method: GOT\n ^\n 5 | source: my-instance", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if !strings.Contains(errStr, tc.err) { + t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err) + } + }) + } + +} diff --git a/internal/tools/looker/lookergetconnections/lookergetconnections.go b/internal/tools/looker/lookergetconnections/lookergetconnections.go new file mode 100644 index 0000000000..171e713658 --- /dev/null +++ b/internal/tools/looker/lookergetconnections/lookergetconnections.go @@ -0,0 +1,169 @@ +// 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 lookergetconnections + +import ( + "context" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon" + "github.com/googleapis/genai-toolbox/internal/util" + + "github.com/looker-open-source/sdk-codegen/go/rtl" + v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4" +) + +const kind string = "looker-get-connections" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type 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"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(*lookersrc.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind) + } + + parameters := tools.Parameters{} + + mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) + + // finish tool setup + return Tool{ + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientOAuth, + Client: s.Client, + ApiSettings: s.ApiSettings, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: parameters.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + mcpManifest: mcpManifest, + }, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + UseClientOAuth bool + Client *v4.LookerSDK + ApiSettings *rtl.ApiSettings + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { + logger, err := util.LoggerFromContext(ctx) + if err != nil { + return nil, fmt.Errorf("unable to get logger from ctx: %s", err) + } + + sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken) + if err != nil { + return nil, fmt.Errorf("error getting sdk: %w", err) + } + resp, err := sdk.AllConnections("name, dialect(name), database, schema",t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connections request: %s", err) + } + + var data []any + for _, v := range resp { + vMap := make(map[string]any) + vMap["name"] = *v.Name + vMap["dialect_name"] = *v.Dialect.Name + if v.Database != nil { + vMap["database"] = *v.Database + } + if v.Schema != nil { + vMap["schema"] = *v.Schema + } + conn, err := sdk.ConnectionFeatures(*v.Name, "multiple_databases", t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connection_features request: %s", err) + } + vMap["supports_multiple_databases"] = *conn.MultipleDatabases + logger.DebugContext(ctx, "Converted to %v\n", vMap) + data = append(data, vMap) + } + logger.DebugContext(ctx, "data = ", data) + + return data, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParamValues{}, nil +} + +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 t.UseClientOAuth +} diff --git a/internal/tools/looker/lookergetconnections/lookergetconnections_test.go b/internal/tools/looker/lookergetconnections/lookergetconnections_test.go new file mode 100644 index 0000000000..b921a57e94 --- /dev/null +++ b/internal/tools/looker/lookergetconnections/lookergetconnections_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 lookergetconnections_test + +import ( + "strings" + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + lkr "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnections" +) + +func TestParseFromYamlLookerGetConnections(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: looker-get-connections + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": lkr.Config{ + Name: "example_tool", + Kind: "looker-get-connections", + 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"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } + +} + +func TestFailParseFromYamlLookerGetConnections(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "Invalid method", + in: ` + tools: + example_tool: + kind: looker-get-connections + source: my-instance + method: GOT + description: some description + `, + err: "unable to parse tool \"example_tool\" as kind \"looker-get-connections\": [4:1] unknown field \"method\"\n 1 | authRequired: []\n 2 | description: some description\n 3 | kind: looker-get-connections\n> 4 | method: GOT\n ^\n 5 | source: my-instance", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if !strings.Contains(errStr, tc.err) { + t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err) + } + }) + } + +} diff --git a/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas.go b/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas.go new file mode 100644 index 0000000000..3b01178373 --- /dev/null +++ b/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas.go @@ -0,0 +1,156 @@ +// 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 lookergetconnectionschemas + +import ( + "context" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon" + + "github.com/looker-open-source/sdk-codegen/go/rtl" + v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4" +) + +const kind string = "looker-get-connection-schemas" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type 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"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(*lookersrc.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind) + } + + connParameter := tools.NewStringParameter("conn", "The connection containing the schemas.") + dbParameter := tools.NewStringParameterWithRequired("db", "The optional database to search", false) + parameters := tools.Parameters{connParameter, dbParameter} + + mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) + + // finish tool setup + return Tool{ + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientOAuth, + Client: s.Client, + ApiSettings: s.ApiSettings, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: parameters.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + mcpManifest: mcpManifest, + }, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + UseClientOAuth bool + Client *v4.LookerSDK + ApiSettings *rtl.ApiSettings + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { + mapParams := params.AsMap() + conn, ok := mapParams["conn"].(string) + if !ok { + return nil, fmt.Errorf("'conn' must be a string, got %T", mapParams["conn"]) + } + db, _ := mapParams["db"].(string) + + sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken) + if err != nil { + return nil, fmt.Errorf("error getting sdk: %w", err) + } + req := v4.RequestConnectionSchemas{ + ConnectionName: conn, + } + if db != "" { + req.Database = &db + } + resp, err := sdk.ConnectionSchemas(req, t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connection_schemas request: %s", err) + } + return resp, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParseParams(t.Parameters, 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 t.UseClientOAuth +} \ No newline at end of file diff --git a/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas_test.go b/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas_test.go new file mode 100644 index 0000000000..1c2951102f --- /dev/null +++ b/internal/tools/looker/lookergetconnectionschemas/lookergetconnectionschemas_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 lookergetconnectionschemas_test + +import ( + "strings" + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + lkr "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectionschemas" +) + +func TestParseFromYamlLookerGetConnectionSchemas(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: looker-get-connection-schemas + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": lkr.Config{ + Name: "example_tool", + Kind: "looker-get-connection-schemas", + 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"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } + +} + +func TestFailParseFromYamlLookerGetConnectionSchemas(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "Invalid method", + in: ` + tools: + example_tool: + kind: looker-get-connection-schemas + source: my-instance + method: GOT + description: some description + `, + err: "unable to parse tool \"example_tool\" as kind \"looker-get-connection-schemas\": [4:1] unknown field \"method\"\n 1 | authRequired: []\n 2 | description: some description\n 3 | kind: looker-get-connection-schemas\n> 4 | method: GOT\n ^\n 5 | source: my-instance", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if !strings.Contains(errStr, tc.err) { + t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err) + } + }) + } + +} diff --git a/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns.go b/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns.go new file mode 100644 index 0000000000..89e90ac3a0 --- /dev/null +++ b/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns.go @@ -0,0 +1,193 @@ +// 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 lookergetconnectiontablecolumns + +import ( + "context" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon" + "github.com/googleapis/genai-toolbox/internal/util" + + "github.com/looker-open-source/sdk-codegen/go/rtl" + v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4" +) + +const kind string = "looker-get-connection-table-columns" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type 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"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(*lookersrc.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind) + } + + connParameter := tools.NewStringParameter("conn", "The connection containing the tables.") + dbParameter := tools.NewStringParameterWithRequired("db", "The optional database to search", false) + schemaParameter := tools.NewStringParameter("schema", "The schema containing the tables.") + tablesParameter := tools.NewStringParameter("tables", "A comma separated list of tables containing the columns.") + parameters := tools.Parameters{connParameter, dbParameter, schemaParameter, tablesParameter} + + mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) + + // finish tool setup + return Tool{ + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientOAuth, + Client: s.Client, + ApiSettings: s.ApiSettings, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: parameters.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + mcpManifest: mcpManifest, + }, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + UseClientOAuth bool + Client *v4.LookerSDK + ApiSettings *rtl.ApiSettings + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { + logger, err := util.LoggerFromContext(ctx) + if err != nil { + return nil, fmt.Errorf("unable to get logger from ctx: %s", err) + } + mapParams := params.AsMap() + conn, ok := mapParams["conn"].(string) + if !ok { + return nil, fmt.Errorf("'conn' must be a string, got %T", mapParams["conn"]) + } + db, _ := mapParams["db"].(string) + schema, ok := mapParams["schema"].(string) + if !ok { + return nil, fmt.Errorf("'schema' must be a string, got %T", mapParams["schema"]) + } + tables, ok := mapParams["tables"].(string) + if !ok { + return nil, fmt.Errorf("'tables' must be a string, got %T", mapParams["tables"]) + } + + sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken) + if err != nil { + return nil, fmt.Errorf("error getting sdk: %w", err) + } + req := v4.RequestConnectionColumns{ + ConnectionName: conn, + SchemaName: &schema, + TableNames: &tables, + } + if db != "" { + req.Database = &db + } + resp, err := sdk.ConnectionColumns(req, t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connection_table_columns request: %s", err) + } + var data []any + for _, t := range resp { + vMap := make(map[string]any) + vMap["table_name"] = *t.Name + vMap["sql_escaped_table_name"] = *t.SqlEscapedName + vMap["schema_name"] = *t.SchemaName + var columnData []any + for _, c := range *t.Columns { + vMap2 := make(map[string]any) + vMap2["column_name"] = *c.Name + vMap2["sql_escaped_column_name"] = *c.SqlEscapedName + vMap2["data_type_database"] = *c.DataTypeDatabase + vMap2["data_type_looker"] = *c.DataTypeLooker + columnData = append(columnData, vMap2) + } + vMap["columns"] = columnData + data = append(data, vMap) + } + logger.DebugContext(ctx, "data = ", data) + + return data, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParseParams(t.Parameters, 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 t.UseClientOAuth +} diff --git a/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns_test.go b/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns_test.go new file mode 100644 index 0000000000..3242a31ffb --- /dev/null +++ b/internal/tools/looker/lookergetconnectiontablecolumns/lookergetconnectiontablecolumns_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 lookergetconnectiontablecolumns_test + +import ( + "strings" + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + lkr "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontablecolumns" +) + +func TestParseFromYamlLookerGetConnectionTableColumns(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: looker-get-connection-table-columns + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": lkr.Config{ + Name: "example_tool", + Kind: "looker-get-connection-table-columns", + 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"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } + +} + +func TestFailParseFromYamlLookerGetConnectionTableColumns(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "Invalid method", + in: ` + tools: + example_tool: + kind: looker-get-connection-table-columns + source: my-instance + method: GOT + description: some description + `, + err: "unable to parse tool \"example_tool\" as kind \"looker-get-connection-table-columns\": [4:1] unknown field \"method\"\n 1 | authRequired: []\n 2 | description: some description\n 3 | kind: looker-get-connection-table-columns\n> 4 | method: GOT\n ^\n 5 | source: my-instance", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if !strings.Contains(errStr, tc.err) { + t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err) + } + }) + } + +} diff --git a/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables.go b/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables.go new file mode 100644 index 0000000000..0ced34f270 --- /dev/null +++ b/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables.go @@ -0,0 +1,184 @@ +// 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 lookergetconnectiontables + +import ( + "context" + "fmt" + + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon" + "github.com/googleapis/genai-toolbox/internal/util" + + "github.com/looker-open-source/sdk-codegen/go/rtl" + v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4" +) + +const kind string = "looker-get-connection-tables" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type 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"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(*lookersrc.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind) + } + + connParameter := tools.NewStringParameter("conn", "The connection containing the tables.") + dbParameter := tools.NewStringParameterWithRequired("db", "The optional database to search", false) + schemaParameter := tools.NewStringParameter("schema", "The schema containing the tables.") + parameters := tools.Parameters{connParameter, dbParameter, schemaParameter} + + mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) + + // finish tool setup + return Tool{ + Name: cfg.Name, + Kind: kind, + Parameters: parameters, + AuthRequired: cfg.AuthRequired, + UseClientOAuth: s.UseClientOAuth, + Client: s.Client, + ApiSettings: s.ApiSettings, + manifest: tools.Manifest{ + Description: cfg.Description, + Parameters: parameters.Manifest(), + AuthRequired: cfg.AuthRequired, + }, + mcpManifest: mcpManifest, + }, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + UseClientOAuth bool + Client *v4.LookerSDK + ApiSettings *rtl.ApiSettings + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) { + logger, err := util.LoggerFromContext(ctx) + if err != nil { + return nil, fmt.Errorf("unable to get logger from ctx: %s", err) + } + mapParams := params.AsMap() + conn, ok := mapParams["conn"].(string) + if !ok { + return nil, fmt.Errorf("'conn' must be a string, got %T", mapParams["conn"]) + } + db, _ := mapParams["db"].(string) + schema, ok := mapParams["schema"].(string) + if !ok { + return nil, fmt.Errorf("'schema' must be a string, got %T", mapParams["schema"]) + } + + sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken) + if err != nil { + return nil, fmt.Errorf("error getting sdk: %w", err) + } + req := v4.RequestConnectionTables{ + ConnectionName: conn, + SchemaName: &schema, + } + if db != "" { + req.Database = &db + } + resp, err := sdk.ConnectionTables(req, t.ApiSettings) + if err != nil { + return nil, fmt.Errorf("error making get_connection_tables request: %s", err) + } + var data []any + for _, s := range resp { + vMap := make(map[string]any) + vMap["schema_name"] = *s.Name + vMap["is_default"] = *s.IsDefault + var tableData []any + for _, t := range *s.Tables { + vMap2 := make(map[string]any) + vMap2["table_name"] = *t.Name + vMap2["sql_escaped_name"] = *t.SqlEscapedName + tableData = append(tableData, vMap2) + } + vMap["tables"] = tableData + data = append(data, vMap) + } + logger.DebugContext(ctx, "data = ", data) + + return data, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) { + return tools.ParseParams(t.Parameters, 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 t.UseClientOAuth +} diff --git a/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables_test.go b/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables_test.go new file mode 100644 index 0000000000..1b6ec50aa0 --- /dev/null +++ b/internal/tools/looker/lookergetconnectiontables/lookergetconnectiontables_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 lookergetconnectiontables_test + +import ( + "strings" + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + lkr "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontables" +) + +func TestParseFromYamlLookerGetConnectionTables(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: looker-get-connection-tables + source: my-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": lkr.Config{ + Name: "example_tool", + Kind: "looker-get-connection-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"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } + +} + +func TestFailParseFromYamlLookerGetConnectionTables(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "Invalid method", + in: ` + tools: + example_tool: + kind: looker-get-connection-tables + source: my-instance + method: GOT + description: some description + `, + err: "unable to parse tool \"example_tool\" as kind \"looker-get-connection-tables\": [4:1] unknown field \"method\"\n 1 | authRequired: []\n 2 | description: some description\n 3 | kind: looker-get-connection-tables\n> 4 | method: GOT\n ^\n 5 | source: my-instance", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if !strings.Contains(errStr, tc.err) { + t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err) + } + }) + } + +} diff --git a/internal/tools/looker/lookergetmodels/lookergetmodels.go b/internal/tools/looker/lookergetmodels/lookergetmodels.go index 510becba48..bd340276e8 100644 --- a/internal/tools/looker/lookergetmodels/lookergetmodels.go +++ b/internal/tools/looker/lookergetmodels/lookergetmodels.go @@ -142,6 +142,7 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken vMap["label"] = *v.Label vMap["name"] = *v.Name vMap["project_name"] = *v.ProjectName + vMap["connections"] = *v.AllowedDbConnectionNames logger.DebugContext(ctx, "Converted to %v\n", vMap) data = append(data, vMap) } diff --git a/tests/looker/looker_integration_test.go b/tests/looker/looker_integration_test.go index d855872444..a27722cc50 100644 --- a/tests/looker/looker_integration_test.go +++ b/tests/looker/looker_integration_test.go @@ -199,6 +199,31 @@ func TestLooker(t *testing.T) { "source": "my-instance", "description": "Simple tool to test end to end functionality.", }, + "get_connections": map[string]any{ + "kind": "looker-get-connections", + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + }, + "get_connection_schemas": map[string]any{ + "kind": "looker-get-connection-schemas", + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + }, + "get_connection_databases": map[string]any{ + "kind": "looker-get-connection-databases", + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + }, + "get_connection_tables": map[string]any{ + "kind": "looker-get-connection-tables", + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + }, + "get_connection_table_columns": map[string]any{ + "kind": "looker-get-connection-table-columns", + "source": "my-instance", + "description": "Simple tool to test end to end functionality.", + }, }, } @@ -992,8 +1017,127 @@ func TestLooker(t *testing.T) { }, }, ) + tests.RunToolGetTestByName(t, "get_connections", + map[string]any{ + "get_connections": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "authRequired": []any{}, + "parameters": []any{}, + }, + }, + ) + tests.RunToolGetTestByName(t, "get_connection_schemas", + map[string]any{ + "get_connection_schemas": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "authRequired": []any{}, + "parameters": []any{ + map[string]any{ + "authSources": []any{}, + "description": "The connection containing the schemas.", + "name": "conn", + "required": true, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "The optional database to search", + "name": "db", + "required": false, + "type": "string", + }, + }, + }, + }, + ) + tests.RunToolGetTestByName(t, "get_connection_databases", + map[string]any{ + "get_connection_databases": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "authRequired": []any{}, + "parameters": []any{ + map[string]any{ + "authSources": []any{}, + "description": "The connection containing the databases.", + "name": "conn", + "required": true, + "type": "string", + }, + }, + }, + }, + ) + tests.RunToolGetTestByName(t, "get_connection_tables", + map[string]any{ + "get_connection_tables": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "authRequired": []any{}, + "parameters": []any{ + map[string]any{ + "authSources": []any{}, + "description": "The connection containing the tables.", + "name": "conn", + "required": true, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "The optional database to search", + "name": "db", + "required": false, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "The schema containing the tables.", + "name": "schema", + "required": true, + "type": "string", + }, + }, + }, + }, + ) + tests.RunToolGetTestByName(t, "get_connection_table_columns", + map[string]any{ + "get_connection_table_columns": map[string]any{ + "description": "Simple tool to test end to end functionality.", + "authRequired": []any{}, + "parameters": []any{ + map[string]any{ + "authSources": []any{}, + "description": "The connection containing the tables.", + "name": "conn", + "required": true, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "The optional database to search", + "name": "db", + "required": false, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "The schema containing the tables.", + "name": "schema", + "required": true, + "type": "string", + }, + map[string]any{ + "authSources": []any{}, + "description": "A comma separated list of tables containing the columns.", + "name": "tables", + "required": true, + "type": "string", + }, + }, + }, + }, + ) - wantResult := "{\"label\":\"System Activity\",\"name\":\"system__activity\",\"project_name\":\"system__activity\"}" + wantResult := "{\"connections\":[],\"label\":\"System Activity\",\"name\":\"system__activity\",\"project_name\":\"system__activity\"}" tests.RunToolInvokeSimpleTest(t, "get_models", wantResult) wantResult = "{\"description\":\"Data about Look and dashboard usage, including frequency of views, favoriting, scheduling, embedding, and access via the API. Also includes details about individual Looks and dashboards.\",\"group_label\":\"System Activity\",\"label\":\"Content Usage\",\"name\":\"content_usage\"}" @@ -1070,6 +1214,21 @@ func TestLooker(t *testing.T) { wantResult = "production" tests.RunToolInvokeParametersTest(t, "dev_mode", []byte(`{"devMode": false}`), wantResult) + + wantResult = "thelook" + tests.RunToolInvokeSimpleTest(t, "get_connections", wantResult) + + wantResult = "{\"name\":\"demo_db\",\"is_default\":true}" + tests.RunToolInvokeParametersTest(t, "get_connection_schemas", []byte(`{"conn": "thelook"}`), wantResult) + + wantResult = "[]" + tests.RunToolInvokeParametersTest(t, "get_connection_databases", []byte(`{"conn": "thelook"}`), wantResult) + + wantResult = "Employees" + tests.RunToolInvokeParametersTest(t, "get_connection_tables", []byte(`{"conn": "thelook", "schema": "demo_db"}`), wantResult) + + wantResult = "{\"column_name\":\"EmpID\",\"data_type_database\":\"int\",\"data_type_looker\":\"number\",\"sql_escaped_column_name\":\"EmpID\"}" + tests.RunToolInvokeParametersTest(t, "get_connection_table_columns", []byte(`{"conn": "thelook", "schema": "demo_db", "tables": "Employees"}`), wantResult) } func runConversationalAnalytics(t *testing.T, modelName, exploreName string) {