localstack · MEPalma · Mar 18, 2025 · Mar 13, 2025 · Mar 13, 2025 · Mar 17, 2025
diff --git a/...re/localstack/services/stepfunctions/asl/static_analyser/usage_metrics_static_analyser.py b/...re/localstack/services/stepfunctions/asl/static_analyser/usage_metrics_static_analyser.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import Final
 
 import localstack.services.stepfunctions.usage as UsageMetrics
 from localstack.services.stepfunctions.asl.antlr.runtime.ASLParser import ASLParser
@@ -10,47 +13,72 @@
 LOG = logging.getLogger(__name__)
 
 
+class QueryLanguage(str):
+    JSONPath = QueryLanguageMode.JSONPath.name
+    JSONata = QueryLanguageMode.JSONata.name
+    Both = "JSONPath+JSONata"
+
+
 class UsageMetricsStaticAnalyser(StaticAnalyser):
     @staticmethod
-    def process(definition: str) -> "UsageMetricsStaticAnalyser":
+    def process(definition: str) -> UsageMetricsStaticAnalyser:
         analyser = UsageMetricsStaticAnalyser()
         try:
+            # Run the static analyser.
             analyser.analyse(definition=definition)
 
-            if analyser.has_jsonata:
-                UsageMetrics.jsonata_create_counter.increment()
+            # Determine which query language is being used in this state machine.
+            query_modes = analyser.query_language_modes
+            if len(query_modes) == 2:
+                language_used = QueryLanguage.Both
+            elif QueryLanguageMode.JSONata in query_modes:
+                language_used = QueryLanguage.JSONata
             else:
-                UsageMetrics.jsonpath_create_counter.increment()
+                language_used = QueryLanguage.JSONPath
+
+            # Determine is the state machine uses the variables feature.
+            uses_variables = analyser.uses_variables
 
-            if analyser.has_variable_sampling:
-                UsageMetrics.variables_create_counter.increment()
+            # Count.
+            UsageMetrics.language_features_counter.labels(
+                query_language=language_used, uses_variables=uses_variables
+            ).increment()
         except Exception as e:
             LOG.warning(
                 "Failed to record Step Functions metrics from static analysis",
                 exc_info=e,
             )
         return analyser
 
+    query_language_modes: Final[set[QueryLanguageMode]]
+    uses_variables: bool
+
     def __init__(self):
         super().__init__()
-        self.has_jsonata: bool = False
-        self.has_variable_sampling = False
+        self.query_language_modes = set()
+        self.uses_variables = False
 
     def visitQuery_language_decl(self, ctx: ASLParser.Query_language_declContext):
-        if self.has_jsonata:
+        if len(self.query_language_modes) == 2:
+            # Both query language modes have been confirmed to be in use.
             return
-
         query_language_mode_int = ctx.children[-1].getSymbol().type
         query_language_mode = QueryLanguageMode(value=query_language_mode_int)
-        if query_language_mode == QueryLanguageMode.JSONata:
-            self.has_jsonata = True
+        self.query_language_modes.add(query_language_mode)
+
+    def visitState_decl(self, ctx: ASLParser.State_declContext):
+        # If before entering a state, no query language was explicitly enforced, then we know
+        # this is the first state operating under the default mode (JSONPath)
+        if not self.query_language_modes:
+            self.query_language_modes.add(QueryLanguageMode.JSONPath)
+        super().visitState_decl(ctx=ctx)
 
     def visitString_literal(self, ctx: ASLParser.String_literalContext):
         # Prune everything parsed as a string literal.
         return
 
     def visitString_variable_sample(self, ctx: ASLParser.String_variable_sampleContext):
-        self.has_variable_sampling = True
+        self.uses_variables = True
 
     def visitAssign_decl(self, ctx: ASLParser.Assign_declContext):
