Repository: aws-cloudformation/cfn-lint

Issue:
Step Functions JSONata isn't fully supported (E3601)
### CloudFormation Lint Version

1.32.0

### What operating system are you using?

Mac

### Describe the bug

The CF template was constructed from YAML export from Step Functions editor. Below E3601 errors shown:

```
E3601 Only one of ['And', 'BooleanEquals', 'BooleanEqualsPath', 'IsBoolean', 'IsNull', 'IsNumeric', 'IsPresent', 'IsString', 'IsTimestamp', 'Not', 'NumericEquals', 'NumericEqualsPath', 'NumericGreaterThan', 'NumericGreaterThanPath', 'NumericGreaterThanEquals', 'NumericGreaterThanEqualsPath', 'NumericLessThan', 'NumericLessThanPath', 'NumericLessThanEquals', 'NumericLessThanEqualsPath', 'Or', 'StringEquals', 'StringEqualsPath', 'StringGreaterThan', 'StringGreaterThanPath', 'StringGreaterThanEquals', 'StringGreaterThanEqualsPath', 'StringLessThan', 'StringLessThanPath', 'StringLessThanEquals', 'StringLessThanEqualsPath', 'StringMatches', 'TimestampEquals', 'TimestampEqualsPath', 'TimestampGreaterThan', 'TimestampGreaterThanPath', 'TimestampGreaterThanEquals', 'TimestampGreaterThanEqualsPath', 'TimestampLessThan', 'TimestampLessThanPath', 'TimestampLessThanEquals', 'TimestampLessThanEqualsPath'] is a required property
test1.yml:14:17

E3601 Additional properties are not allowed ('Output' was unexpected)
test1.yml:21:13

E3601 Additional properties are not allowed ('Arguments' was unexpected)
test1.yml:27:13
```

### Expected behavior

No lint errors

### Reproduction template

```yaml
AWSTemplateFormatVersion: "2010-09-09"
Resources:
  StateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      RoleArn: arn:aws:iam::123456789123:role/myRole
      Definition:
        Comment: Reading messages from an SQS queue and iteratively processing each message.
        QueryLanguage: JSONata
        StartAt: Read messages from SQS queue
        States:
          Are there messages to process?:
            Choices:
              - Condition: '{% $exists($states.input.Messages) %}'
                Next: Pass
            Default: Finish
            Type: Choice
          Pass:
            Type: Pass
            Next: Finish
            Output: '{% $states.input %}'
            Assign:
              timestamp: '{% $now() %}'
          Finish:
            Type: Succeed
          Read messages from SQS queue:
            Arguments:
              AttributeNames:
                - All
              MaxNumberOfMessages: 10
              QueueUrl: https://sqs.us-east-1.amazonaws.com/123456789123/myQueue
              VisibilityTimeout: 30
              WaitTimeSeconds: 20
            Next: Are there messages to process?
            Resource: arn:aws:states:::aws-sdk:sqs:receiveMessage
            Type: Task
```

Hints from issue comments:
@gliptak: #3907 #4049 

