Adding in-memory cache parameters

cmd/baton-vgs/main.go

@@ -9,32 +9,39 @@
 	"github.com/conductorone/baton-sdk/pkg/connectorbuilder"
 	"github.com/conductorone/baton-sdk/pkg/field"
 	"github.com/conductorone/baton-sdk/pkg/types"
+	"github.com/conductorone/baton-vgs/pkg/client"
 	"github.com/conductorone/baton-vgs/pkg/connector"
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 )
 
 const (
-	version                    = "dev"
-	connectorName              = "baton-vgs"
-	serviceAccountClientId     = "service-account-client-id"
-	serviceAccountClientSecret = "service-account-client-secret"
-	organizationId             = "organization-id"
-	vault                      = "vault"
+	version           = "dev"
+	connectorName     = "baton-vgs"
+	batonCacheDisable = "cache-disable"
+	batonCacheTTL     = "cache-ttl"
+	batonCacheMaxSize = "cache-max-size"
 )
 
 var (
-	ServiceAccountClientId     = field.StringField(serviceAccountClientId, field.WithRequired(true), field.WithDescription("The VGS client id."))
-	ServiceAccountClientSecret = field.StringField(serviceAccountClientSecret, field.WithRequired(true), field.WithDescription("The VGS client secret."))
-	OrganizationId             = field.StringField(organizationId, field.WithRequired(true), field.WithDescription("The VGS organization id."))
-	Vault                      = field.StringField(vault, field.WithRequired(true), field.WithDescription("The VGS vault id."))
-	configurationFields        = []field.SchemaField{Vault, ServiceAccountClientId, ServiceAccountClientSecret, OrganizationId}
+	ServiceAccountClientId     = field.StringField(client.ServiceAccountClientIdName, field.WithRequired(true), field.WithDescription("The VGS client id."))
+	ServiceAccountClientSecret = field.StringField(client.ServiceAccountClientSecretName, field.WithRequired(true), field.WithDescription("The VGS client secret."))
+	OrganizationId             = field.StringField(client.OrganizationId, field.WithRequired(true), field.WithDescription("The VGS organization id."))
+	Vault                      = field.StringField(client.VaultId, field.WithRequired(true), field.WithDescription("The VGS vault id."))
+	CacheDisabled              = field.StringField(batonCacheDisable, field.WithRequired(false), field.WithDescription("Verbose mode shows information about new memory allocation."))
+	CacheTTL                   = field.StringField(batonCacheTTL, field.WithRequired(false), field.WithDescription("Time after which entry can be evicted."))
+	CacheMaxSize               = field.StringField(batonCacheMaxSize, field.WithRequired(false), field.WithDescription("It is a limit for BytesQueue size in MB."))
+	configurationFields        = []field.SchemaField{Vault, ServiceAccountClientId, ServiceAccountClientSecret, OrganizationId, CacheDisabled, CacheTTL, CacheMaxSize}
 )
 
 func main() {
 	ctx := context.Background()
-	_, cmd, err := configSchema.DefineConfiguration(ctx, connectorName, getConnector, field.NewConfiguration(configurationFields))
+	_, cmd, err := configSchema.DefineConfiguration(ctx,
+		connectorName,
+		getConnector,
+		field.NewConfiguration(configurationFields),
+	)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
@@ -50,12 +57,7 @@
 
 func getConnector(ctx context.Context, cfg *viper.Viper) (types.ConnectorServer, error) {
 	l := ctxzap.Extract(ctx)
-	cb, err := connector.New(ctx,
-		cfg.GetString(serviceAccountClientId),
-		cfg.GetString(serviceAccountClientSecret),
-		cfg.GetString(organizationId),
-		cfg.GetString(vault),
-	)
+	cb, err := connector.New(ctx, cfg)
 	if err != nil {
 		l.Error("error creating connector", zap.Error(err))
 		return nil, err

pkg/client/client.go

@@ -1,4 +1,4 @@
 package client
 
 import (
 	"bytes"
@@ -16,12 +16,68 @@
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 )
 
-type VGSClient struct {
-	httpClient      *uhttp.BaseHttpClient
-	token           *JWT
-	serviceEndpoint string
-	organizationId  string
-	vaultId         string
+type (
+	VGSClient struct {
+		httpClient      *uhttp.BaseHttpClient
+		token           *JWT
+		serviceEndpoint string
+		organizationId  string
+		vaultId         string
+	}
+
+	Config struct {
+		serviceAccountClientId     string
+		serviceAccountClientSecret string
+		organizationId             string
+		vaultId                    string
+	}
+)
+
+const (
+	ServiceAccountClientIdName     = "service-account-client-id"
+	ServiceAccountClientSecretName = "service-account-client-secret"
+	OrganizationId                 = "organization-id"
+	VaultId                        = "vault"
+	serviceAccountClient           = "serviceAccountClientId"
+	serviceAccountClientSecret     = "serviceAccountClientSecret"
+	organization                   = "organizationId"
+	vault                          = "vaultId"
+	empty                          = ""
+)
+
+func (c *Config) WithServiceAccountClientId(sAccId string) *Config {
+	c.serviceAccountClientId = sAccId
+	return c
+}
+
+func (c *Config) WithServiceAccountClientSecret(sAccSec string) *Config {
+	c.serviceAccountClientSecret = sAccSec
+	return c
+}
+
+func (c *Config) WithOrganizationId(orgId string) *Config {
+	c.organizationId = orgId
+	return c
+}
+
+func (c *Config) WithVaultId(vId string) *Config {
+	c.vaultId = vId
+	return c
+}
+
+func (c *Config) getFieldValue(fieldName string) string {
+	switch fieldName {
+	case serviceAccountClient:
+		return c.serviceAccountClientId
+	case serviceAccountClientSecret:
+		return c.serviceAccountClientSecret
+	case organization:
+		return c.organizationId
+	case vault:
+		return c.vaultId
+	}
+
+	return empty
 }
 
 func WithBody(body string) uhttp.RequestOption {
@@ -95,8 +151,14 @@
 	return uhttp.WithHeader("Authorization", "Basic "+basicAuth(username, password))
 }
 
-func New(ctx context.Context, clientId, clientSecret, orgId, vaultId string) (*VGSClient, error) {
-	var jwt = &JWT{}
+func New(ctx context.Context, cfg Config) (*VGSClient, error) {
+	var (
+		jwt          = &JWT{}
+		clientId     = cfg.getFieldValue(serviceAccountClient)
+		clientSecret = cfg.getFieldValue(serviceAccountClientSecret)
+		orgId        = cfg.getFieldValue(organization)
+		vaultId      = cfg.getFieldValue(vault)
+	)
 	uri, err := url.Parse("https://auth.verygoodsecurity.com/auth/realms/vgs/protocol/openid-connect/token")
 	if err != nil {
 		return nil, err
@@ -107,7 +169,17 @@
 		return nil, err
 	}
 
-	cli := uhttp.NewBaseHttpClient(httpClient)
+	// Setting up in-cache memory parameters, otherwise it takes default values
+	ctx = context.WithValue(ctx, uhttp.ContextKey{}, uhttp.CacheConfig{
+		LogDebug:     true,
+		CacheTTL:     int32(1000),
+		CacheMaxSize: int(1024),
+	})
+	cli, err := uhttp.NewBaseHttpClient(ctx, httpClient)
+	if err != nil {
+		return nil, err
+	}
+
 	req, err := cli.NewRequest(ctx,
 		http.MethodPost,
 		uri,

pkg/client/internal_integration_test.go

@@ -18,6 +18,12 @@ var (
 	clientSecret, _ = os.LookupEnv("BATON_SERVICE_ACCOUNT_CLIENT_SECRET")
 	vaultId, _      = os.LookupEnv("BATON_VAULT")
 	orgId, _        = os.LookupEnv("BATON_ORGANIZATION_ID")
+	cfg             = Config{
+		serviceAccountClientId:     clientId,
+		serviceAccountClientSecret: clientSecret,
+		organizationId:             orgId,
+		vaultId:                    vaultId,
+	}
 )
 
 const (
@@ -48,7 +54,13 @@ func TestOrganizationResources(t *testing.T) {
 		},
 	}
 
-	cli, err := getClientForTesting(ctx, clientId, clientSecret, orgId, vaultId)
+	cfg := Config{
+		serviceAccountClientId:     clientId,
+		serviceAccountClientSecret: clientSecret,
+		organizationId:             orgId,
+		vaultId:                    vaultId,
+	}
+	cli, err := getClientForTesting(ctx, cfg)
 	assert.Nil(t, err)
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
@@ -81,7 +93,7 @@ func TestVaultMembers(t *testing.T) {
 		t.Skip()
 	}
 
-	cli, err := getClientForTesting(ctx, clientId, clientSecret, orgId, vaultId)
+	cli, err := getClientForTesting(ctx, cfg)
 	assert.Nil(t, err)
 
 	endpointUrl, err := url.JoinPath(baseUrl, "vaults", vaultId, "members")
@@ -104,10 +116,23 @@ func TestVaultMembers(t *testing.T) {
 	var data any
 	err = json.Unmarshal(res, &data)
 	assert.Nil(t, err)
+
+	// -- force cache response --
+	resp1, err := cli.httpClient.Do(req)
+	assert.Nil(t, err)
+
+	defer resp1.Body.Close()
+	res1, err := io.ReadAll(resp1.Body)
+	assert.Nil(t, err)
+	assert.NotNil(t, res1)
+
+	var data1 any
+	err = json.Unmarshal(res1, &data1)
+	assert.Nil(t, err)
 }
 
-func getClientForTesting(ctx context.Context, clientId, clientSecret, orgId, vaultId string) (*VGSClient, error) {
-	cli, err := New(ctx, clientId, clientSecret, orgId, vaultId)
+func getClientForTesting(ctx context.Context, cfg Config) (*VGSClient, error) {
+	cli, err := New(ctx, cfg)
 	return cli, err
 }
 
@@ -116,7 +141,7 @@ func TestVaults(t *testing.T) {
 		t.Skip()
 	}
 
-	cli, err := getClientForTesting(ctx, clientId, clientSecret, orgId, vaultId)
+	cli, err := getClientForTesting(ctx, cfg)
 	assert.Nil(t, err)
 
 	endpointUrl, err := url.JoinPath(baseUrl, "vaults")

pkg/connector/connector.go

@@ -8,11 +8,14 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/connectorbuilder"
 	"github.com/conductorone/baton-vgs/pkg/client"
+	"github.com/spf13/viper"
 )
 
-type Connector struct {
-	client *client.VGSClient
-}
+type (
+	Connector struct {
+		client *client.VGSClient
+	}
+)
 
 // ResourceSyncers returns a ResourceSyncer for each resource type that should be synced from the upstream service.
 func (d *Connector) ResourceSyncers(ctx context.Context) []connectorbuilder.ResourceSyncer {
@@ -44,13 +47,21 @@ func (d *Connector) Validate(ctx context.Context) (annotations.Annotations, erro
 }
 
 // New returns a new instance of the connector.
-func New(ctx context.Context, clientId, clientSecret, organizationId, vaultId string) (*Connector, error) {
+func New(ctx context.Context, cfg *viper.Viper) (*Connector, error) {
 	var (
-		vc  *client.VGSClient
-		err error
+		vc             *client.VGSClient
+		config         = client.Config{}
+		clientId       = cfg.GetString(client.ServiceAccountClientIdName)
+		clientSecret   = cfg.GetString(client.ServiceAccountClientSecretName)
+		organizationId = cfg.GetString(client.OrganizationId)
+		vaultId        = cfg.GetString(client.VaultId)
+		err            error
 	)
+
+	config.WithServiceAccountClientId(clientId).WithServiceAccountClientSecret(clientSecret)
+	config.WithOrganizationId(organizationId).WithVaultId(vaultId)
 	if clientId != "" && clientSecret != "" {
-		vc, err = client.New(ctx, clientId, clientSecret, organizationId, vaultId)
+		vc, err = client.New(ctx, config)
 		if err != nil {
 			return nil, err
 		}