Add configurable grouping token limit

sqlparse/engine/__init__.py

@@ -6,11 +6,13 @@
 # the BSD License: https://opensource.org/licenses/BSD-3-Clause
 
 from sqlparse.engine import grouping
+from sqlparse.engine.config import set_max_grouping_tokens
 from sqlparse.engine.filter_stack import FilterStack
 from sqlparse.engine.statement_splitter import StatementSplitter
 
 __all__ = [
     'FilterStack',
     'StatementSplitter',
     'grouping',
+    'set_max_grouping_tokens',
 ]

sqlparse/engine/config.py

@@ -0,0 +1,22 @@
+#
+# Copyright (C) 2009-2020 the sqlparse authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of python-sqlparse and is released under
+# the BSD License: https://opensource.org/licenses/BSD-3-Clause
+
+from sqlparse.engine import grouping
+
+
+def set_max_grouping_tokens(limit: int | None) -> None:
+    """Set the token limit used by grouping operations.
+
+    ``None`` disables the limit and should only be used for trusted input.
+    Positive integer values replace the default limit.
+    """
+    if limit is not None and (
+        isinstance(limit, bool) or not isinstance(limit, int) or limit < 1
+    ):
+        raise ValueError("limit must be a positive integer or None")
+
+    grouping.MAX_GROUPING_TOKENS = limit

tests/test_grouping_config.py

@@ -0,0 +1,38 @@
+import pytest
+import sqlparse
+from sqlparse.engine import grouping, set_max_grouping_tokens
+from sqlparse.exceptions import SQLParseError
+
+
+@pytest.fixture
+def restore_grouping_token_limit():
+    original = grouping.MAX_GROUPING_TOKENS
+    try:
+        yield
+    finally:
+        set_max_grouping_tokens(original)
+
+
+def test_set_max_grouping_tokens_changes_grouping_limit(
+    restore_grouping_token_limit,
+):
+    set_max_grouping_tokens(1)
+
+    with pytest.raises(SQLParseError, match="Maximum number of tokens exceeded"):
+        sqlparse.parse("SELECT 1")
+
+
+def test_set_max_grouping_tokens_can_disable_limit(restore_grouping_token_limit):
+    set_max_grouping_tokens(None)
+
+    statements = sqlparse.parse("SELECT 1")
+
+    assert len(statements) == 1
+
+
+@pytest.mark.parametrize("limit", [0, -1, 1.5, "10", True])
+def test_set_max_grouping_tokens_rejects_invalid_values(
+    restore_grouping_token_limit, limit
+):
+    with pytest.raises(ValueError, match="positive integer or None"):
+        set_max_grouping_tokens(limit)