feat: add kOps support (#27)

castai · Jun 9, 2021 · 0523aeb · 0523aeb
1 parent 8f1df5e
commit 0523aeb
Show file tree

Hide file tree

Showing 9 changed files with 677 additions and 33 deletions.
diff --git a/internal/castai/types.go b/internal/castai/types.go
@@ -18,11 +18,19 @@ type GKEParams struct {
 	ClusterName string `json:"clusterName"`
 }
 
+type KOPSParams struct {
+	CSP         string `json:"cloud"`
+	Region      string `json:"region"`
+	ClusterName string `json:"clusterName"`
+	StateStore  string `json:"stateStore"`
+}
+
 type RegisterClusterRequest struct {
-	ID   uuid.UUID  `json:"id"`
-	Name string     `json:"name"`
-	EKS  *EKSParams `json:"eks"`
-	GKE  *GKEParams `json:"gke"`
+	ID   uuid.UUID   `json:"id"`
+	Name string      `json:"name"`
+	EKS  *EKSParams  `json:"eks"`
+	GKE  *GKEParams  `json:"gke"`
+	KOPS *KOPSParams `json:"kops"`
 }
 
 type Cluster struct {

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -15,6 +15,7 @@ type Config struct {
 	CASTAI     *CASTAI
 	EKS        *EKS
 	GKE        *GKE
+	KOPS       *KOPS
 }
 
 type Log struct {
@@ -43,6 +44,13 @@ type GKE struct {
 	ClusterName string
 }
 
+type KOPS struct {
+	CSP         string
+	Region      string
+	ClusterName string
+	StateStore  string
+}
+
 var cfg *Config
 
 // Get configuration bound to environment variables.
@@ -71,6 +79,11 @@ func Get() Config {
 	_ = viper.BindEnv("gke.projectid", "GKE_PROJECT_ID")
 	_ = viper.BindEnv("gke.clustername", "GKE_CLUSTER_NAME")
 
+	_ = viper.BindEnv("kops.csp", "KOPS_CSP")
+	_ = viper.BindEnv("kops.region", "KOPS_REGION")
+	_ = viper.BindEnv("kops.clustername", "KOPS_CLUSTER_NAME")
+	_ = viper.BindEnv("kops.statestore", "KOPS_STATE_STORE")
+
 	cfg = &Config{}
 	if err := viper.Unmarshal(&cfg); err != nil {
 		panic(fmt.Errorf("parsing configuration: %v", err))
@@ -89,34 +102,49 @@ func Get() Config {
 
 	if cfg.CASTAI != nil {
 		if cfg.CASTAI.ClusterID == "" {
-			requiredDiscoveryDisabled("CASTAI_CLUSTER_ID")
+			requiredWhenDiscoveryDisabled("CASTAI_CLUSTER_ID")
 		}
 		if cfg.CASTAI.OrganizationID == "" {
-			requiredDiscoveryDisabled("CASTAI_ORGANIZATION_ID")
+			requiredWhenDiscoveryDisabled("CASTAI_ORGANIZATION_ID")
 		}
 	}
 
 	if cfg.EKS != nil {
 		if cfg.EKS.AccountID == "" {
-			requiredDiscoveryDisabled("EKS_ACCOUNT_ID")
+			requiredWhenDiscoveryDisabled("EKS_ACCOUNT_ID")
 		}
 		if cfg.EKS.Region == "" {
-			requiredDiscoveryDisabled("EKS_REGION")
+			requiredWhenDiscoveryDisabled("EKS_REGION")
 		}
 		if cfg.EKS.ClusterName == "" {
-			requiredDiscoveryDisabled("EKS_CLUSTER_NAME")
+			requiredWhenDiscoveryDisabled("EKS_CLUSTER_NAME")
 		}
 	}
 
 	if cfg.GKE != nil {
 		if cfg.GKE.Region == "" {
-			requiredDiscoveryDisabled("GKE_REGION")
+			requiredWhenDiscoveryDisabled("GKE_REGION")
 		}
 		if cfg.GKE.ProjectID == "" {
-			requiredDiscoveryDisabled("GKE_PROJECT_ID")
+			requiredWhenDiscoveryDisabled("GKE_PROJECT_ID")
 		}
 		if cfg.GKE.ClusterName == "" {
-			requiredDiscoveryDisabled("GKE_CLUSTER_NAME")
+			requiredWhenDiscoveryDisabled("GKE_CLUSTER_NAME")
+		}
+	}
+
+	if cfg.KOPS != nil {
+		if cfg.KOPS.CSP == "" {
+			requiredWhenDiscoveryDisabled("KOPS_CSP")
+		}
+		if cfg.KOPS.Region == "" {
+			requiredWhenDiscoveryDisabled("KOPS_REGION")
+		}
+		if cfg.KOPS.ClusterName == "" {
+			requiredWhenDiscoveryDisabled("KOPS_CLUSTER_NAME")
+		}
+		if cfg.KOPS.StateStore == "" {
+			requiredWhenDiscoveryDisabled("KOPS_STATE_STORE")
 		}
 	}
 
@@ -132,6 +160,6 @@ func required(variable string) {
 	panic(fmt.Errorf("env variable %s is required", variable))
 }
 
-func requiredDiscoveryDisabled(variable string) {
+func requiredWhenDiscoveryDisabled(variable string) {
 	panic(fmt.Errorf("env variable %s is required when discovery is disabled", variable))
 }
diff --git a/internal/services/providers/gke/gke.go b/internal/services/providers/gke/gke.go
@@ -16,10 +16,10 @@ import (
 const (
 	Name = "gke"
 
-	labelPreemptible = "cloud.google.com/gke-preemptible"
+	LabelPreemptible = "cloud.google.com/gke-preemptible"
 )
 
-func New(_ context.Context, log logrus.FieldLogger) (types.Provider, error) {
+func New(log logrus.FieldLogger) (types.Provider, error) {
 	return &Provider{log: log}, nil
 }
 
@@ -53,7 +53,7 @@ func (p *Provider) IsSpot(_ context.Context, node *corev1.Node) (bool, error) {
 		return true, nil
 	}
 
-	if val, ok := node.Labels[labelPreemptible]; ok && val == "true" {
+	if val, ok := node.Labels[LabelPreemptible]; ok && val == "true" {
 		return true, nil
 	}
 

diff --git a/internal/services/providers/gke/gke_test.go b/internal/services/providers/gke/gke_test.go
@@ -66,7 +66,7 @@ func TestProvider_IsSpot(t *testing.T) {
 		},
 		{
 			name:     "gke spot node",
-			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{labelPreemptible: "true"}}},
+			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{LabelPreemptible: "true"}}},
 			expected: true,
 		},
 		{