From e88ec409d14c85d6b0896c45d9957cce9097912a Mon Sep 17 00:00:00 2001 From: Yuan <45984206+Yuan325@users.noreply.github.com> Date: Thu, 5 Dec 2024 15:05:53 -0800 Subject: [PATCH] feat(source/alloydb-pg): add configuration for public and private IP (#103) Allow user to set if their database uses private or public ip. The reason we add this is because the dialer require different initialization with private and public ip. By default, toolbox will use public ip. --------- Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com> --- internal/server/config.go | 2 +- internal/sources/alloydbpg/alloydb_pg.go | 43 +++++++++---- internal/sources/alloydbpg/alloydb_pg_test.go | 60 +++++++++++++++++++ internal/sources/ip_type.go | 45 ++++++++++++++ 4 files changed, 136 insertions(+), 14 deletions(-) create mode 100644 internal/sources/ip_type.go diff --git a/internal/server/config.go b/internal/server/config.go index 1730a47cd..9d7b44ca9 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -121,7 +121,7 @@ func (c *SourceConfigs) UnmarshalYAML(node *yaml.Node) error { } switch k.Kind { case alloydbpgsrc.SourceKind: - actual := alloydbpgsrc.Config{Name: name} + actual := alloydbpgsrc.Config{Name: name, IP_type: "public"} if err := n.Decode(&actual); err != nil { return fmt.Errorf("unable to parse as %q: %w", k.Kind, err) } diff --git a/internal/sources/alloydbpg/alloydb_pg.go b/internal/sources/alloydbpg/alloydb_pg.go index 7fbc6f3f7..5928cc902 100644 --- a/internal/sources/alloydbpg/alloydb_pg.go +++ b/internal/sources/alloydbpg/alloydb_pg.go @@ -18,6 +18,7 @@ import ( "context" "fmt" "net" + "strings" "cloud.google.com/go/alloydbconn" "github.com/googleapis/genai-toolbox/internal/sources" @@ -30,15 +31,16 @@ const SourceKind string = "alloydb-postgres" var _ sources.SourceConfig = Config{} type Config struct { - Name string `yaml:"name"` - Kind string `yaml:"kind"` - Project string `yaml:"project"` - Region string `yaml:"region"` - Cluster string `yaml:"cluster"` - Instance string `yaml:"instance"` - User string `yaml:"user"` - Password string `yaml:"password"` - Database string `yaml:"database"` + Name string `yaml:"name"` + Kind string `yaml:"kind"` + Project string `yaml:"project"` + Region string `yaml:"region"` + Cluster string `yaml:"cluster"` + Instance string `yaml:"instance"` + IP_type sources.IPType `yaml:"ip_type"` + User string `yaml:"user"` + Password string `yaml:"password"` + Database string `yaml:"database"` } func (r Config) SourceConfigKind() string { @@ -46,7 +48,7 @@ func (r Config) SourceConfigKind() string { } func (r Config) Initialize() (sources.Source, error) { - pool, err := initAlloyDBPgConnectionPool(r.Project, r.Region, r.Cluster, r.Instance, r.User, r.Password, r.Database) + pool, err := initAlloyDBPgConnectionPool(r.Project, r.Region, r.Cluster, r.Instance, r.IP_type.String(), r.User, r.Password, r.Database) if err != nil { return nil, fmt.Errorf("unable to create pool: %w", err) } @@ -80,7 +82,18 @@ func (s *Source) PostgresPool() *pgxpool.Pool { return s.Pool } -func initAlloyDBPgConnectionPool(project, region, cluster, instance, user, pass, dbname string) (*pgxpool.Pool, error) { +func getDialOpts(ip_type string) ([]alloydbconn.DialOption, error) { + switch strings.ToLower(ip_type) { + case "private": + return []alloydbconn.DialOption{alloydbconn.WithPrivateIP()}, nil + case "public": + return []alloydbconn.DialOption{alloydbconn.WithPublicIP()}, nil + default: + return nil, fmt.Errorf("invalid ip_type %s", ip_type) + } +} + +func initAlloyDBPgConnectionPool(project, region, cluster, instance, ip_type, user, pass, dbname string) (*pgxpool.Pool, error) { // Configure the driver to connect to the database dsn := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=disable", user, pass, dbname) config, err := pgxpool.ParseConfig(dsn) @@ -88,8 +101,12 @@ func initAlloyDBPgConnectionPool(project, region, cluster, instance, user, pass, return nil, fmt.Errorf("unable to parse connection uri: %w", err) } - // Create a new dialer with any options - d, err := alloydbconn.NewDialer(context.Background()) + // Create a new dialer with options + dialOpts, err := getDialOpts(ip_type) + if err != nil { + return nil, err + } + d, err := alloydbconn.NewDialer(context.Background(), alloydbconn.WithDefaultDialOptions(dialOpts...)) if err != nil { return nil, fmt.Errorf("unable to parse connection uri: %w", err) } diff --git a/internal/sources/alloydbpg/alloydb_pg_test.go b/internal/sources/alloydbpg/alloydb_pg_test.go index e0b311077..bb5760ce6 100644 --- a/internal/sources/alloydbpg/alloydb_pg_test.go +++ b/internal/sources/alloydbpg/alloydb_pg_test.go @@ -51,6 +51,33 @@ func TestParseFromYamlAlloyDBPg(t *testing.T) { Region: "my-region", Cluster: "my-cluster", Instance: "my-instance", + IP_type: "public", + Database: "my_db", + }, + }, + }, + { + desc: "private ip_type", + in: ` + sources: + my-pg-instance: + kind: alloydb-postgres + project: my-project + region: my-region + cluster: my-cluster + instance: my-instance + ip_type: private + database: my_db + `, + want: map[string]sources.SourceConfig{ + "my-pg-instance": alloydbpg.Config{ + Name: "my-pg-instance", + Kind: alloydbpg.SourceKind, + Project: "my-project", + Region: "my-region", + Cluster: "my-cluster", + Instance: "my-instance", + IP_type: "private", Database: "my_db", }, }, @@ -71,5 +98,38 @@ func TestParseFromYamlAlloyDBPg(t *testing.T) { } }) } +} +func FailParseFromYamlAlloyDBPg(t *testing.T) { + tcs := []struct { + desc string + in string + }{ + { + desc: "invalid ip_type", + in: ` + sources: + my-pg-instance: + kind: alloydb-postgres + project: my-project + region: my-region + cluster: my-cluster + instance: my-instance + ip_type: fail + database: my_db + `, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Sources server.SourceConfigs `yaml:"sources"` + }{} + // Parse contents + err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail: %s", err) + } + }) + } } diff --git a/internal/sources/ip_type.go b/internal/sources/ip_type.go new file mode 100644 index 000000000..546eabd5b --- /dev/null +++ b/internal/sources/ip_type.go @@ -0,0 +1,45 @@ +// Copyright 2024 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 sources + +import ( + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +type IPType string + +func (i *IPType) String() string { + if string(*i) != "" { + return strings.ToLower(string(*i)) + } + return "public" +} + +func (i *IPType) UnmarshalYAML(node *yaml.Node) error { + var ip_type string + if err := node.Decode(&ip_type); err != nil { + return err + } + switch ip_type { + case "private", "public": + *i = IPType(ip_type) + return nil + default: + return fmt.Errorf(`ip_type invalid: must be one of "public", or "private"`) + } +}