Skip to content

Commit 61bcf21

Browse files
DJAndriesDJAndries
committed
Add RDS IAM authentication support
1 parent c751555 commit 61bcf21

5 files changed

Lines changed: 196 additions & 17 deletions

File tree

‎datastore/migrations/20240904202925_init.up.sql‎

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,28 +23,52 @@ CREATE TABLE entities (
2323
mtime BIGINT NOT NULL,
2424
version BIGINT NOT NULL,
2525
data_type INTEGER NOT NULL,
26-
specifics BYTEA STORAGE EXTERNAL NOT NULL,
27-
client_defined_unique_tag TEXT STORAGE PLAIN,
28-
server_defined_unique_tag TEXT STORAGE PLAIN,
29-
name TEXT STORAGE PLAIN,
30-
originator_cache_guid TEXT STORAGE PLAIN,
31-
originator_client_item_id TEXT STORAGE PLAIN,
32-
parent_id TEXT STORAGE PLAIN,
33-
non_unique_name TEXT STORAGE PLAIN,
34-
unique_position BYTEA STORAGE PLAIN,
26+
specifics BYTEA NOT NULL,
27+
client_defined_unique_tag TEXT,
28+
server_defined_unique_tag TEXT,
29+
name TEXT,
30+
originator_cache_guid TEXT,
31+
originator_client_item_id TEXT,
32+
parent_id TEXT,
33+
non_unique_name TEXT,
34+
unique_position BYTEA,
3535
folder BOOLEAN,
3636
deleted BOOLEAN NOT NULL,
3737
PRIMARY KEY (id, chain_id),
3838
UNIQUE (chain_id, client_defined_unique_tag)
3939
)
4040
PARTITION BY RANGE (chain_id);
41+
42+
ALTER TABLE entities ALTER specifics SET STORAGE EXTERNAL;
43+
ALTER TABLE entities ALTER client_defined_unique_tag SET STORAGE PLAIN;
44+
ALTER TABLE entities ALTER server_defined_unique_tag SET STORAGE PLAIN;
45+
ALTER TABLE entities ALTER name SET STORAGE PLAIN;
46+
ALTER TABLE entities ALTER originator_cache_guid SET STORAGE PLAIN;
47+
ALTER TABLE entities ALTER originator_client_item_id SET STORAGE PLAIN;
48+
ALTER TABLE entities ALTER parent_id SET STORAGE PLAIN;
49+
ALTER TABLE entities ALTER non_unique_name SET STORAGE PLAIN;
50+
ALTER TABLE entities ALTER unique_position SET STORAGE PLAIN;
51+
4152
CREATE INDEX entities_chain_id_data_type_mtime_idx ON entities (chain_id, data_type, mtime);
4253

43-
SELECT partman.create_parent(
44-
p_parent_table := 'public.entities',
45-
p_control := 'chain_id',
46-
p_interval := '3500'
47-
);
54+
DO $$
55+
BEGIN
56+
-- for vanilla postgres
57+
PERFORM partman.create_parent(
58+
p_parent_table := 'public.entities',
59+
p_control := 'chain_id',
60+
p_interval := '3500',
61+
p_type := 'range'
62+
);
63+
EXCEPTION WHEN OTHERS THEN
64+
-- for Aurora
65+
PERFORM partman.create_parent(
66+
p_parent_table := 'public.entities',
67+
p_control := 'chain_id',
68+
p_interval := '3500',
69+
p_type := 'native'
70+
);
71+
END $$;
4872

4973
CREATE EXTENSION IF NOT EXISTS pg_cron;
5074