Solution patch:
diff --git a/src/cfnlint/data/schemas/other/step_functions/statemachine.json b/src/cfnlint/data/schemas/other/step_functions/statemachine.json
index f5ff16d0d..e3e2884eb 100644
--- a/src/cfnlint/data/schemas/other/step_functions/statemachine.json
+++ b/src/cfnlint/data/schemas/other/step_functions/statemachine.json
@@ -4,6 +4,63 @@
   "choice": {
    "$ref": "#/definitions/common",
    "additionalProperties": false,
+   "allOf": [
+    {
+     "if": {
+      "properties": {
+       "QueryLanguage": {
+        "const": "JSONata"
+       }
+      },
+      "required": [
+       "QueryLanguage"
+      ]
+     },
+     "then": {
+      "properties": {
+       "Choices": {
+        "items": {
+         "$ref": "#/definitions/choice/definitions/JSONataChoice"
+        },
+        "type": "array"
+       }
+      }
+     }
+    },
+    {
+     "if": {
+      "oneOf": [
+       {
+        "properties": {
+         "QueryLanguage": {
+          "const": "JSONPath"
+         }
+        },
+        "required": [
+         "QueryLanguage"
+        ]
+       },
+       {
+        "not": {
+         "required": [
+          "QueryLanguage"
+         ]
+        }
+       }
+      ]
+     },
+     "then": {
+      "properties": {
+       "Choices": {
+        "items": {
+         "$ref": "#/definitions/choice/definitions/Operator"
+        },
+        "type": "array"
+       }
+      }
+     }
+    }
+   ],
    "definitions": {
     "Assign": {
      "type": "object"
@@ -11,6 +68,22 @@
     "Condition": {
      "type": "string"
     },
+    "JSONataChoice": {
+     "properties": {
+      "Condition": {
+       "type": "string"
+      },
+      "Next": {
+       "pattern": "^.{1,128}$",
+       "type": "string"
+      }
+     },
+     "required": [
+      "Condition",
+      "Next"
+     ],
+     "type": "object"
+    },
     "Operator": {
      "properties": {
       "And": {
@@ -205,9 +278,6 @@
      "type": "object"
     },
     "Choices": {
-     "items": {
-      "$ref": "#/definitions/choice/definitions/Operator"
-     },
      "type": "array"
     },
     "Comment": {
diff --git a/src/cfnlint/rules/resources/stepfunctions/StateMachineDefinition.py b/src/cfnlint/rules/resources/stepfunctions/StateMachineDefinition.py
index 8e69686b6..a80181471 100644
--- a/src/cfnlint/rules/resources/stepfunctions/StateMachineDefinition.py
+++ b/src/cfnlint/rules/resources/stepfunctions/StateMachineDefinition.py
@@ -72,6 +72,23 @@ class StateMachineDefinition(CfnLintJsonSchema):
                 },
             },
         ]
+        schema["definitions"]["choice"]["allOf"] = [
+            {
+                "properties": {
+                    "Choices": {
+                        "items": {
+                            "properties": {
+                                "Condition": {"type": "string"},
+                                "Next": {"pattern": "^.{1,128}$", "type": "string"},
+                            },
+                            "additionalProperties": False,
+                            "required": ["Condition", "Next"],
+                            "type": "object",
+                        }
+                    }
+                }
+            }
+        ]
         return schema
 
     def _clean_schema(self, validator: Validator, instance: Any):


Test patch:
diff --git a/test/unit/rules/resources/stepfunctions/test_state_machine_definition.py b/test/unit/rules/resources/stepfunctions/test_state_machine_definition.py
index c9eea19f3..de7b49f13 100644
--- a/test/unit/rules/resources/stepfunctions/test_state_machine_definition.py
+++ b/test/unit/rules/resources/stepfunctions/test_state_machine_definition.py
@@ -757,6 +757,38 @@ def rule():
                                 }
                             ],
                         },
+                        "Choices1": {
+                            "Type": "Choice",
+                            "Choices": [
+                                {
+                                    "Condition": "BatchJobNotification",
+                                    "Next": "Notify Success",
+                                }
+                            ],
+                            "Next": "Notify Success",
+                            "QueryLanguage": "JSONata",
+                        },
+                        "Choices2": {
+                            "Type": "Choice",
+                            "Choices": [
+                                {
+                                    "Condition": "BatchJobNotification",
+                                    "Next": "Notify Success",
+                                }
+                            ],
+                            "Next": "Notify Success",
+                        },
+                        "Choices3": {
+                            "Type": "Choice",
+                            "Choices": [
+                                {
+                                    "And": "BatchJobNotification",
+                                    "Next": "Notify Success",
+                                }
+                            ],
+                            "Next": "Notify Success",
+                            "QueryLanguage": "JSONata",
+                        },
                         "Notify Success": {
                             "Type": "Task",
                             "Resource": "arn:aws:states:::sns:publish",
@@ -857,6 +889,67 @@ def rule():
                         ["Definition", "States", "Submit Batch Job 3", "Arguments"]
                     ),
                 ),