-        self.has_variable_sampling = True
+        self.uses_variables = True
diff --git a/localstack-core/localstack/services/stepfunctions/usage.py b/localstack-core/localstack/services/stepfunctions/usage.py
@@ -2,19 +2,11 @@
 Usage reporting for StepFunctions service
 """
 
-from localstack.utils.analytics.usage import UsageCounter
-
-# Count of StepFunctions being created with JSONata QueryLanguage
-jsonata_create_counter = UsageCounter("stepfunctions:jsonata:create")
-
-# Count of StepFunctions being created with JSONPath QueryLanguage
-jsonpath_create_counter = UsageCounter("stepfunctions:jsonpath:create")
-
-# Count of StepFunctions being created that use Variable Sampling or the Assign block
-variables_create_counter = UsageCounter("stepfunctions:variables:create")
-
-# Successful invocations (also including expected error cases in line with AWS behaviour)
-invocation_counter = UsageCounter("stepfunctions:invocation")
-
-# Unexpected errors that we do not account for
-error_counter = UsageCounter("stepfunctions:error")
+from localstack.utils.analytics.metrics import Counter
+
+# Initialize a counter to record the usage of language features for each state machine.
+language_features_counter = Counter(
+    namespace="stepfunctions",
+    name="language_features_used",
+    labels=["query_language", "uses_variables"],
+)
diff --git a/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py b/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py
@@ -2,6 +2,7 @@
 
 import pytest
 
+from localstack.services.stepfunctions.asl.component.common.query_language import QueryLanguageMode
 from localstack.services.stepfunctions.asl.static_analyser.usage_metrics_static_analyser import (
     UsageMetricsStaticAnalyser,
 )
@@ -104,22 +105,40 @@ class TestUsageMetricsStaticAnalyser:
     )
     def test_jsonata(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
-        assert analyser.has_jsonata
-        assert not analyser.has_variable_sampling
+        assert not analyser.uses_variables
+        assert QueryLanguageMode.JSONata in analyser.query_language_modes
+
+    @pytest.mark.parametrize(
+        "definition",
+        [
+            BASE_PASS_JSONATA_OVERRIDE,
+            BASE_PASS_JSONATA_OVERRIDE_DEFAULT,
+        ],
+        ids=[
+            "BASE_PASS_JSONATA_OVERRIDE",
+            "BASE_PASS_JSONATA_OVERRIDE_DEFAULT",
+        ],
+    )
+    def test_both_query_languages(self, definition):
+        analyser = UsageMetricsStaticAnalyser.process(definition)
+        assert not analyser.uses_variables
+        assert QueryLanguageMode.JSONata in analyser.query_language_modes
+        assert QueryLanguageMode.JSONPath in analyser.query_language_modes
 
     @pytest.mark.parametrize("definition", [BASE_PASS_JSONPATH], ids=["BASE_PASS_JSONPATH"])
     def test_jsonpath(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
-        assert not analyser.has_jsonata
-        assert not analyser.has_variable_sampling
+        assert QueryLanguageMode.JSONata not in analyser.query_language_modes
+        assert not analyser.uses_variables
 
     @pytest.mark.parametrize(
         "definition", [JSONPATH_TO_JSONATA_DATAFLOW], ids=["JSONPATH_TO_JSONATA_DATAFLOW"]
     )
     def test_jsonata_and_variable_sampling(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
-        assert analyser.has_jsonata
-        assert analyser.has_variable_sampling
+        assert QueryLanguageMode.JSONPath in analyser.query_language_modes
+        assert QueryLanguageMode.JSONata in analyser.query_language_modes
+        assert analyser.uses_variables
 
     @pytest.mark.parametrize(
         "definition",
@@ -134,5 +153,5 @@ def test_jsonata_and_variable_sampling(self, definition):
     )
     def test_jsonpath_and_variable_sampling(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
-        assert not analyser.has_jsonata
-        assert analyser.has_variable_sampling
+        assert QueryLanguageMode.JSONata not in analyser.query_language_modes
+        assert analyser.uses_variables