/
iam_db.go
77 lines (64 loc) · 2.5 KB
/
iam_db.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
package storage
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"fmt"
"net/url"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/rds/rdsutils"
"github.com/jmoiron/sqlx"
"github.com/lib/pq"
)
// iamDb is a custom struct that satisfies the driver.Connector so that it can be used with sql.OpenDB.
// It implements a custom Connect() function that will get a new IAM auth token for each connection.
type iamDb struct {
config DBConfig
awsSession *session.Session
}
// Connect implements the functionality that is executed each time a new connection is made to the database
// For this reason, each new connection actually generates a new auth token, as each auth token only lasts 15 minutes
// Connections made with an auth token will _NOT_ drop when the token expires, so the 15 minute expiry should never matter
// This function helps satisfy the driver.Connector interface
func (idb *iamDb) Connect(ctx context.Context) (driver.Conn, error) {
awsRegion := *idb.awsSession.Config.Region
awsCreds := idb.awsSession.Config.Credentials
dbEndpoint := fmt.Sprintf("%s:%s", idb.config.Host, idb.config.Port)
authToken, err := rdsutils.BuildAuthToken(dbEndpoint, awsRegion, idb.config.Username, awsCreds)
if err != nil {
return nil, err
}
psqlURL := url.URL{
Scheme: "postgres",
Host: dbEndpoint,
User: url.UserPassword(idb.config.Username, authToken),
Path: idb.config.Database,
}
q := psqlURL.Query()
q.Add("sslmode", idb.config.SSLMode)
psqlURL.RawQuery = q.Encode()
psqlDriver := &pq.Driver{}
connector, err := psqlDriver.Open(psqlURL.String())
if err != nil {
return nil, err
}
return connector, nil
}
// Driver returns IAM DB instance, as it satisfies the driver.Driver interface
// This function helps satisfy the driver.Connector interface
func (idb *iamDb) Driver() driver.Driver {
return idb
}
// Open would normally return a new connection to the DB, but this custom IAM DB
// doesn't support it in favor of using `sql.OpenDB` in `newConnectionPoolWithIam`
// This function helps satisfy the driver.Driver interface
func (idb *iamDb) Open(name string) (driver.Conn, error) {
return nil, errors.New("driver open method not supported")
}
// newConnectionPoolWithIam opens a sql.DB using the custom iamDb as a connector.
// It will wrap that sql.DB and return it as a *sqlx.DB
func newConnectionPoolWithIam(awsSession *session.Session, config DBConfig) *sqlx.DB {
db := sql.OpenDB(&iamDb{config, awsSession})
return sqlx.NewDb(db, "postgres")
}