fix: paginate environment-document for local evaluation (#204)

Author: higordasneves

client.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 	"net/http"
+	"net/url"
 	"reflect"
 	"runtime"
 	"strings"
@@ -432,29 +433,54 @@ func (c *Client) pollThenStartRealtime(ctx context.Context) {
 }
 
 func (c *Client) UpdateEnvironment(ctx context.Context) error {
+	start := time.Now()
+
 	var env environments.EnvironmentModel
-	resp, err := c.client.NewRequest().
-		SetContext(ctx).
-		SetResult(&env).
-		ForceContentType("application/json").
-		Get(c.config.baseURL + "environment-document/")
+	nextPage := ""
+	pageCount := 0
 
-	if err != nil {
-		msg := fmt.Sprintf("flagsmith: error performing request to Flagsmith API: %s", err)
-		f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
-		if c.errorHandler != nil {
-			c.errorHandler(f)
+	for {
+		var page environments.EnvironmentModel
+		req := c.client.NewRequest().
+			SetContext(ctx).
+			SetResult(&page).
+			ForceContentType("application/json")
+		if nextPage != "" {
+			req = req.SetQueryParam("page_id", nextPage)
 		}
-		return f
-	}
-	if resp.StatusCode() != 200 {
-		msg := fmt.Sprintf("flagsmith: unexpected response from Flagsmith API: %s", resp.Status())
-		f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
-		if c.errorHandler != nil {
-			c.errorHandler(f)
+
+		resp, err := req.Get(c.config.baseURL + "environment-document/")
+		if err != nil {
+			msg := fmt.Sprintf("flagsmith: error performing request to Flagsmith API: %s", err)
+			f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
+			if c.errorHandler != nil {
+				c.errorHandler(f)
+			}
+			return f
+		}
+		if resp.StatusCode() != 200 {
+			msg := fmt.Sprintf("flagsmith: unexpected response from Flagsmith API: %s", resp.Status())
+			f := &FlagsmithAPIError{Msg: msg, Err: err, ResponseStatusCode: resp.StatusCode(), ResponseStatus: resp.Status()}
+			if c.errorHandler != nil {
+				c.errorHandler(f)
+			}
+			return f
+		}
+
+		pageCount++
+		nextPage = c.ExtractNextPage(resp.Header().Get("link"))
+
+		if pageCount == 1 {
+			env = page
+		} else {
+			env.IdentityOverrides = append(env.IdentityOverrides, page.IdentityOverrides...)
+		}
+
+		if nextPage == "" {
+			break
 		}
-		return f
 	}
+
 	isNew := false
 	previousEnv := c.environment.Load()
 	if previousEnv == nil || env.UpdatedAt.After(previousEnv.(*environments.EnvironmentModel).UpdatedAt) {
@@ -469,5 +495,35 @@ func (c *Client) UpdateEnvironment(ctx context.Context) error {
 		c.log.Info("environment updated", "environment", env.APIKey, "updated_at", env.UpdatedAt)
 	}
 
+	c.log.Debug("IdentityOverrides", "len", len(env.IdentityOverrides))
+
+	if elapsed := time.Since(start); c.config.envRefreshInterval > 0 && elapsed > c.config.envRefreshInterval {
+		c.log.Warn(
+			"fetching environment took longer than the configured refresh interval; raise WithEnvironmentRefreshInterval or trim the environment",
+			"elapsed", elapsed,
+			"refresh_interval", c.config.envRefreshInterval,
+		)
+	}
+
 	return nil
 }
+
+// ExtractNextPage parses the Link header from the environment-document API and
+// returns the decoded page_id value when a next page exists, or empty string otherwise.
+// Expected format: </api/v1/environment-document/?page_id=xxx>; rel="next".
+func (c *Client) ExtractNextPage(linkHeader string) string {
+	parts := strings.SplitN(linkHeader, ">", 2)
+	if len(parts) == 0 {
+		return ""
+	}
+
+	u, err := url.Parse(strings.TrimPrefix(parts[0], "<"))
+	if err != nil {
+		return ""
+	}
+
+	pageID := u.Query().Get("page_id")
+	c.log.Debug("environment-document next page", "link", linkHeader, "page_id", pageID)
+
+	return pageID
+}

client_test.go

@@ -1236,3 +1236,158 @@ func TestCustomClientOptionsShoudPanic(t *testing.T) {
 		})
 	}
 }
+
+func TestExtractNextPage(t *testing.T) {
+	client := flagsmith.NewClient("test-key")
+
+	testCases := []struct {
+		name     string
+		header   string
+		expected string
+	}{
+		{
+			name:     "valid link header with encoded page_id",
+			header:   "</api/v1/environment-document/?page_id=" + fixtures.PageIDEncoded + ">; rel=\"next\"",
+			expected: fixtures.PageID,
+		},
+		{
+			name:     "empty header returns empty string",
+			header:   "",
+			expected: "",
+		},
+		{
+			name:     "header without page_id returns empty string",
+			header:   "</api/v1/environment-document/>; rel=\"next\"",
+			expected: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := client.ExtractNextPage(tc.header)
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}
+
+func TestUpdateEnvironmentPaginatesIdentityOverrides(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.PaginatedEnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then
+	assert.NoError(t, err)
+
+	// Identity from page 1 should be found
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled, "identity from page 1 should have overridden feature disabled")
+
+	// Identity from page 2 should also be found
+	flags2, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifierPage2, nil)
+	assert.NoError(t, err)
+	enabled2, err := flags2.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled2, "identity from page 2 should have overridden feature disabled")
+}
+
+func TestUpdateEnvironmentSinglePageNoLinkHeader(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.EnvironmentDocumentHandler))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then — no pagination, identity from page 1 should still work
+	assert.NoError(t, err)
+
+	flags, err := client.GetIdentityFlags(ctx, fixtures.OverriddenIdentifier, nil)
+	assert.NoError(t, err)
+	enabled, err := flags.IsFeatureEnabled(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.False(t, enabled, "identity override should have feature disabled")
+}
+
+func TestUpdateEnvironmentLogsWarningWhenSlowerThanRefreshInterval(t *testing.T) {
+	// Given: handler delays the response so the fetch takes longer than the
+	// refresh interval; we capture logs to assert the warning is emitted.
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		time.Sleep(20 * time.Millisecond)
+		fixtures.EnvironmentDocumentHandler(rw, req)
+	}))
+	defer server.Close()
+
+	var logOutput strings.Builder
+	var logMu sync.Mutex
+	slogLogger := slog.New(slog.NewTextHandler(writerFunc(func(p []byte) (n int, err error) {
+		logMu.Lock()
+		defer logMu.Unlock()
+		return logOutput.Write(p)
+	}), &slog.HandlerOptions{Level: slog.LevelWarn}))
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithSlogLogger(slogLogger),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"),
+		flagsmith.WithEnvironmentRefreshInterval(1*time.Millisecond))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then
+	assert.NoError(t, err)
+
+	logMu.Lock()
+	logStr := logOutput.String()
+	logMu.Unlock()
+
+	assert.Contains(t, logStr, "fetching environment took longer than the configured refresh interval")
+	assert.Contains(t, logStr, "refresh_interval=1ms")
+	assert.Contains(t, logStr, "elapsed=")
+}
+
+func TestUpdateEnvironmentDoesNotLogWarningWhenWithinRefreshInterval(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.EnvironmentDocumentHandler))
+	defer server.Close()
+
+	var logOutput strings.Builder
+	var logMu sync.Mutex
+	slogLogger := slog.New(slog.NewTextHandler(writerFunc(func(p []byte) (n int, err error) {
+		logMu.Lock()
+		defer logMu.Unlock()
+		return logOutput.Write(p)
+	}), &slog.HandlerOptions{Level: slog.LevelWarn}))
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithSlogLogger(slogLogger),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"),
+		flagsmith.WithEnvironmentRefreshInterval(10*time.Second))
+
+	// When
+	err := client.UpdateEnvironment(ctx)
+
+	// Then
+	assert.NoError(t, err)
+
+	logMu.Lock()
+	logStr := logOutput.String()
+	logMu.Unlock()
+
+	assert.NotContains(t, logStr, "fetching environment took longer")
+}

