feat: add custom entities to dump.Get (#114)

Retrieve a list of custom entities and then retrieve those
entities when performing a dump.Get().

pkg/dump/dump.go

@@ -5,9 +5,11 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"sync"
 
 	"github.com/kong/go-database-reconciler/pkg/utils"
 	"github.com/kong/go-kong/kong"
+	"github.com/kong/go-kong/kong/custom"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -27,6 +29,9 @@ type Config struct {
 	// If true, licenses are exported.
 	IncludeLicenses bool
 
+	// CustomEntityTypes lists types of custom entities to list.
+	CustomEntityTypes []string
+
 	// SelectorTags can be used to export entities tagged with only specific
 	// tags.
 	SelectorTags []string
@@ -334,6 +339,45 @@ func getProxyConfiguration(ctx context.Context, group *errgroup.Group,
 			return nil
 		})
 	}
+
+	if len(config.CustomEntityTypes) > 0 {
+		// Get custom entities with types given in config.CustomEntityTypes.
+		customEntityLock := sync.Mutex{}
+		for _, entityType := range config.CustomEntityTypes {
+			t := entityType
+			group.Go(func() error {
+				// Register entity type.
+				// Because client writes an unprotected map to register entity types, we need to use mutex to protect it.
+				customEntityLock.Lock()
+				err := tryRegisterEntityType(client, custom.Type(t))
+				customEntityLock.Unlock()
+				if err != nil {
+					return fmt.Errorf("custom entity %s: %w", t, err)
+				}
+				// Fetch all entities with the given type.
+				entities, err := GetAllCustomEntitiesWithType(ctx, client, t)
+				if err != nil {
+					return fmt.Errorf("custom entity %s: %w", t, err)
+				}
+				// Add custom entities to rawstate.
+				customEntityLock.Lock()
+				state.CustomEntities = append(state.CustomEntities, entities...)
+				customEntityLock.Unlock()
+				return nil
+			})
+		}
+	}
+}
+
+func tryRegisterEntityType(client *kong.Client, typ custom.Type) error {
+	if client.Lookup(typ) != nil {
+		return nil
+	}
+	return client.Register(typ, &custom.EntityCRUDDefinition{
+		Name:       typ,
+		CRUDPath:   "/" + string(typ),
+		PrimaryKey: "id",
+	})
 }
 
 func getEnterpriseRBACConfiguration(ctx context.Context, group *errgroup.Group,
@@ -375,6 +419,7 @@ func Get(ctx context.Context, client *kong.Client, config Config) (*utils.KongRa
 	} else {
 		// regular case
 		getProxyConfiguration(ctx, group, client, config, &state)
+
 		if !config.SkipConsumers {
 			getConsumerGroupsConfiguration(ctx, group, client, config, &state)
 			getConsumerConfiguration(ctx, group, client, config, &state)
@@ -1001,6 +1046,33 @@ func GetAllLicenses(
 	return licenses, nil
 }
 
+// GetAllCustomEntitiesWithType quries Kong for all Custom entities with the given type.
+func GetAllCustomEntitiesWithType(
+	ctx context.Context, client *kong.Client, entityType string,
+) ([]custom.Entity, error) {
+	entities := []custom.Entity{}
+	opt := newOpt(nil)
+	e := custom.NewEntityObject(custom.Type(entityType))
+	for {
+		s, nextOpt, err := client.CustomEntities.List(ctx, opt, e)
+		if kong.IsNotFoundErr(err) {
+			return entities, nil
+		}
+		if err != nil {
+			return nil, err
+		}
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		entities = append(entities, s...)
+		if nextOpt == nil {
+			break
+		}
+		opt = nextOpt
+	}
+	return entities, nil
+}
+
 // excludeConsumersPlugins filter out consumer plugins
 func excludeConsumersPlugins(plugins []*kong.Plugin) []*kong.Plugin {
 	var filtered []*kong.Plugin

pkg/utils/types.go

@@ -41,7 +41,7 @@ type KongRawState struct {
 
 	Consumers      []*kong.Consumer
 	ConsumerGroups []*kong.ConsumerGroupObject
-	CustomEntities []*custom.Entity
+	CustomEntities []custom.Entity
 
 	Vaults   []*kong.Vault
 	Licenses []*kong.License

tests/integration/dump_test.go

@@ -3,8 +3,12 @@
 package integration
 
 import (
+	"context"
 	"testing"
 
+	deckDump "github.com/kong/go-database-reconciler/pkg/dump"
+	"github.com/kong/go-kong/kong"
+	"github.com/kong/go-kong/kong/custom"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -284,3 +288,48 @@ func Test_Dump_ConsumerGroupConsumersWithCustomID_Konnect(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, expected, output)
 }
+
+func Test_Dump_CustomEntities(t *testing.T) {
+	kong.RunWhenEnterprise(t, ">=3.0.0", kong.RequiredFeatures{})
+	setup(t)
+
+	require.NoError(t, sync("testdata/sync/001-create-a-service/kong3x.yaml"))
+	// Create a degraphql_route attached to the service after created a service by dedicated client
+	// because deck sync does not support custom entities.
+	const serviceID = "58076db2-28b6-423b-ba39-a797193017f7" // ID of the service in the config file
+	client, err := getTestClient()
+	require.NoError(t, err)
+	r, err := client.DegraphqlRoutes.Create(context.Background(), &kong.DegraphqlRoute{
+		Service: &kong.Service{
+			ID: kong.String(serviceID),
+		},
+		URI:   kong.String("/graphql"),
+		Query: kong.String("query{ name }"),
+	})
+	require.NoError(t, err, "Should create degraphql_routes sucessfully")
+	t.Logf("Created degraphql_routes %s attached to service %s", *r.ID, serviceID)
+
+	// Since degraphql_route does not run cascade delete on services, we need to clean it up after the test.
+	t.Cleanup(func() {
+		err := client.DegraphqlRoutes.Delete(context.Background(), kong.String(serviceID), r.ID)
+		require.NoError(t, err, "should delete degraphql_routes in cleanup")
+	})
+
+	// Call dump.Get with custom entities because deck's `dump` command does not support custom entities.
+	rawState, err := deckDump.Get(context.Background(), client, deckDump.Config{
+		CustomEntityTypes: []string{"degraphql_routes"},
+	})
+	require.NoError(t, err, "Should dump from Kong successfully")
+	require.Len(t, rawState.CustomEntities, 1, "Dumped raw state should contain 1 custom entity")
+	// check entity type
+	typ := rawState.CustomEntities[0].Type()
+	require.Equal(t, custom.Type("degraphql_routes"), typ, "Entity should have type degraphql_routes")
+	// check fields of the entity
+	obj := rawState.CustomEntities[0].Object()
+	uri, ok := obj["uri"].(string)
+	require.Truef(t, ok, "'uri' field should have type 'string' but actual '%T'", obj["uri"])
+	require.Equal(t, "/graphql", uri)
+	query, ok := obj["query"].(string)
+	require.Truef(t, ok, "'query' field should have type 'string' but actual '%T'", obj["query"])
+	require.Equal(t, "query{ name }", query)
+}