diff --git a/sqlparse/engine/__init__.py b/sqlparse/engine/__init__.py index 5b6c56da..b74396e4 100644 --- a/sqlparse/engine/__init__.py +++ b/sqlparse/engine/__init__.py @@ -6,6 +6,7 @@ # 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 @@ -13,4 +14,5 @@ 'FilterStack', 'StatementSplitter', 'grouping', + 'set_max_grouping_tokens', ] diff --git a/sqlparse/engine/config.py b/sqlparse/engine/config.py new file mode 100644 index 00000000..a36d5300 --- /dev/null +++ b/sqlparse/engine/config.py @@ -0,0 +1,22 @@ +# +# Copyright (C) 2009-2020 the sqlparse authors and contributors +# +# +# 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 diff --git a/tests/test_grouping_config.py b/tests/test_grouping_config.py new file mode 100644 index 00000000..0b73b569 --- /dev/null +++ b/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)