feat(lua actions): add a flag to Include builtin actions with resourc…

…e overrides (#19708) * feat: include prebuilt action with overrides Signed-off-by: ashutosh16 <[email protected]> Signed-off-by: ashutosh16 <[email protected]> Signed-off-by: Ashu <[email protected]> Co-authored-by: Alexandre Gaudreault <[email protected]>
argoproj · Sep 10, 2024 · 5776554 · 5776554
1 parent 878494f
commit 5776554
Show file tree

Hide file tree

Showing 7 changed files with 877 additions and 752 deletions.
diff --git a/docs/operator-manual/resource_actions.md b/docs/operator-manual/resource_actions.md
@@ -80,6 +80,20 @@ The `discovery.lua` script must return a table where the key name represents the
 
 Each action name must be represented in the list of `definitions` with an accompanying `action.lua` script to control the resource modifications. The `obj` is a global variable which contains the resource. Each action script returns an optionally modified version of the resource. In this example, we are simply setting `.spec.suspend` to either `true` or `false`.
 
+By default, defining a resource action customization will override any built-in action for this resource kind. If you want to retain the built-in actions, you can set the `mergeBuiltinActions` key to `true`. Your custom actions will have precedence over the built-in actions.
+```yaml        
+resource.customizations.actions.argoproj.io_Rollout: |
+  mergeBuiltinActions: true
+  discovery.lua: |
+    actions = {}
+    actions["do-things"] = {}
+    return actions
+  definitions:
+  - name: do-things
+    action.lua: |
+      return obj		
+```
+
 #### Creating new resources with a custom action
 
 !!! important

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
@@ -2144,8 +2144,9 @@ func (o *ResourceOverride) GetActions() (ResourceActions, error) {
 // TODO: describe this type
 // TODO: describe members of this type
 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"`
+	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"`
+	MergeBuiltinActions bool                       `json:"mergeBuiltinActions,omitempty" yaml:"mergeBuiltinActions,omitempty" protobuf:"bytes,3,opt,name=mergeBuiltinActions"`
 }
 
 // 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,62 @@ 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, scripts []string) ([]appv1.ResourceAction, error) {
+	if len(scripts) == 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)
+	availableActionsMap := make(map[string]appv1.ResourceAction)
+
+	for _, script := range scripts {
+		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)
+		returnValue := l.Get(-1)
+		if returnValue.Type() == lua.LTTable {
+			jsonBytes, err := luajson.Encode(returnValue)
+			if err != nil {
+				return nil, fmt.Errorf("error in converting to lua table: %w", err)
+			}
+			if noAvailableActions(jsonBytes) {
 				continue
 			}
-			resourceActionBytes, err := json.Marshal(value)
+			actionsMap := make(map[string]interface{})
+			err = json.Unmarshal(jsonBytes, &actionsMap)
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("error unmarshaling action table: %w", err)
 			}
+			for key, value := range actionsMap {
+				resourceAction := appv1.ResourceAction{Name: key, Disabled: isActionDisabled(value)}
+				if _, exist := availableActionsMap[key]; exist {
+					continue
+				}
+				if emptyResourceActionFromLua(value) {
+					availableActionsMap[key] = resourceAction
+					continue
+				}
+				resourceActionBytes, err := json.Marshal(value)
+				if err != nil {
+					return nil, fmt.Errorf("error marshaling resource action: %w", err)
+				}
 
-			err = json.Unmarshal(resourceActionBytes, &resourceAction)
-			if err != nil {
-				return nil, err
+				err = json.Unmarshal(resourceActionBytes, &resourceAction)
+				if err != nil {
+					return nil, fmt.Errorf("error unmarshaling resource action: %w", err)
+				}
+				availableActionsMap[key] = resourceAction
 			}
-			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())
+	availableActions := make([]appv1.ResourceAction, 0, len(availableActionsMap))
+	for _, action := range availableActionsMap {
+		availableActions = append(availableActions, action)
+	}
+
+	return availableActions, nil
 }
 
 // Actions are enabled by default
@@ -346,22 +360,35 @@ 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
+
+	// Check if there are resource overrides for the given key
 	override, ok := vm.ResourceOverrides[key]
 	if ok && override.Actions != "" {
 		actions, err := override.GetActions()
 		if err != nil {
-			return "", err
+			return nil, err
+		}
+		// Append the action discovery Lua script if built-in actions are to be included
+		if actions.MergeBuiltinActions {
+			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, fmt.Errorf("error while fetching predefined lua scripts: %w", err)
 	}
-	return discoveryScript, nil
+
+	discoveryScripts = append(discoveryScripts, discoveryScript)
+
+	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 TestGetResourceActionsWithBuiltInActionsFlag(t *testing.T) {
+	testObj := StrToUnstructured(objJSON)
+	vm := VM{
+		ResourceOverrides: map[string]appv1.ResourceOverride{
+			"argoproj.io/Rollout": {
+				Actions: string(grpc.MustMarshal(appv1.ResourceActions{
+					ActionDiscoveryLua:  validDiscoveryLua,
+					MergeBuiltinActions: true,
+				})),
+			},
+		},
+	}
+
+	discoveryLua, err := vm.GetResourceActionDiscovery(testObj)
+	require.NoError(t, err)
+	assert.Equal(t, validDiscoveryLua, discoveryLua[0])
 }
 
 const validDiscoveryLua = `
@@ -349,16 +367,25 @@ scale = {name = 'scale', params = scaleParams}
 
 resume = {name = 'resume'}
 
-test = {}
-a = {scale = scale, resume = resume, test = test}
+a = {scale = scale, resume = resume}
+
+return a
+`
+
+const additionalValidDiscoveryLua = `
+scaleParams = { {name = "override", type = "number"} }
+scale = {name = 'scale', params = scaleParams}
+prebuilt = {prebuilt = 'prebuilt', type = 'number'}
+
+a = {scale = scale, prebuilt = prebuilt}
 
 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{
 		{
@@ -369,8 +396,31 @@ func TestExecuteResourceActionDiscovery(t *testing.T) {
 				Name: "replicas",
 				Type: "number",
 			}},
-		}, {
-			Name: "test",
+		},
+	}
+	for _, expectedAction := range expectedActions {
+		assert.Contains(t, actions, expectedAction)
+	}
+}
+
+func TestExecuteResourceActionDiscoveryWithDuplicationActions(t *testing.T) {
+	testObj := StrToUnstructured(objJSON)
+	vm := VM{}
+	actions, err := vm.ExecuteResourceActionDiscovery(testObj, []string{validDiscoveryLua, additionalValidDiscoveryLua})
+	require.NoError(t, err)
+	expectedActions := []appv1.ResourceAction{
+		{
+			Name: "resume",
+		},
+		{
+			Name: "scale",
+			Params: []appv1.ResourceActionParam{{
+				Name: "replicas",
+				Type: "number",
+			}},
+		},
+		{
+			Name: "prebuilt",
 		},
 	}
 	for _, expectedAction := range expectedActions {
@@ -386,7 +436,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 +449,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)
 }