Introduced configuration for K8S API client (#123)

internal/app/app.go

@@ -128,8 +128,8 @@ func (p *PolicyAutomationApp) LoadConfig(config *cfg.Config) (err error) {
 		if p.config.CredentialsFile != "" {
 			builder = builder.WithCredentialsFile(p.config.CredentialsFile)
 		}
-		if p.config.K8SCheck {
-			builder = builder.WithK8SClient(cfg.APIVERSIONS)
+		if p.config.K8SApiConfig.Enabled {
+			builder = builder.WithK8SClient(config.K8SApiConfig.ApiVersions, config.K8SApiConfig.MaxQPS)
 		}
 		p.gke, err = builder.Build()
 		if err != nil {
@@ -351,7 +351,7 @@ func newConfigFromFile(path string) (*cfg.Config, error) {
 func newConfigFromCli(cliConfig *CliConfig) *cfg.Config {
 	config := &cfg.Config{}
 	config.SilentMode = cliConfig.SilentMode
-	config.K8SCheck = cliConfig.K8SCheck
+	config.K8SApiConfig.Enabled = cliConfig.K8SCheck
 	config.CredentialsFile = cliConfig.CredentialsFile
 	config.DumpFile = cliConfig.DumpFile
 	if cliConfig.DiscoveryEnabled {

internal/app/app_test.go

@@ -26,6 +26,7 @@ import (
 
 	cfg "github.com/google/gke-policy-automation/internal/config"
 	"github.com/google/gke-policy-automation/internal/outputs"
+	"github.com/stretchr/testify/assert"
 )
 
 type DiscoveryClientMock struct {
@@ -134,6 +135,10 @@ func TestLoadCliConfig_defaults(t *testing.T) {
 	if !reflect.DeepEqual(policy, defaultPolicy) {
 		t.Error("config policy is not same as default policy")
 	}
+	if pa.config.K8SApiConfig.MaxQPS != cfg.DefaultK8SClientQPS {
+		t.Errorf("K8SApiConfig MaxQPS = %v; want %v", pa.config.K8SApiConfig.MaxQPS, cfg.DefaultK8SClientQPS)
+	}
+	assert.ElementsMatchf(t, pa.config.K8SApiConfig.ApiVersions, cfg.DefaultK8SApiVersions, "K8SApiConfig ApiVersions match")
 }
 
 func TestLoadConfig(t *testing.T) {

internal/app/cli.go

@@ -88,7 +88,7 @@ func createCheckCommand(p PolicyAutomation) *cli.Command {
 				Action: func(c *cli.Context) error {
 					defer p.Close()
 					config.K8SCheck = true
-					if err := p.LoadCliConfig(config, cfg.ValidateClusterCheckConfig); err != nil {
+					if err := p.LoadCliConfig(config, cfg.ValidateScalabilityCheckConfig); err != nil {
 						cli.ShowSubcommandHelp(c)
 						return err
 					}

internal/app/test-fixtures/test_config.yaml

@@ -28,4 +28,9 @@ outputs:
 - file: /some/file.json
 - pubsub:
     project: my-pubsub-project
-    topic: my-topic
+    topic: my-topic
+kubernetesAPIClient:
+  enabled: true
+  clientMaxQPS: 55
+  resourceAPIVersions:
+    - v1

internal/config/config.go

@@ -15,6 +15,7 @@
 package config
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 
@@ -23,28 +24,29 @@ import (
 )
 
 var (
-	APIVERSIONS = []string{"v1", "autoscaling/v1"}
+	DefaultK8SApiVersions = []string{"v1", "autoscaling/v1"}
 )
 
 const (
 	DefaultGitRepository = "https://github.com/google/gke-policy-automation"
 	DefaultGitBranch     = "main"
 	DefaultGitPolicyDir  = "gke-policies"
+	DefaultK8SClientQPS  = 50
 )
 
 type ReadFileFn func(string) ([]byte, error)
 type ValidateConfig func(config Config) error
 
 type Config struct {
 	SilentMode       bool                   `yaml:"silent"`
-	K8SCheck         bool                   `yaml:"k8sCheck"`
 	DumpFile         string                 `yaml:"dumpFile"`
 	CredentialsFile  string                 `yaml:"credentialsFile"`
 	Clusters         []ConfigCluster        `yaml:"clusters"`
 	Policies         []ConfigPolicy         `yaml:"policies"`
 	Outputs          []ConfigOutput         `yaml:"outputs"`
 	ClusterDiscovery ClusterDiscovery       `yaml:"clusterDiscovery"`
 	PolicyExclusions ConfigPolicyExclusions `yaml:"policyExclusions"`
+	K8SApiConfig     K8SApiConfig           `yaml:"kubernetesAPIClient"`
 }
 
 type ConfigPolicy struct {
@@ -96,6 +98,12 @@ type ConfigPolicyExclusions struct {
 	PolicyGroups []string `yaml:"policyGroups"`
 }
 
+type K8SApiConfig struct {
+	Enabled     bool     `yaml:"enabled"`
+	ApiVersions []string `yaml:"resourceAPIVersions"`
+	MaxQPS      int      `yaml:"clientMaxQPS"`
+}
+
 func ReadConfig(path string, readFn ReadFileFn) (*Config, error) {
 	data, err := readFn(path)
 	if err != nil {
@@ -182,6 +190,16 @@ func ValidateGeneratePolicyDocsConfig(config Config) error {
 	return nil
 }
 
+func ValidateScalabilityCheckConfig(config Config) error {
+	if err := ValidateClusterCheckConfig(config); err != nil {
+		return nil
+	}
+	if !config.K8SApiConfig.Enabled {
+		return errors.New("kubernetes API client is disabled")
+	}
+	return nil
+}
+
 func validateClustersConfig(config Config) []error {
 	if config.ClusterDiscovery.Enabled {
 		discovery := config.ClusterDiscovery
@@ -284,4 +302,10 @@ func SetConfigDefaults(config *Config) {
 			GitBranch:     DefaultGitBranch,
 			GitDirectory:  DefaultGitPolicyDir})
 	}
+	if config.K8SApiConfig.MaxQPS == 0 {
+		config.K8SApiConfig.MaxQPS = DefaultK8SClientQPS
+	}
+	if len(config.K8SApiConfig.ApiVersions) == 0 {
+		config.K8SApiConfig.ApiVersions = DefaultK8SApiVersions
+	}
 }

internal/gke/gke.go

@@ -44,6 +44,7 @@ type gkeApiClientBuilder struct {
 	ctx             context.Context
 	credentialsFile string
 	k8sApiVersions  []string
+	k8sMaxQPS       int
 }
 
 func NewGKEApiClientBuilder(ctx context.Context) *gkeApiClientBuilder {
@@ -55,8 +56,9 @@ func (b *gkeApiClientBuilder) WithCredentialsFile(credentialsFile string) *gkeAp
 	return b
 }
 
-func (b *gkeApiClientBuilder) WithK8SClient(apiVersions []string) *gkeApiClientBuilder {
+func (b *gkeApiClientBuilder) WithK8SClient(apiVersions []string, maxQPS int) *gkeApiClientBuilder {
 	b.k8sApiVersions = apiVersions
+	b.k8sMaxQPS = maxQPS
 	return b
 }
 
@@ -71,29 +73,26 @@ func (b *gkeApiClientBuilder) Build() (GKEClient, error) {
 		return nil, err
 	}
 
-	var k8sApiVersions []string
-	if len(b.k8sApiVersions) > 0 {
-		k8sApiVersions = b.k8sApiVersions
-	}
-
 	return &GKEApiClient{
 		ctx:            b.ctx,
 		client:         cli,
 		authTokenFunc:  getClusterToken,
 		k8sClientFunc:  NewKubernetesClient,
-		k8sApiVersions: k8sApiVersions,
+		k8sApiVersions: b.k8sApiVersions,
+		k8sMaxQPS:      b.k8sMaxQPS,
 	}, nil
 }
 
 type authTokenFunc func(ctx context.Context) (string, error)
-type k8sClientFunc func(ctx context.Context, kubeConfig *clientcmdapi.Config) (KubernetesClient, error)
+type k8sClientFunc func(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int) (KubernetesClient, error)
 
 type GKEApiClient struct {
 	ctx            context.Context
 	client         ClusterManagerClient
 	k8sClientFunc  k8sClientFunc
 	authTokenFunc  authTokenFunc
 	k8sApiVersions []string
+	k8sMaxQPS      int
 }
 
 type Cluster struct {
@@ -138,7 +137,7 @@ func (c *GKEApiClient) GetCluster(name string) (*Cluster, error) {
 			log.Debugf("unable to get kubeconfig: %s", err)
 			return nil, err
 		}
-		k8cli, err := c.k8sClientFunc(c.ctx, kubeConfig)
+		k8cli, err := c.k8sClientFunc(c.ctx, kubeConfig, c.k8sMaxQPS)
 		if err != nil {
 			return nil, err
 		}

internal/gke/gke_test.go

@@ -54,7 +54,9 @@ func (mockClusterManagerClient) Close() error {
 
 func TestNewGKEClient(t *testing.T) {
 	testCredsFile := "test-fixtures/test_credentials.json"
-	c, err := NewGKEApiClientBuilder(context.Background()).WithCredentialsFile(testCredsFile).WithK8SClient(config.APIVERSIONS).Build()
+	c, err := NewGKEApiClientBuilder(context.Background()).WithCredentialsFile(testCredsFile).
+		WithK8SClient([]string{"v1"}, config.DefaultK8SClientQPS).
+		Build()
 	if err != nil {
 		t.Fatalf("error when creating client: %v", err)
 	}
@@ -126,9 +128,10 @@ func (mockK8Client) GetResources(resourceType []*ResourceType, namespace []strin
 func TestGKEApiClientBuilder(t *testing.T) {
 	credFile := "test-fixtures/test_credentials.json"
 	apiVersions := []string{"policy/v1", "networking.k8s.io/v1"}
+	maxQPS := 69
 	b := NewGKEApiClientBuilder(context.TODO()).
 		WithCredentialsFile(credFile).
-		WithK8SClient(apiVersions)
+		WithK8SClient(apiVersions, maxQPS)
 	client, err := b.Build()
 	if err != nil {
 		t.Fatalf("err = %v, want nil", err)
@@ -140,6 +143,9 @@ func TestGKEApiClientBuilder(t *testing.T) {
 	if !reflect.DeepEqual(apiClient.k8sApiVersions, apiVersions) {
 		t.Errorf("apiClient k8sApiVersions = %v; want %v", apiClient.k8sApiVersions, apiVersions)
 	}
+	if apiClient.k8sMaxQPS != maxQPS {
+		t.Errorf("apiClient k8sMaxQPS = %v; want %v", apiClient.k8sMaxQPS, maxQPS)
+	}
 	if b.credentialsFile != credFile {
 		t.Errorf("builder credentialsFile = %v; want %v", b.credentialsFile, credFile)
 	}
@@ -149,7 +155,7 @@ func TestGetCluster(t *testing.T) {
 	client := GKEApiClient{
 		ctx:    context.Background(),
 		client: &mockClusterManagerClient{},
-		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config) (KubernetesClient, error) {
+		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int) (KubernetesClient, error) {
 			return &mockK8Client{}, nil
 
 		},
@@ -227,7 +233,7 @@ func TestGetClusterResourcesForEmptyConfig(t *testing.T) {
 	client := GKEApiClient{
 		ctx:    context.Background(),
 		client: &mockClusterManagerClient{},
-		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config) (KubernetesClient, error) {
+		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int) (KubernetesClient, error) {
 			return &mockK8Client{}, nil
 
 		},
@@ -252,7 +258,7 @@ func TestGetClusterResourcesForNonEmptyConfig(t *testing.T) {
 	client := GKEApiClient{
 		ctx:    context.Background(),
 		client: &mockClusterManagerClient{},
-		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config) (KubernetesClient, error) {
+		k8sClientFunc: func(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int) (KubernetesClient, error) {
 			return &mockK8Client{}, nil
 
 		},

internal/gke/kubernetes.go

@@ -21,10 +21,12 @@ import (
 	"sync"
 
 	"github.com/google/gke-policy-automation/internal/log"
+	"github.com/google/gke-policy-automation/internal/version"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/client-go/discovery"
 	"k8s.io/client-go/dynamic"
+	restclient "k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
 )
@@ -71,14 +73,12 @@ const (
 	defaultMaxGoroutines = 10
 )
 
-func NewKubernetesClient(ctx context.Context, kubeConfig *clientcmdapi.Config) (KubernetesClient, error) {
-	return NewKubernetesClientWithConfiguredGoroutines(ctx, kubeConfig, defaultMaxGoroutines)
+func NewKubernetesClient(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int) (KubernetesClient, error) {
+	return NewKubernetesClientWithConfiguredGoroutines(ctx, kubeConfig, maxQPS, defaultMaxGoroutines)
 }
 
-func NewKubernetesClientWithConfiguredGoroutines(ctx context.Context, kubeConfig *clientcmdapi.Config, maxGoRoutines int) (KubernetesClient, error) {
-	config, err := clientcmd.BuildConfigFromKubeconfigGetter("", func() (*clientcmdapi.Config, error) {
-		return kubeConfig, nil
-	})
+func NewKubernetesClientWithConfiguredGoroutines(ctx context.Context, kubeConfig *clientcmdapi.Config, maxQPS int, maxGoRoutines int) (KubernetesClient, error) {
+	config, err := getKubernetesRestClientConfig(kubeConfig, maxQPS)
 	if err != nil {
 		return nil, err
 	}
@@ -233,3 +233,15 @@ func stringSliceContains(hay []string, needle string) bool {
 	}
 	return false
 }
+
+func getKubernetesRestClientConfig(kubeConfig *clientcmdapi.Config, maxQPS int) (*restclient.Config, error) {
+	config, err := clientcmd.BuildConfigFromKubeconfigGetter("", func() (*clientcmdapi.Config, error) {
+		return kubeConfig, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	config.QPS = float32(maxQPS)
+	config.UserAgent = version.UserAgent
+	return config, nil
+}

internal/gke/kubernetes_test.go

@@ -21,6 +21,8 @@ import (
 
 	b64 "encoding/base64"
 
+	"github.com/google/gke-policy-automation/internal/config"
+	"github.com/google/gke-policy-automation/internal/version"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -133,7 +135,7 @@ func TestNewKubernetesClient(t *testing.T) {
 	}
 
 	ctx := context.TODO()
-	cli, err := NewKubernetesClient(ctx, testConfig)
+	cli, err := NewKubernetesClient(ctx, testConfig, config.DefaultK8SClientQPS)
 	if err != nil {
 		t.Fatalf("err is not nil; want nil; err = %s", err)
 	}
@@ -468,3 +470,35 @@ func createKubeNamespaceResourceMockWithResource(resourceName string) *kubeNames
 		},
 	}
 }
+
+func TestGetKubernetesRestClientConfig(t *testing.T) {
+	maxQPS := 111
+	cert, _ := b64.StdEncoding.DecodeString(`LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUVMRENDQXBTZ0F3SUJBZ0lRRThETy9sTllZaUM0SWNhNm04N1EyVEFOQmdrcWhraUc5dzBCQVFzRkFEQXYKTVMwd0t3WURWUVFERXlSbVpUQTFPREF4TkMweU1UUTFMVFF3TURjdFltVTFNUzB6TkRZMVptWTRNREExWm1JdwpJQmNOTWpJd05URTNNVEUwTmpReVdoZ1BNakExTWpBMU1Ea3hNalEyTkRKYU1DOHhMVEFyQmdOVkJBTVRKR1psCk1EVTRNREUwTFRJeE5EVXROREF3TnkxaVpUVXhMVE0wTmpWbVpqZ3dNRFZtWWpDQ0FhSXdEUVlKS29aSWh2Y04KQVFFQkJRQURnZ0dQQURDQ0FZb0NnZ0dCQUs3TGFyT1U5VXFmNkRvM2JMUDR3aHV1ZUs1Sjd5NCtzT3d1aW1KLwpBYk82cG1lRDk0OGJaOXJhUUUwb0trU2RsZFlROEFUY0FSN3p0RjNhQURpZkVUNGJsYzExQ3o5dEZkUE9iMU02CkxoeU92MkZnNjBuMVhneFdhU1NHMVhLdG11OTVhUWRZbmR6TGNnWFM2MXNXWkgrQk90ZktvRWU4bitGY21JczUKWWZha3Iwb3gvVmZKWXMvWEVEUEw1UUdmcEtVY21kd0FiamQycVJ6MEROUlZLVThyRUJSSVp3ODFVUEJiaVpYVgozNlRHMDZCVFBCNnM3cFJQODFBeG10T3Z1dGZIT1l0STBLNU50Z1Vhc2ZLVGV6aTVsYTB1dDFGWnJQcFZ3NmpUCm91RE4wY1hkMjlyTlhKZDg4L0RpdUN4eXNrYk5xU2JkajBVYTV0UUw1ejltbDU2V2ZCMEszc2dzN0l2N1VOUGQKRGFSN2lUZGxoZWFxUmhQRkV3eW56Szg5NmlmSkltTnowTDBKNTFHekR1TFZmRUtyMFl2VytqdW9MUDl1OWhxdwpXSngxaEpsblhVb2NmY25Va1JMb3Y5VThOem5HTjZIZjhKTFBiZlhMdEtoWWFRMjFCZC9XcFV4TUVnRUJRNWdTCk5TOFRaNTVQRy94Vm14YjZpYzE4QjZNYXdRSURBUUFCbzBJd1FEQU9CZ05WSFE4QkFmOEVCQU1DQWdRd0R3WUQKVlIwVEFRSC9CQVV3QXdFQi96QWRCZ05WSFE0RUZnUVVyRWJaQTNkOGJxdHc5SmdxeFU3cU50YnF1N1F3RFFZSgpLb1pJaHZjTkFRRUxCUUFEZ2dHQkFER3BxRzJFR2t5Z3V3bDhZSnZpS1pBN01uMjl5QjJjV2JGZFIxMU9LdE5ZCnRxSm82ZTA3NlJQMjFyQzNVTVNWeHNadXZ4a3hFQkRwL05SZzlpYjVleEcrU3l3dENZcUZKRWpCUlQwck5YWHMKUWE1NVlFak5WRTJYVFN2NzludUVGSVR6aC9PYVV0S2h2SVdoaWJPN2ZQL1lDUEo4SDFlN0NFYW5UbENId3ZqRwpnK241V0ZwVWJWUk1naElFa0pDM0g3MjlZdmplUWs4Z2pxZDZDdlBDb3p6YTBDRkp3ZVBLQXRGRXMxVmJyKzJiClE4MVdlaS9JWWUwTWZROExWSHl2cGVDbnVWR3Z0OERVRmRFdXVXV0pHOGVBYnZkWGRtWDZqbmJNUFJUUkJoUUEKeUQ4aDZkS2x5a2FVckcrM0Y3N00zek8yQk00ZXhqQ0NmTzNXMWZCaEVjUVR6dU5SRFhLY01LY3F5Z1QyNWhHcgpCM0lVa1U2T3crOG43c3hFWnd0cDF4THppWjJzUmdlNkI4MGloZGxCalc2aHYxMVkvRU9LNkI4M2RYSVo4d0lrCkFqV1BBOFlyZWhkSVpBOHhiOWJoQWNuZ2R2NnJqR0loY2h3N25tK3hQa25RM1pOcXZacjY1Y1RPQmVYU2lLYnoKN2hTT25iTGpNTjhZd05VREUvbUpjZz09Ci0tLS0tRU5EIENFUlRJRklDQVRFLS0tLS0K`)
+	kubeConfig := &clientcmdapi.Config{
+		APIVersion: "v1",
+		Kind:       "Config",
+		Clusters: map[string]*clientcmdapi.Cluster{
+			"cluster": {
+				CertificateAuthorityData: cert,
+				Server:                   `https://1.1.1.1`},
+		},
+		AuthInfos: map[string]*clientcmdapi.AuthInfo{"user": {Token: "test-token"}},
+		Contexts: map[string]*clientcmdapi.Context{
+			"context": {
+				Cluster:  "cluster",
+				AuthInfo: "user",
+			},
+		},
+		CurrentContext: "context",
+	}
+	restConfig, err := getKubernetesRestClientConfig(kubeConfig, maxQPS)
+	if err != nil {
+		t.Fatalf("err is = %v; want nil", err)
+	}
+	if restConfig.QPS != float32(maxQPS) {
+		t.Errorf("restConfig QPS = %v; want %v", restConfig.QPS, maxQPS)
+	}
+	if restConfig.UserAgent != version.UserAgent {
+		t.Errorf("restConfig UserAgent = %v; want %v", restConfig.UserAgent, version.UserAgent)
+	}
+}