‎datastore/rds.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package datastore
2+
3+
import (
4+
"context"
5+
"database/sql/driver"
6+
"fmt"
7+
"net/url"
8+
"os"
9+
"sync"
10+
"time"
11+
12+
"github.com/aws/aws-sdk-go-v2/config"
13+
"github.com/aws/aws-sdk-go-v2/feature/rds/auth"
14+
"github.com/jackc/pgx/stdlib"
15+
)
16+
17+
const defaultRegion = "us-west-2"
18+
19+
const (
20+
rdsPortKey = "RDS_DATABASE_PORT"
21+
rdsHostKey = "RDS_WRITER_ENDPOINT"
22+
rdsUserKey = "RDS_USER"
23+
rdsDbNameKey = "RDS_DATABASE_NAME"
24+
regionKey = "AWS_REGION"
25+
)
26+
27+
type rdsConnector struct {
28+
hostAndPort string
29+
dbName string
30+
user string
31+
token string
32+
region string
33+
tokenCacheTime time.Time
34+
mu sync.Mutex
35+
}
36+
37+
func newRDSConnector() *rdsConnector {
38+
port := os.Getenv(rdsPortKey)
39+
host := os.Getenv(rdsHostKey)
40+
user := os.Getenv(rdsUserKey)
41+
dbName := os.Getenv(rdsDbNameKey)
42+
region := os.Getenv(regionKey)
43+
44+
if region == "" {
45+
region = defaultRegion
46+
}
47+
hostAndPort := fmt.Sprintf("%s:%s", host, port)
48+
return &rdsConnector{
49+
hostAndPort: hostAndPort,
50+
dbName: dbName,
51+
user: user,
52+
region: region,
53+
}
54+
}
55+
56+
func (c *rdsConnector) getConnectionString(ctx context.Context) (string, error) {
57+
c.mu.Lock()
58+
defer c.mu.Unlock()
59+
60+
if time.Since(c.tokenCacheTime) > 10*time.Minute {
61+
cfg, err := config.LoadDefaultConfig(ctx)
62+
if err != nil {
63+
return "", fmt.Errorf("failed to load AWS config")
64+
}
65+
66+
token, err := auth.BuildAuthToken(
67+
ctx, c.hostAndPort, c.region, c.user, cfg.Credentials)
68+
if err != nil {
69+
return "", fmt.Errorf("failed to create authentication token: %w", err)
70+
}
71+
c.token = url.QueryEscape(token)
72+
c.tokenCacheTime = time.Now()
73+
}
74+
75+
return fmt.Sprintf("postgres://%s:%s@%s/%s?sslmode=require", c.user, c.token, c.hostAndPort, c.dbName), nil
76+
}
77+
78+
func (c *rdsConnector) Connect(ctx context.Context) (driver.Conn, error) {
79+
connStr, err := c.getConnectionString(ctx)
80+
if err != nil {
81+
return nil, err
82+
}
83+
84+
return stdlib.GetDefaultDriver().Open(connStr)
85+
}
86+
87+
func (c *rdsConnector) Driver() driver.Driver {
88+
return c
89+
}
90+
91+
func (c *rdsConnector) Open(_ string) (driver.Conn, error) {
92+
return nil, fmt.Errorf("open method unsupported")
93+
}

‎datastore/sql.go‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package datastore
22

