[Lineage/Usage] Preserve integer ordinals in GROUP/ORDER BY clauses in query masking (#26303)

* [Lineage/Usage] Preserve integer ordinals in GROUP/ORDER BY clauses in query masking

* Add and reorganize the masker tests

* Preserve ordinals in GROUPING SETS masking

* Handle Keyword.Order explicitly in sqlparse masker

* Defensively allowlist SETS/CUBE/ROLLUP in keyword guard and add ROLLUP/CUBE regression tests

Author: mohittilala

ingestion/src/metadata/ingestion/lineage/masker.py

@@ -24,7 +24,7 @@
 from collate_sqllineage.core.parser.sqlparse.analyzer import SqlParseLineageAnalyzer
 from collate_sqllineage.runner import LineageRunner
 from sqlparse.sql import Comparison
-from sqlparse.tokens import Literal, Number, String
+from sqlparse.tokens import Keyword, Literal, Number, String
 
 from metadata.ingestion.lineage.models import Dialect
 from metadata.utils.execution_time_tracker import (
@@ -51,7 +51,16 @@ def mask_literals_with_sqlparse(
     try:
         parsed = parser._parsed_result
 
-        def mask_token(token):
+        def _is_integer_literal(token) -> bool:
+            """Check if a token is an integer literal (positional reference candidate)."""
+            return token.ttype is Number.Integer
+
+        def mask_token(token, in_groupby_orderby=False):
+            """Mask literal tokens, preserving integer ordinal references in
+            GROUP BY and ORDER BY clauses."""
+            # Skip masking integer positional references in GROUP BY / ORDER BY
+            if in_groupby_orderby and _is_integer_literal(token):
+                return
             # Mask all literals: strings, numbers, or other literal values
             if token.ttype in (
                 String,
@@ -62,18 +71,40 @@ def mask_token(token):
             ):
                 token.value = MASK_TOKEN
             elif token.is_group:
-                # Recursively process grouped tokens
-                for t in token.tokens:
-                    mask_token(t)
-
-        # Process all tokens
-        for token in parsed.tokens:
-            if isinstance(token, Comparison):
-                # In comparisons, mask both sides if literals
-                for t in token.tokens:
-                    mask_token(t)
-            else:
-                mask_token(token)
+                # Recursively process grouped tokens with clause context
+                _mask_group_tokens(token.tokens, in_groupby_orderby)
+
+        def _mask_group_tokens(tokens, in_groupby_orderby=False):
+            """Walk sibling tokens, tracking GROUP BY / ORDER BY context."""
+            current_in_groupby_orderby = in_groupby_orderby
+            for token in tokens:
+                # Track when we enter/leave GROUP BY or ORDER BY clauses
+                if token.ttype is Keyword and token.value.upper() in (
+                    "GROUP BY",
+                    "ORDER BY",
+                ):
+                    current_in_groupby_orderby = True
+                elif token.ttype is Keyword.Order:
+                    # ASC, DESC, NULLS FIRST, NULLS LAST — keep context
+                    pass
+                elif token.ttype is Keyword and token.value.upper() not in (
+                    "GROUPING",
+                    "SETS",
+                    "CUBE",
+                    "ROLLUP",
+                ):
+                    # Any other keyword resets the context (HAVING, LIMIT, WHERE, etc.)
+                    current_in_groupby_orderby = False
+
+                if isinstance(token, Comparison):
+                    # In comparisons, always mask literals (e.g. HAVING COUNT(*) > 1)
+                    for t in token.tokens:
+                        mask_token(t, in_groupby_orderby=False)
+                else:
+                    mask_token(token, current_in_groupby_orderby)
+
+        # Process all tokens starting from the top-level statement
+        _mask_group_tokens(parsed.tokens)
 
         # Return the formatted masked query
         return str(parsed)
@@ -105,14 +136,33 @@ def mask_literals_with_sqlfluff(
             )
             return query
 
-        def replace_literals(segment):
-            """Recursively replace literals with placeholders."""
+        def _is_ordinal_context(segment) -> bool:
+            """Check if the segment is a GROUP BY or ORDER BY clause."""
+            return segment.is_type("groupby_clause", "orderby_clause")
+
+        def _is_integer_ordinal(segment) -> bool:
+            """Check if a numeric_literal segment is an integer ordinal reference.
+            Only pure integers (no decimal point) are valid positional references."""
+            return segment.is_type("numeric_literal") and segment.raw.isdigit()
+
+        def replace_literals(segment, in_groupby_orderby=False):
+            """Recursively replace literals with placeholders,
+            preserving integer ordinal references in GROUP BY / ORDER BY."""
+            # Detect GROUP BY / ORDER BY clause context
+            if _is_ordinal_context(segment):
+                in_groupby_orderby = True
+
+            # Skip masking integer positional references in GROUP BY / ORDER BY
+            if in_groupby_orderby and _is_integer_ordinal(segment):
+                return segment.raw
+
             if segment.is_type("literal", "quoted_literal", "numeric_literal"):
                 return MASK_TOKEN
             if segment.segments:
                 # Recursively process sub-segments
                 return "".join(
-                    replace_literals(sub_seg) for sub_seg in segment.segments
+                    replace_literals(sub_seg, in_groupby_orderby)
+                    for sub_seg in segment.segments
                 )
             return segment.raw
 

ingestion/tests/unit/lineage/masker/__init__.py

@@ -0,0 +1,55 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Test helpers for SQL query masking tests.
+"""
+
+from collate_sqllineage.core.parser.sqlfluff.analyzer import SqlFluffLineageAnalyzer
+from collate_sqllineage.core.parser.sqlglot.analyzer import SqlGlotLineageAnalyzer
+from collate_sqllineage.core.parser.sqlparse.analyzer import SqlParseLineageAnalyzer
+from collate_sqllineage.runner import LineageRunner
+
+from metadata.ingestion.lineage.masker import mask_query, masked_query_cache
+
+PARSER_MAP = {
+    "SqlGlot": SqlGlotLineageAnalyzer,
+    "SqlFluff": SqlFluffLineageAnalyzer,
+    "SqlParse": SqlParseLineageAnalyzer,
+}
+
+
+def assert_masked_query(sql: str, masked_query: str, dialect: str, parser_name: str):
+    """
+    Helper function to test query masking with a specific parser and assert the result.
+
+    :param sql: SQL statement to test
+    :param masked_query: Expected masked query
+    :param dialect: SQL dialect to use (for SqlGlot and SqlFluff)
+    :param parser_name: Name of parser being tested (for error messages)
+    """
+    analyzer_class = PARSER_MAP[parser_name]
+
+    parser_prefix = f"[{parser_name}] " if parser_name else ""
+
+    # clear cache before each test
+    masked_query_cache.clear()
+
+    parser = LineageRunner(sql, dialect=dialect, analyzer=analyzer_class)
+    len(parser.source_tables)  # Force parsing
+
+    actual = mask_query(sql, dialect=dialect, parser=parser)
+    expected = masked_query
+
+    assert actual == expected, (
+        f"\n\t{parser_prefix}Expected Masked Query: {expected}"
+        f"\n\t{parser_prefix}Actual Masked Query:   {actual}"
+    )

ingestion/tests/unit/lineage/masker/helpers.py

@@ -10,16 +10,14 @@
 #  limitations under the License.
 
 """
-Query masking tests
+Query masking tests — core masking logic
 
 Tests for masking SQL queries with different parsers (SqlGlot, SqlFluff, SqlParse).
-Since all parsers now use SqlParse for masking, they should produce identical output.
+Covers: parser dispatch, caching, literal types, ordinal preservation edge cases.
 """
 from unittest import TestCase
 
-import pytest
-
-from ingestion.tests.unit.lineage.queries.helpers import assert_masked_query
+from ingestion.tests.unit.lineage.masker.helpers import assert_masked_query
 from metadata.ingestion.lineage.masker import mask_query, masked_query_cache
 from metadata.ingestion.lineage.models import Dialect
 
@@ -290,174 +288,53 @@ def test_masking_when_no_parser_but_required(self):
 
         assert masked_query is None
 
-    # Dialect specific query masking tests
-
-    def test_postgres_typed_literals_nested_subquery(self):
+    def test_group_by_order_by_ordinal_positions_preserved(self):
         """
-        Test masking of Postgres typed literals in nested subqueries.
-        """
-        query_test_cases = [
-            {
-                "query": "SELECT * FROM orders WHERE order_date = DATE '2023-10-01' AND customer_id IN (SELECT id FROM customers WHERE signup_date = TIMESTAMP '2022-01-15 10:30:00');",  # noqa: E501
-                "expected": "SELECT * FROM orders WHERE order_date = DATE ? AND customer_id IN (SELECT id FROM customers WHERE signup_date = TIMESTAMP ?);",  # noqa: E501
-                "dialect": Dialect.POSTGRES.value,
-            }
-        ]
+        Test that integer ordinal positions in GROUP BY and ORDER BY clauses
+        are NOT masked, while other literals (WHERE, HAVING, LIMIT, OFFSET) still are.
 
-        for test_case in query_test_cases:
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlGlot",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlFluff",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlParse",
-            )
-
-    def test_mysql_implicit_typing_functions_limits_column_if(self):
-        """
-        Test masking of MySQL implicit typing functions like IF, LIMIT.
+        This addresses a reported issue where positional references like
+        GROUP BY 1, 2, 3 were being incorrectly masked to GROUP BY ?, ?, ?
         """
         query_test_cases = [
             {
-                "query": "SELECT IF(status = 'active', 1, 0) AS is_active, DATE(created_at) AS created_day FROM accounts WHERE score > 99.5 AND created_at BETWEEN '2024-01-01' AND '2024-12-31' ORDER BY created_at DESC LIMIT 10 OFFSET 5;",  # noqa: E501
-                "expected": "SELECT IF(status = ?, ?, ?) AS is_active, DATE(created_at) AS created_day FROM accounts WHERE score > ? AND created_at BETWEEN ? AND ? ORDER BY created_at DESC LIMIT ? OFFSET ?;",  # noqa: E501
-                "dialect": Dialect.MYSQL.value,
-            }
-        ]
-
-        for test_case in query_test_cases:
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlGlot",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlFluff",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlParse",
-            )
-
-    def test_bigquery_struct_array_unnest(self):
-        """
-        Test masking of BigQuery STRUCTs, ARRAYs, and UNNEST.
-        """
-        query_test_cases = [
+                "query": "SELECT a, b, c FROM t WHERE x > 5 GROUP BY 1, 2, 3 HAVING COUNT(*) > 1 ORDER BY 1 ASC LIMIT 500 OFFSET 0;",  # noqa: E501
+                "expected": "SELECT a, b, c FROM t WHERE x > ? GROUP BY 1, 2, 3 HAVING COUNT(*) > ? ORDER BY 1 ASC LIMIT ? OFFSET ?;",  # noqa: E501
+                "dialect": Dialect.ANSI.value,
+            },
             {
-                "query": "SELECT u.name, u.age, a.city FROM UNNEST([STRUCT('alice' AS name, 25 AS age, [STRUCT('NY' AS city)])]) AS u, UNNEST(u.f2) AS a WHERE u.age > 21 AND a.city = 'NY';",  # noqa: E501
-                "expected": "SELECT u.name, u.age, a.city FROM UNNEST([STRUCT(? AS name, ? AS age, [STRUCT(? AS city)])]) AS u, UNNEST(u.f2) AS a WHERE u.age > ? AND a.city = ?;",  # noqa: E501
-                "dialect": Dialect.BIGQUERY.value,
+                # Mixed column names and positional integers in GROUP BY / ORDER BY
+                "query": "SELECT a, b FROM t GROUP BY a, 2 ORDER BY b, 1 DESC;",
+                "expected": "SELECT a, b FROM t GROUP BY a, 2 ORDER BY b, 1 DESC;",
+                "dialect": Dialect.ANSI.value,
             },
-        ]
-
-        for test_case in query_test_cases:
-            # TODO: Not masking `'NY'` inside STRUCT whereas `'alice'` and `25` are masked, need to validate
-            # assert_masked_query(
-            #     test_case["query"],
-            #     test_case["expected"],
-            #     test_case["dialect"],
-            #     "SqlGlot",
-            # )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlFluff",
-            )
-            # TODO: Not masking `'NY'` inside STRUCT whereas `'alice'` and `25` are masked, need to validate
-            # assert_masked_query(
-            #     test_case["query"],
-            #     test_case["expected"],
-            #     test_case["dialect"],
-            #     "SqlParse",
-            # )
-
-    def test_snowflake_variant_json_casting(self):
-        """
-        Test masking of Snowflake VARIANT/JSON data types and casting.
-        """
-        query_test_cases = [
             {
-                "query": "SELECT data:id AS user_id, data:profile.name AS user_name, data:profile.age::INT AS user_age FROM events WHERE data:profile.age > 30 AND data:profile.status = 'active';",  # noqa: E501
-                "expected": "SELECT data:id AS user_id, data:profile.name AS user_name, data:profile.age::INT AS user_age FROM events WHERE data:profile.age > ? AND data:profile.status = ?;",  # noqa: E501
-                "dialect": Dialect.SNOWFLAKE.value,
-            }
-        ]
-
-        for test_case in query_test_cases:
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlGlot",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlFluff",
-            )
-            assert_masked_query(
-                test_case["query"],
-                test_case["expected"],
-                test_case["dialect"],
-                "SqlParse",
-            )
-
-    @pytest.mark.skip(
-        reason="SqlGlot and SqlFluff do not support DECLARE statement type yet."
-        " Additionally multi-statement handling needs to be evaluated later."
-    )
-    def test_tsql_variables_convert(self):
-        """
-        Test masking of T-SQL variables and CONVERT function.
-        """
-        query_test_cases = [
+                # CTE with GROUP BY positional references (similar to reported Payoneer query)
+                "query": "WITH cte AS (SELECT a FROM t WHERE x = 'val' GROUP BY 1, 2, 3 HAVING COUNT(*) > 1 ORDER BY 1 ASC) SELECT * FROM cte LIMIT 500 OFFSET 0;",  # noqa: E501
+                "expected": "WITH cte AS (SELECT a FROM t WHERE x = ? GROUP BY 1, 2, 3 HAVING COUNT(*) > ? ORDER BY 1 ASC) SELECT * FROM cte LIMIT ? OFFSET ?;",  # noqa: E501
+                "dialect": Dialect.ANSI.value,
+            },
             {
-                "query": "DECLARE @startDate DATETIME = '2024-01-01'; DECLARE @endDate DATETIME = '2024-12-31'; SELECT * FROM events WHERE event_date BETWEEN @startDate AND @endDate;",  # noqa: E501
-                "expected": "DECLARE @startDate DATETIME = '2024-01-01'; DECLARE @endDate DATETIME = '2024-12-31'; SELECT * FROM events WHERE event_date BETWEEN ? AND ?;",  # noqa: E501
-                "dialect": Dialect.TSQL.value,
-            }
+                # BigQuery dialect with GROUP BY positional references
+                "query": "SELECT full_name, COUNT(*) AS rn FROM t WHERE role = 'admin' AND dept IN ('a', 'b') GROUP BY 1, 2, 3, 4, 5, 6 HAVING COUNT(*) > 1 ORDER BY 1 ASC LIMIT 500 OFFSET 0;",  # noqa: E501
+                "expected": "SELECT full_name, COUNT(*) AS rn FROM t WHERE role = ? AND dept IN (?, ?) GROUP BY 1, 2, 3, 4, 5, 6 HAVING COUNT(*) > ? ORDER BY 1 ASC LIMIT ? OFFSET ?;",  # noqa: E501
+                "dialect": Dialect.BIGQUERY.value,
+            },
         ]
 
         for test_case in query_test_cases:
-            # TODO: sqlglot doesn't support analyzing statement type [declare] for
-            # SQL: DECLARE @startDate DATETIME = '2024-01-01';
             assert_masked_query(
                 test_case["query"],
                 test_case["expected"],
                 test_case["dialect"],
                 "SqlGlot",
             )
-            # TODO: sqlfluff doesn't support analyzing statement type [declare] for
-            # SQL: DECLARE @startDate DATETIME = '2024-01-01';
             assert_masked_query(
                 test_case["query"],
                 test_case["expected"],
                 test_case["dialect"],
                 "SqlFluff",
             )
-            # TODO: since our parser is designed to handle one sql statement at a
-            # time, it returns last masked statement only, need to evaluate later
-            # if multi-statement handling is required
             assert_masked_query(
                 test_case["query"],
                 test_case["expected"],