feat(okta): Add rate limiter config (#10880)

candiduslynx · web-flow · commit 293d9d0cb7cd · 2023-05-22T12:38:56.000Z
Closes #10570
diff --git a/plugins/source/okta/client/client.go b/plugins/source/okta/client/client.go
@@ -2,16 +2,14 @@ package client
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"os"
+	"time"
 
 	"github.com/cloudquery/plugin-pb-go/specs"
 	"github.com/cloudquery/plugin-sdk/v2/plugins/source"
 	"github.com/cloudquery/plugin-sdk/v2/schema"
 	"github.com/okta/okta-sdk-golang/v3/okta"
 	"github.com/rs/zerolog"
-	"github.com/thoas/go-funk"
 )
 
 type Client struct {
@@ -23,8 +21,6 @@ type Client struct {
 	*okta.APIClient
 }
 
-const exampleDomain = "https://<CHANGE_THIS_TO_YOUR_OKTA_DOMAIN>.okta.com"
-
 func (c *Client) Logger() *zerolog.Logger {
 	return &c.logger
 }
@@ -42,48 +38,25 @@ func New(logger zerolog.Logger, s specs.Source, okt *okta.APIClient) *Client {
 	}
 }
 
-func Configure(_ context.Context, logger zerolog.Logger, s specs.Source, _ source.Options) (schema.ClientMeta, error) {
-	oktaSpec := &Spec{}
-	if err := s.UnmarshalSpec(oktaSpec); err != nil {
+func Configure(_ context.Context, logger zerolog.Logger, srcSpec specs.Source, _ source.Options) (schema.ClientMeta, error) {
+	spec := &Spec{}
+	if err := srcSpec.UnmarshalSpec(spec); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal okta spec: %w", err)
 	}
 
-	oktaToken, ok := os.LookupEnv("OKTA_API_TOKEN")
-	if !ok {
-		if oktaSpec.Token == "" {
-			return nil, errors.New("missing OKTA_API_TOKEN, either set it as an environment variable or pass it in the configuration")
-		}
-
-		oktaToken = oktaSpec.Token
-	}
-
-	if oktaSpec.Domain == "" || oktaSpec.Domain == exampleDomain {
-		return nil, errors.New(`failed to configure provider, please set your okta "domain" in okta.yml`)
+	spec.setDefaults()
+	if err := spec.validate(); err != nil {
+		return nil, err
 	}
 
 	cf := okta.NewConfiguration(
-		okta.WithOrgUrl(oktaSpec.Domain),
-		okta.WithToken(oktaToken),
+		okta.WithOrgUrl(spec.Domain),
+		okta.WithToken(spec.Token),
 		okta.WithCache(true),
+		okta.WithRateLimitMaxBackOff(int64(spec.RateLimit.MaxBackoff/time.Second)), // this param takes int64 of seconds
+		okta.WithRateLimitMaxRetries(spec.RateLimit.MaxRetries),
 	)
 	c := okta.NewAPIClient(cf)
 
-	return New(logger, s, c), nil
-}
-
-func ResolveNullableTime(path string) schema.ColumnResolver {
-	return func(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource, c schema.Column) error {
-		data := funk.Get(resource.Item, path)
-		if data == nil {
-			return nil
-		}
-		ts, ok := data.(okta.NullableTime)
-		if !ok {
-			return fmt.Errorf("unexpected type, want \"okta.NullableTime\", have \"%T\"", data)
-		}
-		if !ts.IsSet() {
-			return resource.Set(c.Name, nil)
-		}
-		return resource.Set(c.Name, ts.Get())
-	}
+	return New(logger, srcSpec, c), nil
 }
diff --git a/plugins/source/okta/client/spec.go b/plugins/source/okta/client/spec.go
@@ -1,6 +1,60 @@
 package client
 
-type Spec struct {
-	Token  string `json:"token,omitempty"`
-	Domain string `json:"domain"`
+import (
+	"fmt"
+	"os"
+	"time"
+)
+
+type (
+	Spec struct {
+		Token     string     `json:"token,omitempty"`
+		Domain    string     `json:"domain,omitempty"`
+		RateLimit *RateLimit `json:"rate_limit,omitempty"`
+	}
+	RateLimit struct {
+		MaxBackoff time.Duration `json:"max_backoff,omitempty"`
+		MaxRetries int32         `json:"max_retries,omitempty"`
+	}
+)
+
+const (
+	OktaAPIToken = "OKTA_API_TOKEN"
+)
+
+func (s *Spec) setDefaults() {
+	const (
+		minRetries = int32(3)
+		minBackOff = 5 * time.Second
+	)
+
+	if s.RateLimit == nil {
+		s.RateLimit = new(RateLimit)
+	}
+
+	if s.RateLimit.MaxRetries < minRetries {
+		s.RateLimit.MaxRetries = minRetries
+	}
+
+	if s.RateLimit.MaxBackoff < minBackOff {
+		s.RateLimit.MaxBackoff = minBackOff
+	}
+
+	if len(s.Token) == 0 {
+		s.Token = os.Getenv(OktaAPIToken)
+	}
+}
+
+func (s *Spec) validate() error {
+	if len(s.Token) == 0 {
+		return fmt.Errorf("missing API token (should be set in the configuration or as %q environment variable)", OktaAPIToken)
+	}
+
+	const exampleDomain = "https://<CHANGE_THIS_TO_YOUR_OKTA_DOMAIN>.okta.com"
+	switch s.Domain {
+	case "", exampleDomain:
+		return fmt.Errorf("missing \"domain\" in plugin configuration")
+	}
+
+	return nil
 }
diff --git a/plugins/source/okta/client/transformers.go b/plugins/source/okta/client/transformers.go
@@ -1,11 +1,14 @@
 package client
 
 import (
+	"context"
+	"fmt"
 	"reflect"
 
 	"github.com/cloudquery/plugin-sdk/v2/schema"
 	"github.com/cloudquery/plugin-sdk/v2/transformers"
 	"github.com/okta/okta-sdk-golang/v3/okta"
+	"github.com/thoas/go-funk"
 )
 
 var options = []transformers.StructTransformerOption{
@@ -27,8 +30,25 @@ func typeTransformer(field reflect.StructField) (schema.ValueType, error) {
 
 func resolverTransformer(field reflect.StructField, path string) schema.ColumnResolver {
 	if field.Type == reflect.TypeOf(okta.NullableTime{}) {
-		return ResolveNullableTime(path)
+		return resolveNullableTime(path)
 	}
 
 	return transformers.DefaultResolverTransformer(field, path)
 }
+
+func resolveNullableTime(path string) schema.ColumnResolver {
+	return func(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource, c schema.Column) error {
+		data := funk.Get(resource.Item, path)
+		if data == nil {
+			return nil
+		}
+		ts, ok := data.(okta.NullableTime)
+		if !ok {
+			return fmt.Errorf("unexpected type, want \"okta.NullableTime\", have \"%T\"", data)
+		}
+		if !ts.IsSet() {
+			return resource.Set(c.Name, nil)
+		}
+		return resource.Set(c.Name, ts.Get())
+	}
+}
diff --git a/website/pages/docs/plugins/sources/okta/overview.mdx b/website/pages/docs/plugins/sources/okta/overview.mdx
@@ -22,8 +22,31 @@ The following example sets up the Okta plugin, and connects it to a postgresql d
 
 <Configuration/>
 
-- `domain` (Required) - Specify the Okta domain you are fetching from. [Visit this link](https://developer.okta.com/docs/guides/find-your-domain/findorg/) to find your Okta domain
-- `token` (Optional) - Okta Token to access the API. You can set this with an `OKTA_API_TOKEN` environment variable
+- `domain` (`string`, required)
+
+  Specify the Okta domain you are fetching from.
+  [Visit this link](https://developer.okta.com/docs/guides/find-your-domain/findorg/) to find your Okta domain.
+
+- `token` (`string`, optional)
+
+  Token for Okta API access.
+  You can set this with an `OKTA_API_TOKEN` environment variable.
+
+- `rate_limit` ([Rate limit](#rate-limit-spec) spec, optional. Default: see [rate limit](#rate-limit-spec) spec defaults)
+
+  Rate limit configuration.
+
+### Rate limit spec
+
+- `max_backoff` (`duration`, optional. Default: `5s`)
+
+  Max backoff interval to be used.
+  If the value specified is less than the default one, the default one is used.
+
+- `max_retries` (`int32`, optional. Default: `3`)
+
+  Max retries to be performed.
+  If the value specified is less than the default one, the default one is used.
 
 ## Example Queries
 
