diff --git a/sqlparse/__init__.py b/sqlparse/__init__.py index e62d9785..ae778048 100644 --- a/sqlparse/__init__.py +++ b/sqlparse/__init__.py @@ -16,10 +16,34 @@ from sqlparse import tokens from sqlparse import filters from sqlparse import formatter +from sqlparse.engine import grouping as _grouping __version__ = "0.5.6.dev0" -__all__ = ["engine", "filters", "formatter", "sql", "tokens", "cli"] +__all__ = [ + "engine", + "filters", + "formatter", + "sql", + "tokens", + "cli", + "set_max_grouping_tokens", +] + + +def set_max_grouping_tokens(limit: Optional[int]) -> None: + """Set the maximum token count accepted by the grouping stage. + + Pass ``None`` to disable the token-count limit. Positive integers set a + process-wide limit for subsequent parsing and formatting operations. + Disabling or increasing this limit is not recommended for SQL from + untrusted sources. + """ + if limit is not None and ( + isinstance(limit, bool) or not isinstance(limit, int) or limit < 1 + ): + raise ValueError("Grouping token limit must be a positive integer or None") + _grouping.MAX_GROUPING_TOKENS = limit def parse( diff --git a/tests/test_grouping_config.py b/tests/test_grouping_config.py new file mode 100644 index 00000000..aa4bdc83 --- /dev/null +++ b/tests/test_grouping_config.py @@ -0,0 +1,36 @@ +import pytest + +import sqlparse +from sqlparse.engine import grouping +from sqlparse.exceptions import SQLParseError + + +@pytest.fixture(autouse=True) +def restore_grouping_token_limit(): + original = grouping.MAX_GROUPING_TOKENS + try: + yield + finally: + grouping.MAX_GROUPING_TOKENS = original + + +def test_set_max_grouping_tokens_changes_parser_limit(): + sqlparse.set_max_grouping_tokens(1) + + with pytest.raises(SQLParseError, match="Maximum number of tokens exceeded"): + sqlparse.parse("select value from example") + + +def test_set_max_grouping_tokens_none_disables_limit(): + sqlparse.set_max_grouping_tokens(None) + + assert len(sqlparse.parse("select value from example")) == 1 + + +@pytest.mark.parametrize("limit", (0, -1, True, 1.5, "100")) +def test_set_max_grouping_tokens_rejects_invalid_values(limit): + with pytest.raises( + ValueError, + match="Grouping token limit must be a positive integer or None", + ): + sqlparse.set_max_grouping_tokens(limit) diff --git a/tests/test_grouping_config_format.py b/tests/test_grouping_config_format.py new file mode 100644 index 00000000..d33dfc76 --- /dev/null +++ b/tests/test_grouping_config_format.py @@ -0,0 +1,15 @@ +import pytest + +import sqlparse +from sqlparse.engine import grouping +from sqlparse.exceptions import SQLParseError + + +def test_set_max_grouping_tokens_applies_to_formatting(): + original = grouping.MAX_GROUPING_TOKENS + try: + sqlparse.set_max_grouping_tokens(1) + with pytest.raises(SQLParseError, match="Maximum number of tokens exceeded"): + sqlparse.format("select value from example", reindent=True) + finally: + grouping.MAX_GROUPING_TOKENS = original