fixtures/fixture.go

@@ -14,6 +14,11 @@ const Feature1ID = 1
 const Feature1OverriddenValue = "some-overridden-value"
 const ClientAPIKey = "B62qaMZNwfiqT76p38ggrQ"
 
+const OverriddenIdentifier = "overridden-id"
+const OverriddenIdentifierPage2 = "overridden-id-page2"
+const PageID = "identity_override:1:00000000-0000-0000-0000-000000000001"
+const PageIDEncoded = "identity_override%3A1%3A00000000-0000-0000-0000-000000000001"
+
 const EnvironmentJson = `
 {
 	"api_key": "B62qaMZNwfiqT76p38ggrQ",
@@ -206,6 +211,40 @@ const IdentityResponseJson = `
 
 `
 
+// EnvironmentJsonPage2 contains only identity_overrides — the base environment fields
+// are irrelevant for subsequent pages since only IdentityOverrides are merged.
+const EnvironmentJsonPage2 = `
+{
+	"api_key": "B62qaMZNwfiqT76p38ggrQ",
+	"updated_at": "2023-12-06T10:21:54.079725Z",
+	"project": {"name": "Test project", "organisation": {"feature_analytics": false, "name": "Test Org", "id": 1, "persist_trait_data": true, "stop_serving_flags": false}, "id": 1, "hide_disabled_flags": false, "segments": []},
+	"segment_overrides": [],
+	"id": 1,
+	"feature_states": [],
+	"identity_overrides": [
+		{
+			"identifier": "overridden-id-page2",
+			"identity_uuid": "1a2b3c4d-5e6f-7890-abcd-ef1234567890",
+			"created_date": "2019-08-27T14:53:45.698555Z",
+			"updated_at": "2023-07-14 16:12:00.000000",
+			"environment_api_key": "B62qaMZNwfiqT76p38ggrQ",
+			"identity_features": [
+				{
+					"id": 1,
+					"feature": {"id": 1, "name": "feature_1", "type": "STANDARD"},
+					"featurestate_uuid": "00000000-0000-0000-0000-000000000002",
+					"feature_state_value": "some-overridden-value",
+					"enabled": false,
+					"environment": 1,
+					"identity": null,
+					"feature_segment": null
+				}
+			]
+		}
+	]
+}
+`
+
 func EnvironmentDocumentHandler(rw http.ResponseWriter, req *http.Request) {
 	if req.URL.Path != "/api/v1/environment-document/" {
 		panic("Wrong path")
@@ -223,6 +262,35 @@ func EnvironmentDocumentHandler(rw http.ResponseWriter, req *http.Request) {
 	}
 }
 
+// PaginatedEnvironmentDocumentHandler serves two pages of environment document.
+// Page 1 includes a Link header pointing to page 2. Page 2 has no Link header.
+func PaginatedEnvironmentDocumentHandler(rw http.ResponseWriter, req *http.Request) {
+	if req.URL.Path != "/api/v1/environment-document/" {
+		panic("Wrong path")
+	}
+	if req.Header.Get("X-Environment-Key") != EnvironmentAPIKey {
+		panic("Wrong API key")
+	}
+
+	rw.Header().Set("Content-Type", "application/json")
+
+	if req.URL.Query().Get("page_id") == "" {
+		rw.Header().Set("link", "</api/v1/environment-document/?page_id="+PageIDEncoded+">; rel=\"next\"")
+		rw.WriteHeader(http.StatusOK)
+		_, err := io.WriteString(rw, EnvironmentJson)
+		if err != nil {
+			panic(err)
+		}
+		return
+	}
+
+	rw.WriteHeader(http.StatusOK)
+	_, err := io.WriteString(rw, EnvironmentJsonPage2)
+	if err != nil {
+		panic(err)
+	}
+}
+
 func FlagsAPIHandlerWithInternalServerError(rw http.ResponseWriter, req *http.Request) {
 	if req.URL.Path != "/api/v1/flags/" {
 		panic("Wrong path")