Merge pull request #218 from igomez06/main-vai-extension

Vai toggle extension.

extensions/features/features.go

@@ -0,0 +1,5 @@
+package features
+
+const (
+	ManagementFeature = "management.cattle.io.feature"
+)

extensions/features/update.go

@@ -0,0 +1,24 @@
+package features
+
+import (
+	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
+	v1 "github.com/rancher/shepherd/clients/rancher/v1"
+)
+
+// UpdateFeatureFlag is a helper function that uses the steve client to update a Global setting
+func UpdateFeatureFlag(steveclient *v1.Client, featureFlag *v1.SteveAPIObject, value bool) (*v1.SteveAPIObject, error) {
+	updateFeature := &v3.Feature{}
+	err := v1.ConvertToK8sType(featureFlag.JSONResp, updateFeature)
+	if err != nil {
+		return nil, err
+	}
+
+	updateFeature.Spec.Value = &value
+
+	updateFeatureFlag, err := steveclient.SteveType(ManagementFeature).Update(featureFlag, updateFeature)
+	if err != nil {
+		return nil, err
+	}
+
+	return updateFeatureFlag, nil
+}

extensions/settings/update.go

@@ -5,7 +5,24 @@ import (
 	v1 "github.com/rancher/shepherd/clients/rancher/v1"
 )
 
-// UpdateGlobalSettings is a helper function that uses the steve client to update a Global setting
+// ResetUGlobalSettings is a helper function that uses the steve client to reset a specific
+// Global Setting to its default value.
+func ResetUGlobalSettings(steveclient *v1.Client, globalSetting *v1.SteveAPIObject) (*v1.SteveAPIObject, error) {
+	updateSetting := &v3.Setting{}
+	err := v1.ConvertToK8sType(globalSetting.JSONResp, updateSetting)
+	if err != nil {
+		return nil, err
+	}
+
+	updateSetting.Value = updateSetting.Default
+	updateGlobalSetting, err := steveclient.SteveType(ManagementSetting).Update(globalSetting, updateSetting)
+	if err != nil {
+		return nil, err
+	}
+	return updateGlobalSetting, nil
+}
+
+// UpdateGlobalSettings is a helper function that uses the steve client to update a Global setting.
 func UpdateGlobalSettings(steveclient *v1.Client, globalSetting *v1.SteveAPIObject, value string) (*v1.SteveAPIObject, error) {
 	updateSetting := &v3.Setting{}
 	err := v1.ConvertToK8sType(globalSetting.JSONResp, updateSetting)

extensions/vai/vai.go

@@ -0,0 +1,108 @@
+package vai
+
+import (
+	"fmt"
+
+	"github.com/rancher/shepherd/clients/rancher"
+	"github.com/rancher/shepherd/extensions/features"
+	"github.com/rancher/shepherd/extensions/workloads/pods"
+	"k8s.io/apimachinery/pkg/types"
+)
+
+const (
+	cattleSystemNamespace    = "cattle-system"
+	cattleAgentDeployment    = "cattle-cluster-agent"
+	clusterRegisterContainer = "cluster-register"
+	uiSQLCacheResource       = "ui-sql-cache"
+	cattleFeaturesEnvVar     = "CATTLE_FEATURES"
+)
+
+func updateSQLCache(adminClient *rancher.Client, value bool) error {
+	managementClient := adminClient.Steve.SteveType("management.cattle.io.feature")
+
+	steveCacheFlagResp, err := managementClient.ByID(uiSQLCacheResource)
+	if err != nil {
+		return err
+	}
+
+	_, err = features.UpdateFeatureFlag(adminClient.Steve, steveCacheFlagResp, value)
+	if err != nil {
+		return err
+	}
+
+	errors := pods.StatusPods(adminClient, "local")
+	if len(errors) > 1 {
+		return fmt.Errorf("error when restarting pods")
+	}
+
+	return nil
+}
+
+// EnableVaiCaching is the extension that sets all the appropriate global/performance settings, and feature flags
+// to enable the vai sql caching.
+func EnableVaiCaching(adminClient *rancher.Client) error {
+	err := updateSQLCache(adminClient, true)
+	if err != nil {
+		return err
+	}
+
+	adminClient.Session.RegisterCleanupFunc(func() error {
+		err := DisableVaiCaching(adminClient)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+
+	return nil
+}
+
+// DisableVaiCaching is the extension that sets all the appropriate global/performance settings, and feature flags
+// to disable the vai sql caching.
+func DisableVaiCaching(adminClient *rancher.Client) error {
+	err := updateSQLCache(adminClient, false)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func setDownstreamClusterSQLCaching(adminClient *rancher.Client, clusterID, value string) error {
+	downStreamController, err := adminClient.WranglerContext.DownStreamClusterWranglerContext(clusterID)
+	if err != nil {
+		return err
+	}
+
+	patchedResource := fmt.Sprintf(`{"spec":{"template":{"spec":{"containers":[{"name":"cluster-register","env":[{"name":"CATTLE_FEATURES","value":"embedded-cluster-api=false,fleet=false,multi-cluster-management=false,multi-cluster-management-agent=true,provisioningv2=false,rke2=false%v"}]}]}}}}`, value)
+
+	_, err = downStreamController.Apps.Deployment().Patch(cattleSystemNamespace, cattleAgentDeployment, types.StrategicMergePatchType, []byte(patchedResource))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// EnableDownstreamClusterSQLCaching is a function that propagates the caching feature flag to the downstream cluster
+func EnableDownstreamClusterSQLCaching(adminClient *rancher.Client, clusterID string) error {
+	err := setDownstreamClusterSQLCaching(adminClient, clusterID, fmt.Sprintf(",%v=true", uiSQLCacheResource))
+	if err != nil {
+		return err
+	}
+
+	adminClient.Session.RegisterCleanupFunc(func() error {
+		err := DisableDownstreamClusterSQLCaching(adminClient, clusterID)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+
+	return nil
+}
+
+// DisableDownstreamClusterSQLCaching is a function that removes, if it's there, the ui-sql-cache=true flag.
+func DisableDownstreamClusterSQLCaching(adminClient *rancher.Client, clusterID string) error {
+	return setDownstreamClusterSQLCaching(adminClient, clusterID, "")
+}

go.mod

@@ -32,7 +32,6 @@ require (
 	github.com/rancher/rancher/pkg/apis v0.0.0-20240529155832-64dbde013b4c
 	github.com/rancher/rke v1.6.0-rc4
 	github.com/rancher/system-upgrade-controller/pkg/apis v0.0.0-20210727200656-10b094e30007
-	github.com/rancher/types v0.0.0-20220328215343-4370ff10ecd5
 	github.com/rancher/wrangler v1.1.1
 	github.com/rancher/wrangler/v2 v2.2.0-rc6
 	github.com/sirupsen/logrus v1.9.3