Add code to support creating storage with separate reader/writer (#281)

* Add ProvideRW

* comment

* confi -> c

Co-authored-by: Yuchen Wang <[email protected]>

go/config/config.go

@@ -23,8 +23,9 @@ const (
 // json tags are required because
 // config.ConvertGenericConfigToSpecificType internally uses json package.
 type Config struct {
-	Host string `json:"host"`
-	Port int    `json:"port"`
+	Host   string `json:"host"`
+	Reader string `json:"reader"`
+	Port   int    `json:"port"`
 	// For rds_prostgres, DBName has to alrady exist and can be accessed by User.
 	DBName   string `json:"db_name"`
 	User     string `json:"user"`

go/config/config_test.go

@@ -26,6 +26,7 @@ func TestNewConfig(t *testing.T) {
 			file: "valid.yaml",
 			wantConfig: Config{
 				Host:          "some-host.rds.amazonaws.com",
+				Reader:        "some-host-ro.rds.amazonaws.com",
 				Port:          defaultPort,
 				DBName:        defaultDBName,
 				User:          "grafeas_rw",

go/config/testdata/invalid_api_endpoint.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/config/testdata/invalid_missing_athenz_domain.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/config/testdata/invalid_missing_iam_role.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/config/testdata/invalid_missing_region.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/config/testdata/invalid_missing_ssl_root_cert.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     pagination_key: "some_random_key"
     conn_pool:

go/config/testdata/invalid_missing_user.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"
     conn_pool:

go/config/testdata/invalid_port.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     port: -1
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"

go/config/testdata/invalid_renew_threshold.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/config/testdata/valid.yaml

@@ -5,6 +5,7 @@ grafeas:
   storage_type: "rds"
   rds:
     host: "some-host.rds.amazonaws.com"
+    reader: "some-host-ro.rds.amazonaws.com"
     user: "grafeas_rw"
     ssl_root_cert: "/opt/rds-ca-2019-root.pem"
     pagination_key: "some_random_key"

go/v1beta1/storage/connector.go

@@ -49,9 +49,13 @@ type connector struct {
 	dsnLock sync.RWMutex
 }
 
-func newConnector(ctx context.Context, conf *config.Config, driver driver.Driver, cc CredentialsCreator, logger *log.Logger) (*connector, error) {
+func newConnector(ctx context.Context, conf *config.Config, driver driver.Driver, cc CredentialsCreator, logger *log.Logger, overwriteHost string) (*connector, error) {
+	host := conf.Host
+	if overwriteHost != "" {
+		host = overwriteHost
+	}
 	c := &connector{
-		host:        conf.Host,
+		host:        host,
 		port:        conf.Port,
 		dbName:      conf.DBName,
 		user:        conf.User,

go/v1beta1/storage/connector_test.go

@@ -63,7 +63,7 @@ func TestNewConnector(t *testing.T) {
 			}
 			var buf bytes.Buffer
 			logger := log.New(&buf, "", 0)
-			c, err := newConnector(ctx, &conf, mockDriver, cc, logger)
+			c, err := newConnector(ctx, &conf, mockDriver, cc, logger, "")
 			if (err == nil) != (tt.wantErrMsg == "") {
 				if err == nil {
 					t.Error("want error, but no error is returned")

go/v1beta1/storage/storage.go

@@ -44,6 +44,7 @@ type Storage interface {
 // StorageCreator can be implemented based on the backend storage types (e.g. PostgreSQL, MySQL, etc.).
 type StorageCreator interface {
 	Create(connector driver.Connector, paginationKey string) (Storage, error)
+	CreateRW(readerConnector driver.Connector, writerConnector driver.Connector, paginationKey string) (Storage, error)
 }
 
 // CredentialsCreator can be implemented to create Credentials based on different types of providers.
@@ -77,7 +78,7 @@ func (p GrafeasStorageProvider) Provide(_ string, confi *config.StorageConfigura
 
 	// TODO: Use the context passed from main after
 	// the signature of RegisterStorageTypeProvider is updated to include it.
-	connector, err := newConnector(context.Background(), conf, p.drv, p.credentialsCreator, log.Default())
+	connector, err := newConnector(context.Background(), conf, p.drv, p.credentialsCreator, log.Default(), "")
 	if err != nil {
 		return nil, fmt.Errorf("%s, err: %v", errMsgInitConnector, err)
 	}
@@ -95,6 +96,37 @@ func (p GrafeasStorageProvider) Provide(_ string, confi *config.StorageConfigura
 	return grafeasStorage, nil
 }
 
+// ProvideRW returns a storage which is configured based on the receiver's fields. The storage connects to different reader/writer. If no reader is provided, then it will only connect to the writer.
+func (p GrafeasStorageProvider) ProvideRW(_ string, c *config.StorageConfiguration) (*storage.Storage, error) {
+	conf, err := rdsconfig.New(c)
+	if err != nil {
+		return nil, fmt.Errorf("%s, err: %v", errMsgInitConfig, err)
+	}
+
+	// TODO: Use the context passed from main after
+	// the signature of RegisterStorageTypeProvider is updated to include it.
+	writerConnector, err := newConnector(context.Background(), conf, p.drv, p.credentialsCreator, log.Default(), "")
+	if err != nil {
+		return nil, fmt.Errorf("%s, err: %v", errMsgInitConnector, err)
+	}
+	readerConnector, err := newConnector(context.Background(), conf, p.drv, p.credentialsCreator, log.Default(), conf.Reader)
+	if err != nil {
+		return nil, fmt.Errorf("%s, err: %v", errMsgInitConnector, err)
+	}
+
+	rdsStorage, err := p.storageCreator.CreateRW(readerConnector, writerConnector, conf.PaginationKey)
+	if err != nil {
+		return nil, fmt.Errorf("%s, err: %v", errMsgInitStorage, err)
+	}
+	setConnPoolParams(rdsStorage, conf.ConnPool)
+
+	grafeasStorage := &storage.Storage{
+		Ps: rdsStorage,
+		Gs: rdsStorage,
+	}
+	return grafeasStorage, nil
+}
+
 func setConnPoolParams(mgr ConnPoolMgr, conf rdsconfig.ConnPoolConfig) {
 	mgr.SetMaxOpenConns(conf.MaxOpenConns)
 	mgr.SetMaxIdleConns(conf.MaxIdleConns)

go/v1beta1/storage/storage_test.go

@@ -120,3 +120,108 @@ func TestStorageProviderProvide(t *testing.T) {
 		})
 	}
 }
+
+func TestStorageProviderProvideRW(t *testing.T) {
+	t.Parallel()
+
+	mockCtrl := gomock.NewController(t)
+	validConf := config.StorageConfiguration(rdsconfig.Config{
+		Host:        "some-host.rds.amazonaws.com",
+		Reader:      "some-host-ro.rds.amazonaws.com",
+		User:        "grafeas_rw",
+		Password:    "dummy-password-for-unit-tests-only",
+		SSLRootCert: "/opt/rds-ca-2019-root.pem",
+		// ConnPool is populated in order to test if setConnPoolParams is invoked in Provide.
+		ConnPool: rdsconfig.ConnPoolConfig{
+			MaxOpenConns:             1,
+			MaxIdleConns:             2,
+			ConnMaxLifetimeInSeconds: 3,
+			ConnMaxIdleTimeInSeconds: 4,
+		},
+	})
+
+	type testCase struct {
+		name         string
+		expect       func(*testCase)
+		conf         config.StorageConfiguration
+		store        *mocks.MockStorage
+		storeCreator *MockStorageCreator
+		credsCreator *mocks.MockCredentialsCreator
+		wantErrMsg   string
+	}
+	tests := []testCase{
+		{
+			name: "happy path",
+			expect: func(tt *testCase) {
+				tt.credsCreator.EXPECT().Create(gomock.Any()).Times(2).Return(credentials.NewStaticCredentials("a", "b", "c"), nil)
+				tt.storeCreator.EXPECT().CreateRW(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(tt.store, nil)
+
+				conf := tt.conf.(rdsconfig.Config).ConnPool
+				tt.store.EXPECT().SetMaxOpenConns(conf.MaxOpenConns).Times(1)
+				tt.store.EXPECT().SetMaxIdleConns(conf.MaxIdleConns).Times(1)
+				tt.store.EXPECT().SetConnMaxLifetime(time.Duration(conf.ConnMaxLifetimeInSeconds) * time.Second).Times(1)
+				tt.store.EXPECT().SetConnMaxIdleTime(time.Duration(conf.ConnMaxIdleTimeInSeconds) * time.Second).Times(1)
+			},
+			conf:         validConf,
+			store:        mocks.NewMockStorage(mockCtrl),
+			storeCreator: NewMockStorageCreator(mockCtrl),
+			credsCreator: mocks.NewMockCredentialsCreator(mockCtrl),
+		},
+		{
+			name: "invalid config",
+			// An empty Config is invalid because the Host field does not have a default value.
+			conf:       config.StorageConfiguration(rdsconfig.Config{}),
+			wantErrMsg: errMsgInitConfig,
+		},
+		{
+			name: "invalid connector",
+			expect: func(tt *testCase) {
+				tt.credsCreator.EXPECT().Create(gomock.Any()).Times(1).Return(credentials.AnonymousCredentials, nil)
+			},
+			conf:         validConf,
+			credsCreator: mocks.NewMockCredentialsCreator(mockCtrl),
+			wantErrMsg:   errMsgInitConnector,
+		},
+		{
+			name: "invalid store",
+			expect: func(tt *testCase) {
+				tt.credsCreator.EXPECT().Create(gomock.Any()).Times(2).Return(credentials.NewStaticCredentials("a", "b", "c"), nil)
+				tt.storeCreator.EXPECT().CreateRW(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil, errors.New("random error"))
+			},
+			conf:         validConf,
+			storeCreator: NewMockStorageCreator(mockCtrl),
+			credsCreator: mocks.NewMockCredentialsCreator(mockCtrl),
+			wantErrMsg:   errMsgInitStorage,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			if tt.expect != nil {
+				tt.expect(&tt)
+			}
+			storageProvider := NewGrafeasStorageProvider(mocks.NewMockDriver(mockCtrl), tt.credsCreator, tt.storeCreator)
+			storage, err := storageProvider.ProvideRW("", &tt.conf)
+			if (err != nil) != (tt.wantErrMsg != "") {
+				if err != nil {
+					t.Errorf("don't want error, but got %q", err)
+				} else {
+					t.Errorf("got nil error, but want error to include %q", tt.wantErrMsg)
+				}
+				return
+			}
+			if err != nil {
+				if !strings.Contains(err.Error(), tt.wantErrMsg) {
+					t.Errorf("want %q to include %q", err.Error(), tt.wantErrMsg)
+				}
+				return
+			}
+
+			if storage.Gs != tt.store || storage.Ps != tt.store {
+				t.Errorf("unexpected fields: %v", storage)
+			}
+		})
+	}
+}