+                ValidationError(
+                    (
+                        "Only one of ['And', 'BooleanEquals', 'BooleanEqualsPath', 'IsBoolean', 'IsNull', "
+                        "'IsNumeric', 'IsPresent', 'IsString', 'IsTimestamp', 'Not', 'NumericEquals', "
+                        "'NumericEqualsPath', 'NumericGreaterThan', 'NumericGreaterThanPath', "
+                        "'NumericGreaterThanEquals', 'NumericGreaterThanEqualsPath', 'NumericLessThan', "
+                        "'NumericLessThanPath', 'NumericLessThanEquals', 'NumericLessThanEqualsPath', 'Or', "
+                        "'StringEquals', 'StringEqualsPath', 'StringGreaterThan', 'StringGreaterThanPath', "
+                        "'StringGreaterThanEquals', 'StringGreaterThanEqualsPath', 'StringLessThan', "
+                        "'StringLessThanPath', 'StringLessThanEquals', 'StringLessThanEqualsPath', "
+                        "'StringMatches', 'TimestampEquals', 'TimestampEqualsPath', 'TimestampGreaterThan', "
+                        "'TimestampGreaterThanPath', 'TimestampGreaterThanEquals', 'TimestampGreaterThanEqualsPath', "
+                        "'TimestampLessThan', 'TimestampLessThanPath', 'TimestampLessThanEquals', "
+                        "'TimestampLessThanEqualsPath'] is a required property"
+                    ),
+                    rule=StateMachineDefinition(),
+                    validator="requiredXor",
+                    schema_path=deque(
+                        [
+                            "properties",
+                            "States",
+                            "patternProperties",
+                            "^.{1,128}$",
+                            "allOf",
+                            0,
+                            "then",
+                            "allOf",
+                            1,
+                            "then",
+                            "properties",
+                            "Choices",
+                            "items",
+                            "requiredXor",
+                        ]
+                    ),
+                    path=deque(["Definition", "States", "Choices2", "Choices", 0]),
+                ),
+                ValidationError(
+                    ("'Condition' is a required property"),
+                    rule=StateMachineDefinition(),
+                    validator="required",
+                    schema_path=deque(
+                        [
+                            "properties",
+                            "States",
+                            "patternProperties",
+                            "^.{1,128}$",
+                            "allOf",
+                            0,
+                            "then",
+                            "allOf",
+                            0,
+                            "then",
+                            "properties",
+                            "Choices",
+                            "items",
+                            "required",
+                        ]
+                    ),
+                    path=deque(["Definition", "States", "Choices3", "Choices", 0]),
+                ),
             ],
         ),
         (
@@ -932,6 +1025,26 @@ def rule():
                             },
                             "End": True,
                         },
+                        "Choices1": {
+                            "Type": "Choice",
+                            "Choices": [
+                                {
+                                    "Condition": "{% $states.input.type != 'Private' %}",
+                                    "Next": "Notify Failure",
+                                },
+                            ],
+                            "End": True,
+                        },
+                        "Choices2": {
+                            "Type": "Choice",
+                            "Choices": [
+                                {
+                                    "And": "{% $states.input.type != 'Private' %}",
+                                    "Next": "Notify Failure",
+                                },
+                            ],
+                            "End": True,
+                        },
                         "Notify Failure": {
                             "Type": "Task",
                             "Resource": "arn:aws:states:::sns:publish",
@@ -1019,6 +1132,54 @@ def rule():
                         ["Definition", "States", "Notify Success", "Parameters"]
                     ),
                 ),
+                ValidationError(
+                    ("'Condition' is a required property"),
+                    rule=StateMachineDefinition(),
+                    validator="required",
+                    schema_path=deque(
+                        [
+                            "properties",
+                            "States",
+                            "patternProperties",
+                            "^.{1,128}$",
+                            "allOf",
+                            0,
+                            "then",
+                            "allOf",
+                            0,
+                            "properties",
+                            "Choices",
+                            "items",
+                            "required",
+                        ]
+                    ),
+                    path=deque(["Definition", "States", "Choices2", "Choices", 0]),
+                ),
+                ValidationError(
+                    ("Additional properties are not allowed ('And' was unexpected)"),
+                    rule=StateMachineDefinition(),
+                    validator="additionalProperties",
+                    schema_path=deque(
+                        [
+                            "properties",
+                            "States",
+                            "patternProperties",
+                            "^.{1,128}$",
+                            "allOf",
+                            0,
+                            "then",
+                            "allOf",
+                            0,
+                            "properties",
+                            "Choices",
+                            "items",
+                            "additionalProperties",
+                        ]
+                    ),
+                    path=deque(
+                        ["Definition", "States", "Choices2", "Choices", 0, "And"]
+                    ),
+                ),
                 ValidationError(
                     (
                         "Additional properties are not allowed ('Parameters' was unexpected)"
@@ -1056,4 +1217,10 @@ def test_validate(
     validator,
 ):
     errs = list(rule.validate(validator, {}, instance, {}))
+    for i, err in enumerate(errs):
+        print(i)
+        print(err.message, err.message == expected[i].message)
+        print(err.path, err.path == expected[i].path)
+        print(err.schema_path, err.schema_path == expected[i].schema_path)
+        print(err.validator, err.validator == expected[i].validator)
     assert errs == expected, f"{name!r} test failed with {errs!r}"