feat: add support for pragma, set, call (#668)

Author: tconbeer

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file.
 ### Formatting Changes and Bug Fixes
 
 - sqlfmt will no longer attempt to format quoted names in blocks of unsupported DDL ([#653](https://github.com/tconbeer/sqlfmt/issues/653) - thank you [@stiabh](https://github.com/stiabh)!)
+- adds support for formatting `pragma`, `set`, and `call` statements([#656](https://github.com/tconbeer/sqlfmt/issues/656) - thank you [@aersam](https://github.com/aersam)!)
 
 ## [0.25.0] - 2025-01-24
 

src/sqlfmt/rules/__init__.py

@@ -9,12 +9,14 @@
     CREATE_CLONABLE,
     CREATE_FUNCTION,
     CREATE_WAREHOUSE,
+    PRAGMA_SET_CALL,
     group,
 )
 from sqlfmt.rules.core import CORE as CORE
 from sqlfmt.rules.function import FUNCTION as FUNCTION
 from sqlfmt.rules.grant import GRANT as GRANT
 from sqlfmt.rules.jinja import JINJA as JINJA  # noqa
+from sqlfmt.rules.pragma import PRAGMA as PRAGMA
 from sqlfmt.rules.unsupported import UNSUPPORTED as UNSUPPORTED
 from sqlfmt.rules.warehouse import WAREHOUSE as WAREHOUSE
 from sqlfmt.tokens import TokenType
@@ -262,6 +264,15 @@
             ),
         ),
     ),
+    Rule(
+        name="pragma",
+        priority=2005,
+        pattern=group(PRAGMA_SET_CALL) + group(r"\W", r"$"),
+        action=partial(
+            actions.handle_nonreserved_top_level_keyword,
+            action=partial(actions.lex_ruleset, new_ruleset=PRAGMA),
+        ),
+    ),
     Rule(
         name="grant",
         priority=2010,

src/sqlfmt/rules/common.py

@@ -53,3 +53,5 @@ def group(*choices: str) -> str:
     )
     + r"(\s+if\s+not\s+exists)?"
 )
+
+PRAGMA_SET_CALL = group(r"pragma", r"set", r"call")

src/sqlfmt/rules/pragma.py

@@ -0,0 +1,27 @@
+from functools import partial
+
+from sqlfmt import actions
+from sqlfmt.rule import Rule
+from sqlfmt.rules.common import PRAGMA_SET_CALL, group
+from sqlfmt.rules.core import CORE
+from sqlfmt.tokens import TokenType
+
+# Covers simple PRAGMA, SET, and CALL statements, e.g.:
+# PRAGMA database_list;
+# PRAGMA storage_info('table_name');
+# CALL pragma_storage_info('table_name');
+# SET default_collation = 'nocase';
+PRAGMA = [
+    *CORE,
+    Rule(
+        name="unterm_keyword",
+        priority=1300,
+        pattern=group(PRAGMA_SET_CALL) + group(r"\W", r"$"),
+        action=partial(
+            actions.handle_reserved_keyword,
+            action=partial(
+                actions.add_node_to_buffer, token_type=TokenType.UNTERM_KEYWORD
+            ),
+        ),
+    ),
+]

tests/data/unformatted/412_pragma.sql

@@ -0,0 +1,48 @@
+-- source: https://duckdb.org/docs/configuration/pragmas.html
+PRAGMA collations;
+SET threads = 4;
+SET default_collation = 'nocase';
+SET default_null_order = 'NULLS_FIRST';
+SET default_null_order = 'NULLS_LAST_ON_ASC_FIRST_ON_DESC';
+PRAGMA default_collation = 'nocase';
+PRAGMA default_null_order = 'NULLS_FIRST';
+PRAGMA default_null_order = 'NULLS_LAST_ON_ASC_FIRST_ON_DESC';
+PRAGMA order_by_non_integer_literal = true;
+PRAGMA version;
+CALL pragma_version();
+PRAGMA enable_progress_bar;
+PRAGMA explain_output = 'physical_only';
+PRAGMA table_info('table_name');
+CALL pragma_table_info('table_name');
+)))))__SQLFMT_OUTPUT__(((((
+-- source: https://duckdb.org/docs/configuration/pragmas.html
+pragma collations
+;
+set threads = 4
+;
+set default_collation = 'nocase'
+;
+set default_null_order = 'NULLS_FIRST'
+;
+set default_null_order = 'NULLS_LAST_ON_ASC_FIRST_ON_DESC'
+;
+pragma default_collation = 'nocase'
+;
+pragma default_null_order = 'NULLS_FIRST'
+;
+pragma default_null_order = 'NULLS_LAST_ON_ASC_FIRST_ON_DESC'
+;
+pragma order_by_non_integer_literal = true
+;
+pragma version
+;
+call pragma_version()
+;
+pragma enable_progress_bar
+;
+pragma explain_output = 'physical_only'
+;
+pragma table_info('table_name')
+;
+call pragma_table_info('table_name')
+;

tests/functional_tests/test_general_formatting.py

@@ -90,6 +90,7 @@
         "unformatted/409_create_external_function.sql",
         "unformatted/410_create_warehouse.sql",
         "unformatted/411_create_clone.sql",
+        "unformatted/412_pragma.sql",
         "unformatted/900_create_view.sql",
         "unformatted/999_unsupported_ddl.sql",
     ],