33
import (
4+
"context"
5+
"database/sql"
46
"embed"
57
"errors"
68
"fmt"
@@ -11,6 +13,7 @@ import (
1113
// import postgres package for migrations
1214
_ "github.com/golang-migrate/migrate/v4/database/postgres"
1315
"github.com/golang-migrate/migrate/v4/source/iofs"
16+
1417
// import pgx so it can be used with sqlx
1518
_ "github.com/jackc/pgx/stdlib"
1619
"github.com/jmoiron/sqlx"
@@ -58,12 +61,22 @@ func NewSQLDB(isTesting bool) (*SQLDB, error) {
5861
envKey = sqlURLEnvKey
5962
}
6063

64+
var rdsConnector *rdsConnector
65+
if os.Getenv(rdsHostKey) != "" {
66+
rdsConnector = newRDSConnector()
67+
}
68+
6169
sqlURL := os.Getenv(envKey)
62-
if sqlURL == "" {
70+
if rdsConnector != nil {
71+
sqlURL, err = rdsConnector.getConnectionString(context.Background())
72+
if err != nil {
73+
return nil, err
74+
}
75+
} else if sqlURL == "" {
6376
if isTesting {
6477
sqlURL = defaultSQLTestURL
6578
} else {
66-
return nil, fmt.Errorf("%s must be defined", envKey)
79+
return nil, fmt.Errorf("%s or %s must be defined", envKey, rdsHostKey)
6780
}
6881
}
6982
iofsDriver, err := iofs.New(migrationFiles, "migrations")
@@ -82,9 +95,16 @@ func NewSQLDB(isTesting bool) (*SQLDB, error) {
8295
if !errors.Is(err, migrate.ErrNoChange) {
8396
return nil, fmt.Errorf("Failed to run migrations: %w", err)
8497
}
98+
err = nil
8599
}
86100

87-
db, err := sqlx.Connect("pgx", sqlURL)
101+
var db *sqlx.DB
102+
if rdsConnector != nil {
103+
baseDB := sql.OpenDB(rdsConnector)
104+
db = sqlx.NewDb(baseDB, "pgx")
105+
} else {
106+
db, err = sqlx.Connect("pgx", sqlURL)
107+
}
88108
if err != nil {
89109
return nil, fmt.Errorf("Failed to connect to SQL DB: %w", err)
90110
}

‎go.mod‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,20 @@ require (
2525

2626
require (
2727
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
28+
github.com/aws/aws-sdk-go-v2 v1.32.2 // indirect
29+
github.com/aws/aws-sdk-go-v2/config v1.28.0 // indirect
30+
github.com/aws/aws-sdk-go-v2/credentials v1.17.41 // indirect
31+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.17 // indirect
32+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.4.21 // indirect
33+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.21 // indirect
34+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.21 // indirect
35+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 // indirect
36+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0 // indirect
37+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.2 // indirect
38+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.2 // indirect
39+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2 // indirect
40+
github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 // indirect
41+
github.com/aws/smithy-go v1.22.0 // indirect
2842
github.com/beorn7/perks v1.0.1 // indirect
2943
github.com/brave-intl/bat-go/libs v0.0.0-20240909083638-be56e4a5398e // indirect
3044
github.com/cespare/xxhash/v2 v2.2.0 // indirect

‎go.sum‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,34 @@ github.com/aws/aws-sdk-go v1.54.19 h1:tyWV+07jagrNiCcGRzRhdtVjQs7Vy41NwsuOcl0IbV
3131
github.com/aws/aws-sdk-go v1.54.19/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
3232
github.com/aws/aws-sdk-go v1.55.5 h1:KKUZBfBoyqy5d3swXyiC7Q76ic40rYcbqH7qjh59kzU=
3333
github.com/aws/aws-sdk-go v1.55.5/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
34+
github.com/aws/aws-sdk-go-v2 v1.32.2 h1:AkNLZEyYMLnx/Q/mSKkcMqwNFXMAvFto9bNsHqcTduI=
35+
github.com/aws/aws-sdk-go-v2 v1.32.2/go.mod h1:2SK5n0a2karNTv5tbP1SjsX0uhttou00v/HpXKM1ZUo=
36+
github.com/aws/aws-sdk-go-v2/config v1.28.0 h1:FosVYWcqEtWNxHn8gB/Vs6jOlNwSoyOCA/g/sxyySOQ=
37+
github.com/aws/aws-sdk-go-v2/config v1.28.0/go.mod h1:pYhbtvg1siOOg8h5an77rXle9tVG8T+BWLWAo7cOukc=
38+
github.com/aws/aws-sdk-go-v2/credentials v1.17.41 h1:7gXo+Axmp+R4Z+AK8YFQO0ZV3L0gizGINCOWxSLY9W8=
39+
github.com/aws/aws-sdk-go-v2/credentials v1.17.41/go.mod h1:u4Eb8d3394YLubphT4jLEwN1rLNq2wFOlT6OuxFwPzU=
40+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.17 h1:TMH3f/SCAWdNtXXVPPu5D6wrr4G5hI1rAxbcocKfC7Q=
41+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.17/go.mod h1:1ZRXLdTpzdJb9fwTMXiLipENRxkGMTn1sfKexGllQCw=
42+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.4.21 h1:wRH9E07mfYqZ1EPphNTUIkrZ/7wcbZAGcjhrBlkWy4c=
43+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.4.21/go.mod h1:6m/MDzT+aFxaIo46f2MYV4d+qG9J9keLlHL0qKnQFgA=
44+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.21 h1:UAsR3xA31QGf79WzpG/ixT9FZvQlh5HY1NRqSHBNOCk=
45+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.21/go.mod h1:JNr43NFf5L9YaG3eKTm7HQzls9J+A9YYcGI5Quh1r2Y=
46+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.21 h1:6jZVETqmYCadGFvrYEQfC5fAQmlo80CeL5psbno6r0s=
47+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.21/go.mod h1:1SR0GbLlnN3QUmYaflZNiH1ql+1qrSiB2vwcJ+4UM60=
48+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 h1:VaRN3TlFdd6KxX1x3ILT5ynH6HvKgqdiXoTxAF4HQcQ=
49+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc=
50+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0 h1:TToQNkvGguu209puTojY/ozlqy2d/SFNcoLIqTFi42g=
51+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.0/go.mod h1:0jp+ltwkf+SwG2fm/PKo8t4y8pJSgOCO4D8Lz3k0aHQ=
52+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.2 h1:s7NA1SOw8q/5c0wr8477yOPp0z+uBaXBnLE0XYb0POA=
53+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.2/go.mod h1:fnjjWyAW/Pj5HYOxl9LJqWtEwS7W2qgcRLWP+uWbss0=
54+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.2 h1:bSYXVyUzoTHoKalBmwaZxs97HU9DWWI3ehHSAMa7xOk=
55+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.2/go.mod h1:skMqY7JElusiOUjMJMOv1jJsP7YUg7DrhgqZZWuzu1U=
56+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2 h1:AhmO1fHINP9vFYUE0LHzCWg/LfUWUF+zFPEcY9QXb7o=
57+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2/go.mod h1:o8aQygT2+MVP0NaV6kbdE1YnnIM8RRVQzoeUH45GOdI=
58+
github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 h1:CiS7i0+FUe+/YY1GvIBLLrR/XNGZ4CtM1Ll0XavNuVo=
59+
github.com/aws/aws-sdk-go-v2/service/sts v1.32.2/go.mod h1:HtaiBI8CjYoNVde8arShXb94UbQQi9L4EMr6D+xGBwo=
60+
github.com/aws/smithy-go v1.22.0 h1:uunKnWlcoL3zO7q+gG2Pk53joueEOsnNB28QdMsmiMM=
61+
github.com/aws/smithy-go v1.22.0/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
3462
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
3563
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
3664
github.com/brave-intl/bat-go/libs v0.0.0-20231020145457-cc9860c87bae h1:CGUFAtMXAsGajLeobq6ep+5wREYS+lepZSdPckY+Ba0=

0 commit comments

Comments
 (0)