|
1 | 1 | package provider |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "context" |
4 | 5 | "database/sql" |
5 | 6 | "fmt" |
| 7 | + "net" |
| 8 | + "net/url" |
6 | 9 | "sync" |
| 10 | + |
| 11 | + "cloud.google.com/go/cloudsqlconn" |
| 12 | + "cloud.google.com/go/cloudsqlconn/postgres/pgxv4" |
| 13 | + "github.com/hashicorp/terraform-plugin-log/tflog" |
| 14 | + "golang.org/x/net/proxy" |
| 15 | + "golang.org/x/oauth2" |
7 | 16 | ) |
8 | 17 |
|
9 | 18 | type Config struct { |
10 | | - dsnTemplate string |
11 | 19 | dbRegistry map[string]*sql.DB |
12 | 20 | dbRegistryMutex sync.Mutex |
| 21 | + dbDrivers map[string]bool |
| 22 | + dbDriversMutex sync.Mutex |
| 23 | + connections map[string]*ConnectionConfig |
13 | 24 | } |
14 | 25 |
|
15 | | -func NewConfig(dsnTemplate string) *Config { |
| 26 | +func NewConfig() *Config { |
16 | 27 | return &Config{ |
17 | | - dsnTemplate: dsnTemplate, |
18 | 28 | dbRegistry: make(map[string]*sql.DB), |
| 29 | + dbDrivers: make(map[string]bool), |
| 30 | + connections: make(map[string]*ConnectionConfig), |
19 | 31 | } |
20 | 32 | } |
21 | 33 |
|
22 | | -func (c *Config) connectToPostgresqlDb(dbName string) (*sql.DB, error) { |
23 | | - dsn := fmt.Sprintf(c.dsnTemplate, "dbname="+dbName) |
24 | | - return c.connectToPostgresql(dsn) |
25 | | -} |
26 | | - |
27 | | -func (c *Config) connectToPostgresqlNoDb() (*sql.DB, error) { |
28 | | - dsn := fmt.Sprintf(c.dsnTemplate, "dbname=postgres") |
29 | | - return c.connectToPostgresql(dsn) |
30 | | -} |
31 | | - |
32 | | -func (c *Config) connectToPostgresql(dsn string) (*sql.DB, error) { |
| 34 | +func (c *Config) connectToPostgresql(ctx context.Context, cc *ConnectionConfig) (*sql.DB, error) { |
33 | 35 | c.dbRegistryMutex.Lock() |
34 | 36 | defer c.dbRegistryMutex.Unlock() |
35 | 37 |
|
36 | | - if c.dbRegistry[dsn] != nil { |
37 | | - return c.dbRegistry[dsn], nil |
| 38 | + key := cc.DsnKey() |
| 39 | + |
| 40 | + if c.dbRegistry[key] != nil { |
| 41 | + return c.dbRegistry[key], nil |
| 42 | + } |
| 43 | + |
| 44 | + err := c.registerDriver(ctx, cc) |
| 45 | + if err != nil { |
| 46 | + return nil, err |
38 | 47 | } |
39 | 48 |
|
40 | | - db, err := sql.Open("cloudsql-postgres", dsn) |
| 49 | + db, err := sql.Open(cc.DriverKey(), cc.Dsn()) |
41 | 50 | if err != nil { |
42 | 51 | return nil, err |
43 | 52 | } |
44 | | - c.dbRegistry[dsn] = db |
45 | | - return c.dbRegistry[dsn], nil |
| 53 | + |
| 54 | + c.dbRegistry[key] = db |
| 55 | + return c.dbRegistry[key], nil |
| 56 | +} |
| 57 | + |
| 58 | +func (c *Config) registerDriver(ctx context.Context, cc *ConnectionConfig) error { |
| 59 | + c.dbDriversMutex.Lock() |
| 60 | + defer c.dbDriversMutex.Unlock() |
| 61 | + |
| 62 | + key := cc.DriverKey() |
| 63 | + |
| 64 | + if c.dbDrivers[key] { |
| 65 | + return nil |
| 66 | + } |
| 67 | + |
| 68 | + var ( |
| 69 | + dialOptions []cloudsqlconn.DialOption |
| 70 | + options []cloudsqlconn.Option |
| 71 | + ) |
| 72 | + |
| 73 | + if cc.PrivateIP.ValueBool() { |
| 74 | + dialOptions = append(dialOptions, cloudsqlconn.WithPrivateIP()) |
| 75 | + } |
| 76 | + |
| 77 | + if cc.PSC.ValueBool() { |
| 78 | + dialOptions = append(dialOptions, cloudsqlconn.WithPSC()) |
| 79 | + } |
| 80 | + |
| 81 | + options = append(options, cloudsqlconn.WithDefaultDialOptions(dialOptions...)) |
| 82 | + |
| 83 | + if cc.GoogleApiAccessToken.ValueString() != "" { |
| 84 | + token := &oauth2.Token{AccessToken: cc.GoogleApiAccessToken.ValueString()} |
| 85 | + options = append(options, cloudsqlconn.WithTokenSource(oauth2.StaticTokenSource(token))) |
| 86 | + } |
| 87 | + |
| 88 | + if !cc.Proxy.IsNull() { |
| 89 | + options = append(options, cloudsqlconn.WithDialFunc(createDialer(cc.Proxy.ValueString(), ctx))) |
| 90 | + } |
| 91 | + |
| 92 | + _, err := pgxv4.RegisterDriver(key, options...) |
| 93 | + if err != nil { |
| 94 | + return err |
| 95 | + } |
| 96 | + |
| 97 | + c.dbDrivers[key] = true |
| 98 | + return nil |
| 99 | +} |
| 100 | + |
| 101 | +func createDialer(proxyInput string, ctxProvider context.Context) func(ctx context.Context, network, addr string) (net.Conn, error) { |
| 102 | + return func(ctx context.Context, network, address string) (net.Conn, error) { |
| 103 | + tflog.Info(ctxProvider, "Creating Dialer with proxy: "+proxyInput) |
| 104 | + if len(proxyInput) == 0 { |
| 105 | + return nil, fmt.Errorf("proxy is empty") |
| 106 | + } |
| 107 | + |
| 108 | + proxyURL, err := url.Parse(proxyInput) |
| 109 | + if err != nil { |
| 110 | + return nil, err |
| 111 | + } |
| 112 | + d, err := proxy.FromURL(proxyURL, proxy.Direct) |
| 113 | + if err != nil { |
| 114 | + return nil, err |
| 115 | + } |
| 116 | + |
| 117 | + if xd, ok := d.(proxy.ContextDialer); ok { |
| 118 | + return xd.DialContext(ctx, network, address) |
| 119 | + } |
| 120 | + |
| 121 | + tflog.Warn(ctxProvider, "net.Conn created without context.Context") |
| 122 | + return d.Dial(network, address) // TODO: force use of context? |
| 123 | + } |
46 | 124 | } |
0 commit comments