feat: include prebuilt action with overrides

Signed-off-by: ashutosh16 <[email protected]>
argoproj · Aug 28, 2024 · fb7525f · fb7525f
1 parent f597912
commit fb7525f
Show file tree

Hide file tree

Showing 6 changed files with 295 additions and 218 deletions.
diff --git a/pkg/apis/application/v1alpha1/generated.pb.go b/pkg/apis/application/v1alpha1/generated.pb.go
diff --git a/pkg/apis/application/v1alpha1/generated.proto b/pkg/apis/application/v1alpha1/generated.proto
diff --git a/pkg/apis/application/v1alpha1/types.go b/pkg/apis/application/v1alpha1/types.go
@@ -2092,6 +2092,7 @@ func (o *ResourceOverride) GetActions() (ResourceActions, error) {
 type ResourceActions struct {
 	ActionDiscoveryLua string                     `json:"discovery.lua,omitempty" yaml:"discovery.lua,omitempty" protobuf:"bytes,1,opt,name=actionDiscoveryLua"`
 	Definitions        []ResourceActionDefinition `json:"definitions,omitempty" protobuf:"bytes,2,rep,name=definitions"`
+	AddBuildInActions  bool                       `json:"addBuiltInActions,omitempty" protobuf:"bytes,3,opt,name=addBuiltInActions"`
 }
 
 // TODO: describe this type

diff --git a/server/application/application.go b/server/application/application.go
@@ -2368,14 +2368,14 @@ func (s *Server) getAvailableActions(resourceOverrides map[string]appv1.Resource
 		ResourceOverrides: resourceOverrides,
 	}
 
-	discoveryScript, err := luaVM.GetResourceActionDiscovery(obj)
+	discoveryScripts, err := luaVM.GetResourceActionDiscovery(obj)
 	if err != nil {
 		return nil, fmt.Errorf("error getting Lua discovery script: %w", err)
 	}
-	if discoveryScript == "" {
+	if len(discoveryScripts) == 0 {
 		return []appv1.ResourceAction{}, nil
 	}
-	availableActions, err := luaVM.ExecuteResourceActionDiscovery(obj, discoveryScript)
+	availableActions, err := luaVM.ExecuteResourceActionDiscovery(obj, discoveryScripts)
 	if err != nil {
 		return nil, fmt.Errorf("error executing Lua discovery script: %w", err)
 	}

diff --git a/util/lua/lua.go b/util/lua/lua.go
@@ -275,48 +275,57 @@ func cleanReturnedArray(newObj, obj []interface{}) []interface{} {
 	return arrayToReturn
 }
 
-func (vm VM) ExecuteResourceActionDiscovery(obj *unstructured.Unstructured, script string) ([]appv1.ResourceAction, error) {
-	l, err := vm.runLua(obj, script)
-	if err != nil {
-		return nil, err
+func (vm VM) ExecuteResourceActionDiscovery(obj *unstructured.Unstructured, script []string) ([]appv1.ResourceAction, error) {
+	if len(script) == 0 {
+		return nil, fmt.Errorf("no action discovery script provided")
 	}
-	returnValue := l.Get(-1)
-	if returnValue.Type() == lua.LTTable {
-		jsonBytes, err := luajson.Encode(returnValue)
-		if err != nil {
-			return nil, err
-		}
-		availableActions := make([]appv1.ResourceAction, 0)
-		if noAvailableActions(jsonBytes) {
-			return availableActions, nil
-		}
-		availableActionsMap := make(map[string]interface{})
-		err = json.Unmarshal(jsonBytes, &availableActionsMap)
+	availableActions := make([]appv1.ResourceAction, 0)
+
+	for _, script := range script {
+		l, err := vm.runLua(obj, script)
 		if err != nil {
 			return nil, err
 		}
-		for key := range availableActionsMap {
-			value := availableActionsMap[key]
-			resourceAction := appv1.ResourceAction{Name: key, Disabled: isActionDisabled(value)}
-			if emptyResourceActionFromLua(value) {
-				availableActions = append(availableActions, resourceAction)
-				continue
-			}
-			resourceActionBytes, err := json.Marshal(value)
+		returnValue := l.Get(-1)
+		if returnValue.Type() == lua.LTTable {
+			jsonBytes, err := luajson.Encode(returnValue)
 			if err != nil {
 				return nil, err
 			}
-
-			err = json.Unmarshal(resourceActionBytes, &resourceAction)
+			if noAvailableActions(jsonBytes) {
+				return availableActions, nil
+			}
+			availableActionsMap := make(map[string]interface{})
+			err = json.Unmarshal(jsonBytes, &availableActionsMap)
 			if err != nil {
 				return nil, err
 			}
-			availableActions = append(availableActions, resourceAction)
+			for key := range availableActionsMap {
+				value := availableActionsMap[key]
+				resourceAction := appv1.ResourceAction{Name: key, Disabled: isActionDisabled(value)}
+				if emptyResourceActionFromLua(value) {
+					availableActions = append(availableActions, resourceAction)
+					continue
+				}
+				resourceActionBytes, err := json.Marshal(value)
+				if err != nil {
+					return nil, err
+				}
+
+				err = json.Unmarshal(resourceActionBytes, &resourceAction)
+				if err != nil {
+					return nil, err
+				}
+				availableActions = append(availableActions, resourceAction)
+			}
+
+		} else {
+			return nil, fmt.Errorf(incorrectReturnType, "table", returnValue.Type().String())
+
 		}
-		return availableActions, err
 	}
 
-	return nil, fmt.Errorf(incorrectReturnType, "table", returnValue.Type().String())
+	return availableActions, nil
 }
 
 // Actions are enabled by default
@@ -346,22 +355,38 @@ func noAvailableActions(jsonBytes []byte) bool {
 	return string(jsonBytes) == "[]"
 }
 
-func (vm VM) GetResourceActionDiscovery(obj *unstructured.Unstructured) (string, error) {
+func (vm VM) GetResourceActionDiscovery(obj *unstructured.Unstructured) ([]string, error) {
 	key := GetConfigMapKey(obj.GroupVersionKind())
+	var discoveryScripts []string
+
 	override, ok := vm.ResourceOverrides[key]
 	if ok && override.Actions != "" {
 		actions, err := override.GetActions()
 		if err != nil {
-			return "", err
+			return nil, err
+		}
+		if actions.AddBuildInActions {
+			discoveryScripts = append(discoveryScripts, actions.ActionDiscoveryLua)
+		} else {
+			return []string{actions.ActionDiscoveryLua}, nil
 		}
-		return actions.ActionDiscoveryLua, nil
 	}
+
+	// Fetch predefined Lua scripts
 	discoveryKey := fmt.Sprintf("%s/actions/", key)
 	discoveryScript, err := vm.getPredefinedLuaScripts(discoveryKey, actionDiscoveryScriptFile)
 	if err != nil {
-		return "", err
+		return nil, err
+	}
+
+	// Append or return the discovery script based on the presence of built-in actions
+	if len(discoveryScripts) > 0 {
+		discoveryScripts = append(discoveryScripts, discoveryScript)
+	} else {
+		return []string{discoveryScript}, nil
 	}
-	return discoveryScript, nil
+
+	return discoveryScripts, nil
 }
 
 // GetResourceAction attempts to read lua script from config and then filesystem for that resource

diff --git a/util/lua/lua_test.go b/util/lua/lua_test.go
@@ -324,7 +324,7 @@ func TestGetResourceActionDiscoveryNoPredefined(t *testing.T) {
 	vm := VM{}
 	discoveryLua, err := vm.GetResourceActionDiscovery(testObj)
 	require.NoError(t, err)
-	assert.Empty(t, discoveryLua)
+	assert.Equal(t, discoveryLua[0], "")
 }
 
 func TestGetResourceActionDiscoveryWithOverride(t *testing.T) {
@@ -340,7 +340,25 @@ func TestGetResourceActionDiscoveryWithOverride(t *testing.T) {
 	}
 	discoveryLua, err := vm.GetResourceActionDiscovery(testObj)
 	require.NoError(t, err)
-	assert.Equal(t, validDiscoveryLua, discoveryLua)
+	assert.Equal(t, validDiscoveryLua, discoveryLua[0])
+}
+
+func TestGetResourceActionsWithAddBuiltInActionsFlag(t *testing.T) {
+	testObj := StrToUnstructured(objJSON)
+	vm := VM{
+		ResourceOverrides: map[string]appv1.ResourceOverride{
+			"argoproj.io/Rollout": {
+				Actions: string(grpc.MustMarshal(appv1.ResourceActions{
+					ActionDiscoveryLua: validDiscoveryLua,
+					AddBuildInActions:  true,
+				})),
+			},
+		},
+	}
+
+	discoveryLua, err := vm.GetResourceActionDiscovery(testObj)
+	require.NoError(t, err)
+	assert.Equal(t, validDiscoveryLua, discoveryLua[0])
 }
 
 const validDiscoveryLua = `
@@ -358,7 +376,7 @@ return a
 func TestExecuteResourceActionDiscovery(t *testing.T) {
 	testObj := StrToUnstructured(objJSON)
 	vm := VM{}
-	actions, err := vm.ExecuteResourceActionDiscovery(testObj, validDiscoveryLua)
+	actions, err := vm.ExecuteResourceActionDiscovery(testObj, []string{validDiscoveryLua})
 	require.NoError(t, err)
 	expectedActions := []appv1.ResourceAction{
 		{
@@ -386,7 +404,7 @@ return a`
 func TestExecuteResourceActionDiscoveryInvalidResourceAction(t *testing.T) {
 	testObj := StrToUnstructured(objJSON)
 	vm := VM{}
-	actions, err := vm.ExecuteResourceActionDiscovery(testObj, discoveryLuaWithInvalidResourceAction)
+	actions, err := vm.ExecuteResourceActionDiscovery(testObj, []string{discoveryLuaWithInvalidResourceAction})
 	require.Error(t, err)
 	assert.Nil(t, actions)
 }
@@ -399,7 +417,7 @@ return a
 func TestExecuteResourceActionDiscoveryInvalidReturn(t *testing.T) {
 	testObj := StrToUnstructured(objJSON)
 	vm := VM{}
-	actions, err := vm.ExecuteResourceActionDiscovery(testObj, invalidDiscoveryLua)
+	actions, err := vm.ExecuteResourceActionDiscovery(testObj, []string{invalidDiscoveryLua})
 	assert.Nil(t, actions)
 	require.Error(t, err)
 }