feat: dump custom entities in config fetcher (#6305)

* dump custom entities in config fetcher

* add integration test case for last valid config

* update changelog

* use custom.Type in test util function buildCustomEntityWithObject

* Apply suggestions from code review

Co-authored-by: Jakub Warczarek <jakub.warczarek@konghq.com>

* fix CustomEntityTypes

---------

Co-authored-by: Jakub Warczarek <jakub.warczarek@konghq.com>

CHANGELOG.md

@@ -116,6 +116,9 @@ Adding a new version? You'll need three changes:
 - Added `duration` field in logs after successfully sent configuration to Kong
   gateway or Konnect.
   [#6360](https://github.com/Kong/kubernetes-ingress-controller/pull/6360)
+- `KongCustomEntity` is now included in last valid configuration retrieved from
+  Kong gateways.
+  [#6305](https://github.com/Kong/kubernetes-ingress-controller/pull/6305)
 
 ### Fixed
 

internal/dataplane/configfetcher/config_fetcher.go

@@ -19,7 +19,7 @@ import (
 
 type LastValidConfigFetcher interface {
 	// TryFetchingValidConfigFromGateways tries to fetch a valid configuration from all gateways and persists it if found.
-	TryFetchingValidConfigFromGateways(ctx context.Context, logger logr.Logger, gatewayClients []*adminapi.Client) error
+	TryFetchingValidConfigFromGateways(ctx context.Context, logger logr.Logger, gatewayClients []*adminapi.Client, customEntityTypes []string) error
 
 	// LastValidConfig returns the last valid config and true if there's one available. Otherwise, second return value is false.
 	LastValidConfig() (*kongstate.KongState, bool)
@@ -108,6 +108,7 @@ func (cf *DefaultKongLastGoodConfigFetcher) TryFetchingValidConfigFromGateways(
 	ctx context.Context,
 	logger logr.Logger,
 	gatewayClients []*adminapi.Client,
+	customEntityTypes []string,
 ) error {
 	logger.V(logging.DebugLevel).Info("Fetching last good configuration from gateway clients", "count", len(gatewayClients))
 
@@ -118,7 +119,10 @@ func (cf *DefaultKongLastGoodConfigFetcher) TryFetchingValidConfigFromGateways(
 	)
 	for _, client := range gatewayClients {
 		logger.V(logging.DebugLevel).Info("Fetching configuration", "url", client.BaseRootURL())
-		rs, err := cf.getKongRawState(ctx, client.AdminAPIClient())
+		// Copy the dump configuration and add custom entity types to fetch config with custom entities.
+		config := cf.config
+		config.CustomEntityTypes = customEntityTypes
+		rs, err := cf.getKongRawState(ctx, client.AdminAPIClient(), config)
 		if err != nil {
 			errs = errors.Join(errs, err)
 		}
@@ -152,8 +156,8 @@ func (cf *DefaultKongLastGoodConfigFetcher) TryFetchingValidConfigFromGateways(
 	return errs
 }
 
-func (cf *DefaultKongLastGoodConfigFetcher) getKongRawState(ctx context.Context, client *kong.Client) (*utils.KongRawState, error) {
-	return dump.Get(ctx, client, cf.config)
+func (cf *DefaultKongLastGoodConfigFetcher) getKongRawState(ctx context.Context, client *kong.Client, config dump.Config) (*utils.KongRawState, error) {
+	return dump.Get(ctx, client, config)
 }
 
 func (cf *DefaultKongLastGoodConfigFetcher) getKongStatus(ctx context.Context, client *kong.Client) (*kong.Status, error) {

internal/dataplane/configfetcher/config_fetcher_test.go

@@ -82,7 +82,7 @@ func TestTryFetchingValidConfigFromGateways(t *testing.T) {
 			require.Nil(t, state)
 
 			ctx := context.Background()
-			err := fetcher.TryFetchingValidConfigFromGateways(ctx, zapr.NewLogger(zap.NewNop()), tc.adminAPIClients(t, ctx))
+			err := fetcher.TryFetchingValidConfigFromGateways(ctx, zapr.NewLogger(zap.NewNop()), tc.adminAPIClients(t, ctx), nil)
 			if tc.expectError {
 				require.Error(t, err)
 				assert.False(t, ok)

internal/dataplane/configfetcher/kongrawstate.go

@@ -170,6 +170,15 @@ func KongRawStateToKongState(rawstate *utils.KongRawState) *kongstate.KongState
 		}
 	}
 
+	for _, entity := range rawstate.CustomEntities {
+		entityType := entity.Type()
+		obj := entity.Object()
+		ksEntity := kongstate.CustomEntity{
+			Object: obj,
+		}
+		kongState.AddCustomEntity(string(entityType), kongstate.EntitySchema{}, ksEntity)
+	}
+
 	return kongState
 }
 

internal/dataplane/configfetcher/kongrawstate_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/kong/go-database-reconciler/pkg/utils"
 	"github.com/kong/go-kong/kong"
+	"github.com/kong/go-kong/kong/custom"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -15,6 +16,12 @@ import (
 	"github.com/kong/kubernetes-ingress-controller/v3/internal/dataplane/kongstate"
 )
 
+func buildCustomEntityWithObject(entityType custom.Type, obj custom.Object) custom.Entity {
+	e := custom.NewEntityObject(entityType)
+	e.SetObject(obj)
+	return e
+}
+
 func TestKongRawStateToKongState(t *testing.T) {
 	// This is to gather all the fields in KongRawState that are tested in this suite.
 	testedKongRawStateFields := sets.New[string]()
@@ -170,6 +177,16 @@ func TestKongRawStateToKongState(t *testing.T) {
 						SubjectName: kong.String("subjectName"),
 					},
 				},
+				CustomEntities: []custom.Entity{
+					buildCustomEntityWithObject("degraphql_routes", custom.Object{
+						"id":    "degraphql-route-1",
+						"uri":   "/graphql",
+						"query": "query{name}",
+						"service": map[string]any{
+							"id": "service",
+						},
+					}),
+				},
 			},
 			expectedKongState: &kongstate.KongState{
 				Services: []kongstate.Service{
@@ -292,6 +309,22 @@ func TestKongRawStateToKongState(t *testing.T) {
 						},
 					},
 				},
+				CustomEntities: map[string]*kongstate.KongCustomEntityCollection{
+					"degraphql_routes": {
+						Entities: []kongstate.CustomEntity{
+							{
+								Object: custom.Object{
+									"id":    "degraphql-route-1",
+									"uri":   "/graphql",
+									"query": "query{name}",
+									"service": map[string]any{
+										"id": "service",
+									},
+								},
+							},
+						},
+					},
+				},
 			},
 		},
 		{
@@ -344,9 +377,6 @@ func ensureAllKongStateFieldsAreTested(t *testing.T, testedFields []string) {
 		"Plugins",
 		// Licenses are injected from the license getter rather than extracted from the last state.
 		"Licenses",
-		// CustomEntities are not supported yet because go-database-reconciler does not include custom entities.
-		// TODO: support custom entities: https://github.com/Kong/kubernetes-ingress-controller/issues/6054
-		"CustomEntities",
 	}
 	allKongStateFields := func() []string {
 		var fields []string

internal/dataplane/kong_client.go

@@ -68,6 +68,7 @@ const (
 type KongConfigBuilder interface {
 	BuildKongConfig() translator.KongConfigBuildingResult
 	UpdateCache(store.CacheStores)
+	CustomEntityTypes() []string
 }
 
 // FallbackConfigGenerator generates a fallback configuration based on a cache snapshot and a set of broken objects.
@@ -427,7 +428,7 @@ func (c *KongClient) Update(ctx context.Context) error {
 		// configuration already stored in memory. This can happen when KIC restarts and there
 		// already is a Kong Proxy with a valid configuration loaded.
 		if _, found := c.kongConfigFetcher.LastValidConfig(); !found {
-			if err := c.kongConfigFetcher.TryFetchingValidConfigFromGateways(ctx, c.logger, c.clientsProvider.GatewayClients()); err != nil {
+			if err := c.kongConfigFetcher.TryFetchingValidConfigFromGateways(ctx, c.logger, c.clientsProvider.GatewayClients(), c.kongConfigBuilder.CustomEntityTypes()); err != nil {
 				// If the client fails to fetch the last good configuration, we log it
 				// and carry on, as this is a condition that can be recovered with the following steps.
 				c.logger.Error(err, "Failed to fetch last good configuration from gateways")

internal/dataplane/kong_client_test.go

@@ -551,6 +551,10 @@ func (p *mockKongConfigBuilder) returnTranslationFailures(enabled bool) {
 	}
 }
 
+func (p *mockKongConfigBuilder) CustomEntityTypes() []string {
+	return nil
+}
+
 func (p *mockKongConfigBuilder) returnTranslationFailuresForAllButFirstCall(failures []failures.ResourceFailure) {
 	p.onlyFirstBuildCallWithNoTranslationFailures = true
 	p.translationFailuresToReturn = failures
@@ -1020,7 +1024,7 @@ func (cf *mockKongLastValidConfigFetcher) StoreLastValidConfig(s *kongstate.Kong
 	cf.lastKongState = s
 }
 
-func (cf *mockKongLastValidConfigFetcher) TryFetchingValidConfigFromGateways(context.Context, logr.Logger, []*adminapi.Client) error {
+func (cf *mockKongLastValidConfigFetcher) TryFetchingValidConfigFromGateways(context.Context, logr.Logger, []*adminapi.Client, []string) error {
 	if cf.kongRawState != nil {
 		cf.lastKongState = configfetcher.KongRawStateToKongState(cf.kongRawState)
 	}

internal/dataplane/kongstate/customentity.go

@@ -208,15 +208,18 @@ func (ks *KongState) FillCustomEntities(
 			continue
 		}
 		for _, generatedEntity := range generatedEntities {
-			ks.addCustomEntity(entity.Spec.EntityType, schema, generatedEntity)
+			ks.AddCustomEntity(entity.Spec.EntityType, schema, generatedEntity)
 		}
 	}
 
 	ks.sortCustomEntities()
 }
 
-// addCustomEntity adds a custom entity into the collection of its type.
-func (ks *KongState) addCustomEntity(entityType string, schema EntitySchema, e CustomEntity) {
+// AddCustomEntity adds a custom entity into the collection of its type.
+func (ks *KongState) AddCustomEntity(entityType string, schema EntitySchema, e CustomEntity) {
+	if ks.CustomEntities == nil {
+		ks.CustomEntities = map[string]*KongCustomEntityCollection{}
+	}
 	// Put the entity into the custom collection to store the entities of its type.
 	if _, ok := ks.CustomEntities[entityType]; !ok {
 		ks.CustomEntities[entityType] = &KongCustomEntityCollection{
@@ -227,6 +230,11 @@ func (ks *KongState) addCustomEntity(entityType string, schema EntitySchema, e C
 	collection.Entities = append(collection.Entities, e)
 }
 
+// CustomEntityTypes returns types of translated custom entities included in the KongState.
+func (ks *KongState) CustomEntityTypes() []string {
+	return lo.Keys(ks.CustomEntities)
+}
+
 // fetchEntitySchema fetches schema of an entity by its type and stores the schema in its custom entity collection
 // as a cache to avoid excessive calling of Kong admin APIs.
 func (ks *KongState) fetchEntitySchema(schemaGetter SchemaGetter, entityType string) (EntitySchema, error) {

internal/dataplane/translator/translator.go

@@ -88,6 +88,7 @@ type Translator struct {
 
 	// schemaServiceProvider provides the schema service required for fetching schemas of custom entities.
 	schemaServiceProvider SchemaServiceProvider
+	customEntityTypes     []string
 
 	failuresCollector          *failures.ResourceFailuresCollector
 	translatedObjectsCollector *ObjectsCollector
@@ -215,6 +216,10 @@ func (t *Translator) BuildKongConfig() KongConfigBuildingResult {
 				t.registerSuccessfullyTranslatedObject(collection.Entities[i].K8sKongCustomEntity)
 			}
 		}
+		// Update types of translated custom entities in the round of translation
+		// for dumping them from Kong gateway in config fetcher,
+		// because running full build of Kong configuration to get KongState is a heavy operation.
+		t.customEntityTypes = result.CustomEntityTypes()
 	}
 
 	// generate Certificates and SNIs
@@ -266,6 +271,13 @@ func (t *Translator) InjectLicenseGetter(licenseGetter license.Getter) {
 	t.licenseGetter = licenseGetter
 }
 
+func (t *Translator) CustomEntityTypes() []string {
+	if t.featureFlags.KongCustomEntity {
+		return t.customEntityTypes
+	}
+	return nil
+}
+
 // -----------------------------------------------------------------------------
 // Translator - Private Methods
 // -----------------------------------------------------------------------------

internal/konnect/controlplanes/client.go

@@ -534,7 +534,6 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 	queryValues := queryURL.Query()
 
 	if params.PageSize != nil {
-
 		if queryFrag, err := runtime.StyleParamWithLocation("form", true, "page[size]", runtime.ParamLocationQuery, *params.PageSize); err != nil {
 			return nil, err
 		} else if parsed, err := url.ParseQuery(queryFrag); err != nil {
@@ -546,11 +545,9 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 				}
 			}
 		}
-
 	}
 
 	if params.PageNumber != nil {
-
 		if queryFrag, err := runtime.StyleParamWithLocation("form", true, "page[number]", runtime.ParamLocationQuery, *params.PageNumber); err != nil {
 			return nil, err
 		} else if parsed, err := url.ParseQuery(queryFrag); err != nil {
@@ -562,11 +559,9 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 				}
 			}
 		}
-
 	}
 
 	if params.FilterNameEq != nil {
-
 		if queryFrag, err := runtime.StyleParamWithLocation("form", true, "filter[name][eq]", runtime.ParamLocationQuery, *params.FilterNameEq); err != nil {
 			return nil, err
 		} else if parsed, err := url.ParseQuery(queryFrag); err != nil {
@@ -578,11 +573,9 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 				}
 			}
 		}
-
 	}
 
 	if params.FilterName != nil {
-
 		if queryFrag, err := runtime.StyleParamWithLocation("form", true, "filter[name]", runtime.ParamLocationQuery, *params.FilterName); err != nil {
 			return nil, err
 		} else if parsed, err := url.ParseQuery(queryFrag); err != nil {
@@ -594,11 +587,9 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 				}
 			}
 		}
-
 	}
 
 	if params.FilterNameContains != nil {
-
 		if queryFrag, err := runtime.StyleParamWithLocation("form", true, "filter[name][contains]", runtime.ParamLocationQuery, *params.FilterNameContains); err != nil {
 			return nil, err
 		} else if parsed, err := url.ParseQuery(queryFrag); err != nil {
@@ -610,7 +601,6 @@ func NewListControlPlanesRequest(server string, params *ListControlPlanesParams)
 				}
 			}
 		}
-
 	}
 
 	queryURL.RawQuery = queryValues.Encode()