Implement rate limiters for every Postman client function that uses a GET request (#3964)

* Made a second rate limiter for the general rate limit and moved limiters to the functions that make the GET request

* Updated tests

Author: casey-tran

pkg/sources/postman/postman.go

@@ -197,7 +197,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
 			continue
 		}
 
-		collection, err := s.client.GetCollection(collectionID)
+		collection, err := s.client.GetCollection(ctx, collectionID)
 		if err != nil {
 			return fmt.Errorf("error getting collection %s: %w", collectionID, err)
 		}
@@ -260,7 +260,7 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
 
 	// gather and scan environment variables
 	for _, envID := range workspace.Environments {
-		envVars, err := s.client.GetEnvironmentVariables(envID.UUID)
+		envVars, err := s.client.GetEnvironmentVariables(ctx, envID.UUID)
 		if err != nil {
 			ctx.Logger().Error(err, "could not get env variables", "environment_uuid", envID.UUID)
 			continue
@@ -297,7 +297,7 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
 		if shouldSkip(collectionID.UUID, s.conn.IncludeCollections, s.conn.ExcludeCollections) {
 			continue
 		}
-		collection, err := s.client.GetCollection(collectionID.UUID)
+		collection, err := s.client.GetCollection(ctx, collectionID.UUID)
 		if err != nil {
 			return err
 		}

pkg/sources/postman/postman_client.go

@@ -187,11 +187,12 @@ type Client struct {
 	// Headers to attach to every requests made with the client.
 	Headers map[string]string
 
-	// Rate limiter needed for Postman API workspace enumeration. Postman API general rate limit 300 requests per minute
-	// but is 10 calls in 10 seconds for GET /collections, GET /workspaces, and GET /workspaces/{id} endpoints.
-	WorkspaceEnumerationRateLimiter *rate.Limiter
+	// Rate limiter needed for Postman API workspace and collection requests. Postman API rate limit
+	// is 10 calls in 10 seconds for GET /collections, GET /workspaces, and GET /workspaces/{id} endpoints.
+	WorkspaceAndCollectionRateLimiter *rate.Limiter
 
-	// TODO: Add rate limiter(s) for Postman API calls inside of other loops
+	// Rate limiter needed for Postman API. General rate limit is 300 requests per minute.
+	GeneralRateLimiter *rate.Limiter
 }
 
 // NewClient returns a new Postman API client.
@@ -203,9 +204,10 @@ func NewClient(postmanToken string) *Client {
 	}
 
 	c := &Client{
-		HTTPClient:                      http.DefaultClient,
-		Headers:                         bh,
-		WorkspaceEnumerationRateLimiter: rate.NewLimiter(rate.Every(time.Second), 1),
+		HTTPClient:                        http.DefaultClient,
+		Headers:                           bh,
+		WorkspaceAndCollectionRateLimiter: rate.NewLimiter(rate.Every(time.Second), 1),
+		GeneralRateLimiter:                rate.NewLimiter(rate.Every(time.Second/5), 1),
 	}
 
 	return c
@@ -281,9 +283,6 @@ func (c *Client) EnumerateWorkspaces(ctx context.Context) ([]Workspace, error) {
 	}
 
 	for i, workspace := range workspacesObj.Workspaces {
-		if err := c.WorkspaceEnumerationRateLimiter.Wait(ctx); err != nil {
-			return nil, fmt.Errorf("could not wait for rate limiter during workspace enumeration: %w", err)
-		}
 		tempWorkspace, err := c.GetWorkspace(ctx, workspace.ID)
 		if err != nil {
 			return nil, fmt.Errorf("could not get workspace %q (%s) during enumeration: %w", workspace.Name, workspace.ID, err)
@@ -304,6 +303,9 @@ func (c *Client) GetWorkspace(ctx context.Context, workspaceUUID string) (Worksp
 	}{}
 
 	url := fmt.Sprintf(WORKSPACE_URL, workspaceUUID)
+	if err := c.WorkspaceAndCollectionRateLimiter.Wait(ctx); err != nil {
+		return Workspace{}, fmt.Errorf("could not wait for rate limiter during workspace getting: %w", err)
+	}
 	r, err := c.getPostmanReq(url, nil)
 	if err != nil {
 		return Workspace{}, fmt.Errorf("could not get workspace (%s): %w", workspaceUUID, err)
@@ -323,12 +325,15 @@ func (c *Client) GetWorkspace(ctx context.Context, workspaceUUID string) (Worksp
 }
 
 // GetEnvironmentVariables returns the environment variables for a given environment
-func (c *Client) GetEnvironmentVariables(environment_uuid string) (VariableData, error) {
+func (c *Client) GetEnvironmentVariables(ctx context.Context, environment_uuid string) (VariableData, error) {
 	obj := struct {
 		VariableData VariableData `json:"environment"`
 	}{}
 
 	url := fmt.Sprintf(ENVIRONMENTS_URL, environment_uuid)
+	if err := c.GeneralRateLimiter.Wait(ctx); err != nil {
+		return VariableData{}, fmt.Errorf("could not wait for rate limiter during environment variable getting: %w", err)
+	}
 	r, err := c.getPostmanReq(url, nil)
 	if err != nil {
 		return VariableData{}, fmt.Errorf("could not get env variables for environment (%s): %w", environment_uuid, err)
@@ -347,12 +352,15 @@ func (c *Client) GetEnvironmentVariables(environment_uuid string) (VariableData,
 }
 
 // GetCollection returns the collection for a given collection
-func (c *Client) GetCollection(collection_uuid string) (Collection, error) {
+func (c *Client) GetCollection(ctx context.Context, collection_uuid string) (Collection, error) {
 	obj := struct {
 		Collection Collection `json:"collection"`
 	}{}
 
 	url := fmt.Sprintf(COLLECTIONS_URL, collection_uuid)
+	if err := c.WorkspaceAndCollectionRateLimiter.Wait(ctx); err != nil {
+		return Collection{}, fmt.Errorf("could not wait for rate limiter during collection getting: %w", err)
+	}
 	r, err := c.getPostmanReq(url, nil)
 	if err != nil {
 		return Collection{}, fmt.Errorf("could not get collection (%s): %w", collection_uuid, err)

pkg/sources/postman/postman_test.go

@@ -242,7 +242,7 @@ func TestSource_ScanVariableData(t *testing.T) {
 	}
 }
 
-func TestSource_ScanEnumerateRateLimit(t *testing.T) {
+func TestSource_ScanEnumerateWorkspaceRateLimit(t *testing.T) {
 	defer gock.Off()
 	// Mock the API response for workspaces
 	numWorkspaces := 3
@@ -265,9 +265,9 @@ func TestSource_ScanEnumerateRateLimit(t *testing.T) {
 			Get(fmt.Sprintf("/workspaces/%d", i)).
 			Reply(200).
 			BodyString(fmt.Sprintf(`{"workspace":{"id":"%d","name":"workspace-%d","type":"personal","description":"Test workspace number %d",
-		"visibility":"personal","createdBy":"1234","updatedBy":"1234","createdAt":"2024-12-12T23:32:27.000Z","updatedAt":"2024-12-12T23:33:01.000Z",
-		"collections":[{"id":"abc%d","name":"test-collection-1","uid":"1234-abc%d"},{"id":"def%d","name":"test-collection-2","uid":"1234-def%d"}],
-		"environments":[{"id":"ghi%d","name":"test-environment-1","uid":"1234-ghi%d"},{"id":"jkl%d","name":"test-environment-2","uid":"1234-jkl%d"}]}}`, i, i, i, i, i, i, i, i, i, i, i))
+			"visibility":"personal","createdBy":"1234","updatedBy":"1234","createdAt":"2024-12-12T23:32:27.000Z","updatedAt":"2024-12-12T23:33:01.000Z",
+			"collections":[{"id":"abc%d","name":"test-collection-1","uid":"1234-abc%d"},{"id":"def%d","name":"test-collection-2","uid":"1234-def%d"}],
+			"environments":[{"id":"ghi%d","name":"test-environment-1","uid":"1234-ghi%d"},{"id":"jkl%d","name":"test-environment-2","uid":"1234-jkl%d"}]}}`, i, i, i, i, i, i, i, i, i, i, i))
 	}
 
 	ctx := context.Background()
@@ -281,6 +281,7 @@ func TestSource_ScanEnumerateRateLimit(t *testing.T) {
 		t.Fatalf("init error: %v", err)
 	}
 	gock.InterceptClient(s.client.HTTPClient)
+	defer gock.RestoreClient(s.client.HTTPClient)
 
 	start := time.Now()
 	_, err = s.client.EnumerateWorkspaces(ctx)
@@ -291,6 +292,44 @@ func TestSource_ScanEnumerateRateLimit(t *testing.T) {
 	// With <numWorkspaces> requests at 1 per second rate limit,
 	// elapsed time should be at least <numWorkspaces - 1> seconds
 	if elapsed < time.Duration(numWorkspaces-1)*time.Second {
-		t.Errorf("Rate limiting not working as expected. Elapsed time: %v, expected at least %d seconds", elapsed, numWorkspaces-1)
+		t.Errorf("Rate limiting not working as expected. Elapsed time: %v seconds, expected at least %d seconds", elapsed.Seconds(), numWorkspaces-1)
+	}
+}
+
+func TestSource_ScanGeneralRateLimit(t *testing.T) {
+	defer gock.Off()
+	// Mock the API response for a specific environment id
+	gock.New("https://api.getpostman.com").
+		Get("environments/abc").
+		Persist().
+		Reply(200).
+		BodyString(`{"environment":{"uid":"1234-abc","id":"abc","name":"test-environment","owner":"1234","createdAt":"2025-02-13T23:17:36.000Z",
+		"updatedAt":"2025-02-13T23:18:14.000Z","values":[{"key":"test-key","value":"a-secret-value","enabled":true,"type":"default"}],"isPublic":false}}`)
+	ctx := context.Background()
+	s, conn := createTestSource(&sourcespb.Postman{
+		Credential: &sourcespb.Postman_Token{
+			Token: "super-secret-token",
+		},
+	})
+	err := s.Init(ctx, "test - postman", 0, 1, false, conn, 1)
+	if err != nil {
+		t.Fatalf("init error: %v", err)
+	}
+	gock.InterceptClient(s.client.HTTPClient)
+	defer gock.RestoreClient(s.client.HTTPClient)
+
+	numRequests := 3
+	start := time.Now()
+	for i := 0; i < numRequests; i++ {
+		_, err = s.client.GetEnvironmentVariables(ctx, "abc")
+		if err != nil {
+			t.Fatalf("get environment variables error: %v", err)
+		}
+	}
+	elapsed := time.Since(start)
+	// With number of requests at 5 per second rate limit,
+	// elapsed time should be at least <(numRequests - 1)/5> seconds
+	if elapsed < time.Duration((numRequests-1)/5)*time.Second {
+		t.Errorf("Rate limiting not working as expected. Elapsed time: %v seconds, expected at least %v seconds", elapsed.Seconds(), (float64(numRequests)-1)/5)
 	}
 }