fix(admin): validate cron expression in launch plan schedule

- Add cron expression validation using robfig/cron parser - Fix test cases to use correct 5-field cron format - Update invalid cron expression tests to expect errors - Remove extra asterisk from cron patterns in tests Signed-off-by: peterxcli <[email protected]>
flyteorg · Nov 2, 2024 · 5a6cdbe · 5a6cdbe
1 parent 553a702
commit 5a6cdbe
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 6 deletions.
diff --git a/flyteadmin/pkg/manager/impl/validation/launch_plan_validator.go b/flyteadmin/pkg/manager/impl/validation/launch_plan_validator.go
@@ -13,6 +13,7 @@ import (
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/flyteorg/flyte/flytepropeller/pkg/compiler/validators"
+	"github.com/robfig/cron/v3"
 )
 
 func ValidateLaunchPlan(ctx context.Context,
@@ -91,6 +92,19 @@ func validateSchedule(request *admin.LaunchPlanCreateRequest, expectedInputs *co
 					"KickoffTimeInputArg must reference a datetime input. [%v] is a [%v]", schedule.GetKickoffTimeInputArg(), param.GetVar().GetType())
 			}
 		}
+
+		// validate cron expression
+		var cronExpression string
+		if schedule.GetCronExpression() != "" {
+			cronExpression = schedule.GetCronExpression()
+		} else if schedule.GetCronSchedule() != nil {
+			cronExpression = schedule.GetCronSchedule().GetSchedule()
+		}
+		if cronExpression != "" {
+			if _, err := cron.ParseStandard(cronExpression); err != nil {
+				return errors.NewFlyteAdminErrorf(codes.InvalidArgument, "Invalid cron expression: %v", err)
+			}
+		}
 	}
 	return nil
 }

diff --git a/flyteadmin/pkg/manager/impl/validation/launch_plan_validator_test.go b/flyteadmin/pkg/manager/impl/validation/launch_plan_validator_test.go
@@ -358,7 +358,7 @@ func TestValidateSchedule_ArgNotFixed(t *testing.T) {
 		},
 	}
 	t.Run("with deprecated cron expression", func(t *testing.T) {
-		request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * * *")
+		request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * *")
 
 		err := validateSchedule(request, inputMap)
 		assert.NotNil(t, err)
@@ -370,15 +370,15 @@ func TestValidateSchedule_ArgNotFixed(t *testing.T) {
 		assert.NotNil(t, err)
 	})
 	t.Run("with cron schedule", func(t *testing.T) {
-		request := testutils.GetLaunchPlanRequestWithCronSchedule("* * * * * *")
+		request := testutils.GetLaunchPlanRequestWithCronSchedule("* * * * *")
 
 		err := validateSchedule(request, inputMap)
 		assert.NotNil(t, err)
 	})
 }
 
 func TestValidateSchedule_KickoffTimeArgDoesNotExist(t *testing.T) {
-	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * * *")
+	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * *")
 	inputMap := &core.ParameterMap{
 		Parameters: map[string]*core.Parameter{},
 	}
@@ -389,7 +389,7 @@ func TestValidateSchedule_KickoffTimeArgDoesNotExist(t *testing.T) {
 }
 
 func TestValidateSchedule_KickoffTimeArgPointsAtWrongType(t *testing.T) {
-	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * * *")
+	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * *")
 	inputMap := &core.ParameterMap{
 		Parameters: map[string]*core.Parameter{
 			foo: {
@@ -409,7 +409,7 @@ func TestValidateSchedule_KickoffTimeArgPointsAtWrongType(t *testing.T) {
 }
 
 func TestValidateSchedule_NoRequired(t *testing.T) {
-	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * * *")
+	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * *")
 	inputMap := &core.ParameterMap{
 		Parameters: map[string]*core.Parameter{
 			foo: {
@@ -428,7 +428,7 @@ func TestValidateSchedule_NoRequired(t *testing.T) {
 }
 
 func TestValidateSchedule_KickoffTimeBound(t *testing.T) {
-	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * * *")
+	request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * *")
 	inputMap := &core.ParameterMap{
 		Parameters: map[string]*core.Parameter{
 			foo: {
@@ -446,3 +446,34 @@ func TestValidateSchedule_KickoffTimeBound(t *testing.T) {
 	err := validateSchedule(request, inputMap)
 	assert.Nil(t, err)
 }
+
+func TestValidateSchedule_InvalidCronExpression(t *testing.T) {
+	inputMap := &core.ParameterMap{
+		Parameters: map[string]*core.Parameter{
+			foo: {
+				Var: &core.Variable{
+					Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_DATETIME}},
+				},
+				Behavior: &core.Parameter_Required{
+					Required: true,
+				},
+			},
+		},
+	}
+
+	t.Run("with unsupported cron special characters on deprecated cron schedule: #", func(t *testing.T) {
+		request := testutils.GetLaunchPlanRequestWithDeprecatedCronSchedule("* * * * MON#1")
+		request.Spec.EntityMetadata.Schedule.KickoffTimeInputArg = foo
+
+		err := validateSchedule(request, inputMap)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("with unsupported cron special characters: #", func(t *testing.T) {
+		request := testutils.GetLaunchPlanRequestWithCronSchedule("* * * * MON#1")
+		request.Spec.EntityMetadata.Schedule.KickoffTimeInputArg = foo
+
+		err := validateSchedule(request, inputMap)
+		assert.NotNil(t, err)
+	})
+}