feat(source/alloydb-pg): add configuration for public and private IP

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)
 			}

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,23 +31,24 @@ 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 {
 	return SourceKind
 }
 
 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,16 +82,31 @@ 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)
 	if err != nil {
 		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)
 	}

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)
+			}
+		})
+	}
 }

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"`)
+	}
+}