From c5e8e71ee3a670001bb88cbf29a8bef840adbb84 Mon Sep 17 00:00:00 2001
From: Todd Treece <360020+toddtreece@users.noreply.github.com>
Date: Wed, 11 Dec 2024 12:13:42 -0500
Subject: [PATCH] K8s: Add APIGroupPostStartHookProvider (#97801)

---
 pkg/services/apiserver/builder/common.go      |   5 +
 pkg/services/apiserver/builder/helper.go      |  27 +++++
 pkg/services/apiserver/builder/helper_test.go | 109 ++++++++++++++++++
 3 files changed, 141 insertions(+)
 create mode 100644 pkg/services/apiserver/builder/helper_test.go

diff --git a/pkg/services/apiserver/builder/common.go b/pkg/services/apiserver/builder/common.go
index 4dfff1abc2d..5761479c092 100644
--- a/pkg/services/apiserver/builder/common.go
+++ b/pkg/services/apiserver/builder/common.go
@@ -62,6 +62,11 @@ type APIGroupRouteProvider interface {
 	GetAPIRoutes() *APIRoutes
 }
 
+type APIGroupPostStartHookProvider interface {
+	// GetPostStartHooks returns a list of functions that will be called after the server has started
+	GetPostStartHooks() (map[string]genericapiserver.PostStartHookFunc, error)
+}
+
 type APIGroupOptions struct {
 	Scheme           *runtime.Scheme
 	OptsGetter       generic.RESTOptionsGetter
diff --git a/pkg/services/apiserver/builder/helper.go b/pkg/services/apiserver/builder/helper.go
index 911f76093af..88c4b2eb8ef 100644
--- a/pkg/services/apiserver/builder/helper.go
+++ b/pkg/services/apiserver/builder/helper.go
@@ -141,6 +141,10 @@ func SetupConfig(
 
 	serverConfig.EffectiveVersion = utilversion.DefaultKubeEffectiveVersion()
 
+	if err := AddPostStartHooks(serverConfig, builders); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -261,3 +265,26 @@ func InstallAPIs(
 
 	return nil
 }
+
+// AddPostStartHooks adds post start hooks to a generic API server config
+func AddPostStartHooks(
+	config *genericapiserver.RecommendedConfig,
+	builders []APIGroupBuilder,
+) error {
+	for _, b := range builders {
+		hookProvider, ok := b.(APIGroupPostStartHookProvider)
+		if !ok {
+			continue
+		}
+		hooks, err := hookProvider.GetPostStartHooks()
+		if err != nil {
+			return err
+		}
+		for name, hook := range hooks {
+			if err := config.AddPostStartHook(name, hook); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
diff --git a/pkg/services/apiserver/builder/helper_test.go b/pkg/services/apiserver/builder/helper_test.go
new file mode 100644
index 00000000000..f3414b47579
--- /dev/null
+++ b/pkg/services/apiserver/builder/helper_test.go
@@ -0,0 +1,109 @@
+package builder_test
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apiserver/pkg/authorization/authorizer"
+	"k8s.io/apiserver/pkg/server"
+	"k8s.io/kube-openapi/pkg/common"
+
+	"github.com/grafana/grafana/pkg/services/apiserver"
+	"github.com/grafana/grafana/pkg/services/apiserver/builder"
+)
+
+func TestAddPostStartHooks(t *testing.T) {
+	tests := []struct {
+		name      string
+		builders  []builder.APIGroupBuilder
+		wantErr   bool
+		wantHooks []string
+	}{
+		{
+			name:     "no builders",
+			builders: []builder.APIGroupBuilder{},
+			wantErr:  false,
+		},
+		{
+			name: "builder without post start hooks",
+			builders: []builder.APIGroupBuilder{
+				&mockAPIGroupPostStartHookProvider{},
+			},
+			wantErr: false,
+		},
+		{
+			name: "builder with post start hooks",
+			builders: []builder.APIGroupBuilder{
+				&mockAPIGroupPostStartHookProvider{
+					hooks: map[string]server.PostStartHookFunc{
+						"test-hook": func(server.PostStartHookContext) error { return nil },
+					},
+				},
+			},
+			wantErr:   false,
+			wantHooks: []string{"test-hook"},
+		},
+		{
+			name: "builder with post start hook provider error",
+			builders: []builder.APIGroupBuilder{
+				&mockAPIGroupPostStartHookProvider{
+					hooks: map[string]server.PostStartHookFunc{},
+					err:   errors.New("hook provider error"),
+				},
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := server.NewRecommendedConfig(apiserver.Codecs)
+			err := builder.AddPostStartHooks(config, tt.builders)
+			if tt.wantErr {
+				require.Error(t, err)
+			}
+
+			if len(tt.wantHooks) > 0 {
+				for _, hookName := range tt.wantHooks {
+					_, ok := config.PostStartHooks[hookName]
+					require.True(t, ok)
+				}
+			}
+		})
+	}
+}
+
+var _ builder.APIGroupBuilder = &mockAPIGroupPostStartHookProvider{}
+var _ builder.APIGroupPostStartHookProvider = &mockAPIGroupPostStartHookProvider{}
+
+type mockAPIGroupPostStartHookProvider struct {
+	hooks map[string]server.PostStartHookFunc
+	err   error
+}
+
+func (m *mockAPIGroupPostStartHookProvider) GetPostStartHooks() (map[string]server.PostStartHookFunc, error) {
+	return m.hooks, m.err
+}
+
+func (m *mockAPIGroupPostStartHookProvider) GetGroupVersion() schema.GroupVersion {
+	return schema.GroupVersion{}
+}
+
+func (m *mockAPIGroupPostStartHookProvider) InstallSchema(scheme *runtime.Scheme) error {
+	return nil
+}
+
+func (m *mockAPIGroupPostStartHookProvider) UpdateAPIGroupInfo(apiGroupInfo *server.APIGroupInfo, opts builder.APIGroupOptions) error {
+	return nil
+}
+
+func (m *mockAPIGroupPostStartHookProvider) GetOpenAPIDefinitions() common.GetOpenAPIDefinitions {
+	return nil
+}
+
+func (m *mockAPIGroupPostStartHookProvider) GetAuthorizer() authorizer.Authorizer {
+	return nil
+}