sqlnbfmt 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
sqlnbfmt-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Qing Ye
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.1
2
+ Name: sqlnbfmt
3
+ Version: 0.1.0
4
+ Summary: A tool to format SQL code in Jupyter notebooks.
5
+ Author-email: Qing <qingye779@gmail.com>
6
+ License: MIT
7
+ Requires-Python: >=3.12
8
+ Description-Content-Type: text/markdown
9
+ License-File: LICENSE
10
+ Requires-Dist: astor>=0.8.1
11
+ Requires-Dist: nbformat>=5.10.4
12
+ Requires-Dist: pyyaml>=6.0.2
13
+ Requires-Dist: sqlglot>=25.27.0
14
+
15
+ # sqlnbfmt
16
+ A SQL formatter for Jupyter Notebooks
@@ -0,0 +1,2 @@
1
+ # sqlnbfmt
2
+ A SQL formatter for Jupyter Notebooks
@@ -0,0 +1,34 @@
1
+ [project]
2
+ name = "sqlnbfmt"
3
+ version = "0.1.0"
4
+ description = "A tool to format SQL code in Jupyter notebooks."
5
+ authors = [
6
+ { name = "Qing", email = "qingye779@gmail.com" }
7
+ ]
8
+ license = { text = "MIT" }
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ dependencies = [
12
+ "astor>=0.8.1",
13
+ "nbformat>=5.10.4",
14
+ "pyyaml>=6.0.2",
15
+ "sqlglot>=25.27.0",
16
+ ]
17
+
18
+ [tool.uv]
19
+ dev-dependencies = [
20
+ "pytest-cov>=5.0.0",
21
+ "pytest>=8.3.3",
22
+ ]
23
+
24
+ [build-system]
25
+ requires = ["setuptools>=61.0", "wheel"]
26
+ build-backend = "setuptools.build_meta"
27
+
28
+ [project.scripts]
29
+ sqlnbfmt = "sqlnbfmt.formatter:main"
30
+
31
+ [tool.pytest.ini_options]
32
+ markers = [
33
+ "integration: mark a test as an integration test."
34
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
File without changes
@@ -0,0 +1,409 @@
1
+ import ast
2
+ import logging
3
+ import re
4
+ import sys
5
+ import uuid
6
+ from dataclasses import dataclass
7
+ from pathlib import Path
8
+ from typing import Set, Optional, Dict, Union, Tuple, List
9
+
10
+ import nbformat
11
+ import yaml
12
+ from sqlglot import parse_one, parse, errors
13
+ from sqlglot.errors import TokenError
14
+
15
+
16
+ @dataclass
17
+ class FormattingConfig:
18
+ """Configuration for SQL formatting."""
19
+
20
+ sql_keywords: Set[str]
21
+ function_names: Set[str]
22
+ sql_decorators: Set[str]
23
+ single_line_threshold: int = 80
24
+ indent_width: int = 4
25
+
26
+
27
+ class SQLFormattingError(Exception):
28
+ """Custom exception for SQL formatting errors."""
29
+ pass
30
+
31
+
32
+ def setup_logging(level: str = "INFO") -> logging.Logger:
33
+ """Sets up logging with the specified level."""
34
+ logger = logging.getLogger('formatter')
35
+ logger.setLevel(level.upper())
36
+ if not logger.handlers:
37
+ handler = logging.StreamHandler(sys.stdout)
38
+ formatter = logging.Formatter("[%(levelname)s] %(message)s")
39
+ handler.setFormatter(formatter)
40
+ logger.addHandler(handler)
41
+ return logger
42
+
43
+
44
+ def load_config(config_path: Union[str, Path] = "config.yaml") -> FormattingConfig:
45
+ """Loads configuration from a YAML file."""
46
+ config_file = Path(config_path)
47
+ if not config_file.is_file():
48
+ raise FileNotFoundError(f"Configuration file '{config_file}' not found.")
49
+
50
+ try:
51
+ with open(config_file) as file:
52
+ config = yaml.safe_load(file)
53
+ return FormattingConfig(
54
+ sql_keywords=set(config.get("sql_keywords", [])),
55
+ function_names=set(config.get("function_names", [])),
56
+ sql_decorators=set(config.get("sql_decorators", [])),
57
+ single_line_threshold=config.get("formatting_options", {}).get(
58
+ "single_line_threshold", 80
59
+ ),
60
+ indent_width=config.get("formatting_options", {}).get(
61
+ "indent_width", 4
62
+ ),
63
+ )
64
+ except yaml.YAMLError as e:
65
+ raise yaml.YAMLError(f"Error parsing the configuration file: {e}")
66
+
67
+
68
+ def format_sql_code(
69
+ sql_code: str,
70
+ dialect: Optional[str],
71
+ config: FormattingConfig,
72
+ placeholders: Optional[Dict[str, str]] = None,
73
+ force_single_line: bool = False,
74
+ ignore_placeholders: bool = False,
75
+ ) -> str:
76
+ """Formats SQL code using sqlglot's native formatting capabilities while preserving placeholders."""
77
+ try:
78
+ logger = logging.getLogger('formatter')
79
+
80
+ if not sql_code.strip():
81
+ return sql_code
82
+
83
+ temp_sql = sql_code
84
+ placeholder_mapping = {}
85
+
86
+ # Handle placeholders for expressions
87
+ if placeholders and not ignore_placeholders:
88
+ for expr_placeholder, expr in placeholders.items():
89
+ marker = f"'__EXPR_PLACEHOLDER_{uuid.uuid4()}__'"
90
+ temp_sql = temp_sql.replace(expr_placeholder, marker)
91
+ placeholder_mapping[marker.strip("'")] = expr_placeholder
92
+ logger.debug(f"Replaced placeholder {expr_placeholder} with marker {marker}")
93
+
94
+ # Automatically replace common placeholders (%s and ?) with unique markers
95
+ if not ignore_placeholders:
96
+ auto_placeholders = {'%s', '?'}
97
+ for placeholder in auto_placeholders:
98
+ unique_id = str(uuid.uuid4())
99
+ marker = f"'__AUTO_PLACEHOLDER_{unique_id}__'"
100
+ temp_sql = temp_sql.replace(placeholder, marker)
101
+ placeholder_mapping[marker.strip("'")] = placeholder
102
+ logger.debug(f"Auto-replaced placeholder {placeholder} with marker {marker}")
103
+
104
+ temp_sql = temp_sql.strip()
105
+
106
+ # Prevent multiple statements
107
+ statements = parse(temp_sql, read=dialect)
108
+ if len(statements) > 1:
109
+ raise SQLFormattingError("Multiple SQL statements are not supported.")
110
+
111
+ # Parse SQL
112
+ parsed = statements[0]
113
+
114
+ # Generate formatted SQL with uppercase keywords and proper quoting
115
+ formatted_sql = parsed.sql(
116
+ pretty=not force_single_line,
117
+ indent=config.indent_width,
118
+ dialect=dialect,
119
+ )
120
+
121
+ # Condense to single line if necessary
122
+ if force_single_line or (
123
+ config.single_line_threshold
124
+ and len(formatted_sql.replace('\n', ' ')) <= config.single_line_threshold
125
+ ):
126
+ formatted_sql = " ".join(formatted_sql.split())
127
+
128
+ # Restore placeholders
129
+ if placeholder_mapping:
130
+ for marker, original_placeholder in placeholder_mapping.items():
131
+ formatted_sql = formatted_sql.replace(f"'{marker}'", original_placeholder)
132
+ logger.debug(f"Restored marker '{marker}' to placeholder {original_placeholder}")
133
+
134
+ # Logging
135
+ logger.debug(f"Original SQL:\n{sql_code}")
136
+ logger.debug(f"Formatted SQL:\n{formatted_sql}")
137
+
138
+ return formatted_sql
139
+
140
+ except errors.ParseError as e:
141
+ logger.warning(f"SQL parsing failed: {e}. Returning original SQL code.")
142
+ return sql_code # Return the original SQL code unmodified
143
+ except TokenError as te:
144
+ logger.warning(f"Tokenization failed: {te}. Returning original SQL code.")
145
+ return sql_code # Return the original SQL code unmodified
146
+ except Exception as e:
147
+ raise SQLFormattingError(f"Unexpected error during SQL formatting: {e}") from e
148
+
149
+
150
+ class SQLStringFormatter(ast.NodeTransformer):
151
+ """AST NodeTransformer that formats SQL strings."""
152
+
153
+ def __init__(
154
+ self, config: FormattingConfig, dialect: Optional[str], logger: logging.Logger
155
+ ):
156
+ super().__init__()
157
+ self.config = config
158
+ self.dialect = dialect
159
+ self.logger = logger
160
+ self.changed = False
161
+
162
+ def is_likely_sql(self, code: str) -> bool:
163
+ """Enhanced SQL detection with better heuristics."""
164
+ if not code or len(code.strip()) < 10:
165
+ return False
166
+
167
+ if re.match(r"^\s*(/|[a-zA-Z]:\\|https?://|<!DOCTYPE|<html)", code.strip()):
168
+ return False
169
+
170
+ upper_code = code.upper()
171
+ keyword_count = sum(
172
+ 1
173
+ for keyword in self.config.sql_keywords
174
+ if re.search(rf"\b{re.escape(keyword)}\b", upper_code)
175
+ )
176
+
177
+ has_sql_pattern = bool(
178
+ re.search(
179
+ r"\bSELECT\b.*\bFROM\b|\bUPDATE\b.*\bSET\b|\bINSERT\b.*\bINTO\b|\bDELETE\b.*\bFROM\b",
180
+ upper_code,
181
+ re.DOTALL,
182
+ )
183
+ )
184
+
185
+ return keyword_count >= 2 or has_sql_pattern
186
+
187
+ def extract_fstring_parts(self, node: ast.JoinedStr) -> Tuple[str, Dict[str, str]]:
188
+ """Extracts parts of an f-string, preserving expressions."""
189
+ sql_parts = []
190
+ placeholders = {}
191
+
192
+ for value in node.values:
193
+ if isinstance(value, ast.Constant):
194
+ sql_parts.append(value.value)
195
+ elif isinstance(value, ast.FormattedValue):
196
+ expr = ast.unparse(value.value).strip()
197
+ placeholder = f"__EXPR_PLACEHOLDER_{uuid.uuid4()}__"
198
+ sql_parts.append(placeholder)
199
+ placeholders[placeholder] = f"{{{expr}}}"
200
+
201
+ return "".join(sql_parts), placeholders
202
+
203
+ def format_sql_node(
204
+ self, node: Union[ast.Constant, ast.JoinedStr], force_single_line: bool = False
205
+ ) -> Optional[ast.AST]:
206
+ """Formats SQL code in AST nodes."""
207
+ try:
208
+ if isinstance(node, ast.Constant) and isinstance(node.value, str):
209
+ if not self.is_likely_sql(node.value):
210
+ return None
211
+ formatted_sql = format_sql_code(
212
+ node.value,
213
+ self.dialect,
214
+ self.config,
215
+ force_single_line=force_single_line,
216
+ )
217
+ if formatted_sql != node.value:
218
+ self.changed = True
219
+ return ast.Constant(value=formatted_sql)
220
+
221
+ elif isinstance(node, ast.JoinedStr):
222
+ sql_str, placeholders = self.extract_fstring_parts(node)
223
+ if not self.is_likely_sql(sql_str):
224
+ return None
225
+ formatted_sql = format_sql_code(
226
+ sql_str,
227
+ self.dialect,
228
+ self.config,
229
+ placeholders=placeholders,
230
+ force_single_line=force_single_line,
231
+ )
232
+ if formatted_sql != sql_str:
233
+ self.changed = True
234
+ # Reconstruct the f-string
235
+ new_values = []
236
+ pattern = re.compile('|'.join(re.escape(k) for k in placeholders.keys()))
237
+ idx = 0
238
+ while idx < len(formatted_sql):
239
+ match = pattern.search(formatted_sql, idx)
240
+ if match:
241
+ if match.start() > idx:
242
+ new_values.append(ast.Constant(value=formatted_sql[idx:match.start()]))
243
+ expr_placeholder = match.group()
244
+ expr_str = placeholders[expr_placeholder]
245
+ expr_ast = ast.parse(expr_str).body[0].value
246
+ new_values.append(ast.FormattedValue(
247
+ value=expr_ast,
248
+ conversion=-1,
249
+ format_spec=None
250
+ ))
251
+ idx = match.end()
252
+ else:
253
+ new_values.append(ast.Constant(value=formatted_sql[idx:]))
254
+ break
255
+ return ast.JoinedStr(values=new_values)
256
+ return None
257
+ except SQLFormattingError as e:
258
+ self.logger.warning(f"SQL formatting skipped: {e}")
259
+ return None
260
+
261
+ def visit_Constant(self, node: ast.Constant) -> ast.AST:
262
+ """Handles string constants."""
263
+ formatted_node = self.format_sql_node(node)
264
+ return formatted_node or node
265
+
266
+ def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.AST:
267
+ """Handles f-strings."""
268
+ formatted_node = self.format_sql_node(node)
269
+ return formatted_node or node
270
+
271
+ def visit_Assign(self, node: ast.Assign) -> ast.AST:
272
+ """Handles assignments."""
273
+ self.generic_visit(node)
274
+ return node
275
+
276
+ def visit_Call(self, node: ast.Call) -> ast.AST:
277
+ """Handles function calls."""
278
+ func_name = self.get_full_func_name(node.func)
279
+ if any(name in func_name for name in self.config.function_names):
280
+ for idx, arg in enumerate(node.args):
281
+ if isinstance(arg, (ast.Constant, ast.JoinedStr)):
282
+ formatted_node = self.format_sql_node(arg)
283
+ if formatted_node:
284
+ node.args[idx] = formatted_node
285
+ for keyword in node.keywords:
286
+ if isinstance(keyword.value, (ast.Constant, ast.JoinedStr)):
287
+ formatted_node = self.format_sql_node(keyword.value)
288
+ if formatted_node:
289
+ keyword.value = formatted_node
290
+ self.generic_visit(node)
291
+ return node
292
+
293
+ @staticmethod
294
+ def get_full_func_name(node: ast.AST) -> str:
295
+ """Gets the full function name from an AST node."""
296
+ parts = []
297
+ while isinstance(node, ast.Attribute):
298
+ parts.append(node.attr)
299
+ node = node.value
300
+ if isinstance(node, ast.Name):
301
+ parts.append(node.id)
302
+ return ".".join(reversed(parts))
303
+
304
+
305
+ def process_notebook(
306
+ notebook_path: Path,
307
+ config: FormattingConfig,
308
+ dialect: Optional[str],
309
+ logger: logging.Logger
310
+ ) -> bool:
311
+ """Processes a Jupyter notebook, formatting SQL code within code cells."""
312
+ try:
313
+ nb = nbformat.read(notebook_path, as_version=4)
314
+ changed = False
315
+ failed = False # Flag to track failures
316
+
317
+ for idx, cell in enumerate(nb.cells):
318
+ if cell.cell_type != "code":
319
+ continue
320
+
321
+ original_source = cell.source
322
+
323
+ # Handle magic SQL cells
324
+ if cell.source.lstrip().startswith("%%sql"):
325
+ stripped_source = cell.source.lstrip()
326
+ lines = stripped_source.split('\n', 1)
327
+ magic_command = lines[0]
328
+ sql_code = lines[1] if len(lines) > 1 else ''
329
+ try:
330
+ formatted_sql = format_sql_code(
331
+ sql_code, dialect, config, ignore_placeholders=True
332
+ )
333
+ cell.source = f"{magic_command}\n{formatted_sql}"
334
+ if cell.source != original_source:
335
+ changed = True
336
+ except SQLFormattingError as e:
337
+ logger.warning(f"Failed to format magic SQL in cell {idx}: {e}")
338
+ failed = True
339
+ continue
340
+
341
+ # Process regular code cells
342
+ try:
343
+ tree = ast.parse(cell.source)
344
+ formatter = SQLStringFormatter(config, dialect, logger)
345
+ new_tree = formatter.visit(tree)
346
+ if formatter.changed:
347
+ cell.source = ast.unparse(new_tree)
348
+ changed = True
349
+ except SyntaxError as e:
350
+ logger.warning(f"Syntax error in cell {idx}: {e}")
351
+ failed = True
352
+ except Exception as e:
353
+ logger.error(f"Error processing cell {idx}: {e}")
354
+ failed = True
355
+
356
+ if changed and not failed:
357
+ nbformat.write(nb, notebook_path)
358
+ logger.info(f"Updated notebook: {notebook_path}")
359
+ return True
360
+ elif failed:
361
+ logger.error(f"Notebook processing failed for: {notebook_path}")
362
+ return False
363
+ else:
364
+ logger.info(f"No changes made to notebook: {notebook_path}")
365
+ return True
366
+
367
+ except Exception as e:
368
+ logger.error(f"Error processing notebook {notebook_path}: {e}")
369
+ return False
370
+
371
+
372
+ def main():
373
+ """Main entry point for the SQL formatter."""
374
+ import argparse
375
+
376
+ parser = argparse.ArgumentParser(description="Format SQL code in Jupyter notebooks")
377
+ parser.add_argument(
378
+ "notebooks", nargs="+", type=Path, help="Notebook paths to process"
379
+ )
380
+ parser.add_argument(
381
+ "--config", type=Path, default="config.yaml", help="Configuration file path"
382
+ )
383
+ parser.add_argument(
384
+ "--dialect", type=str, help="SQL dialect (e.g., mysql, postgres)"
385
+ )
386
+ parser.add_argument(
387
+ "--log-level",
388
+ choices=["DEBUG", "INFO", "WARNING", "ERROR"],
389
+ default="INFO",
390
+ help="Logging level",
391
+ )
392
+
393
+ args = parser.parse_args()
394
+ logger = setup_logging(args.log_level)
395
+
396
+ try:
397
+ config = load_config(args.config)
398
+ all_succeeded = True
399
+
400
+ for notebook in args.notebooks:
401
+ success = process_notebook(notebook, config, args.dialect, logger)
402
+ if not success:
403
+ all_succeeded = False
404
+
405
+ sys.exit(0 if all_succeeded else 1)
406
+
407
+ except Exception as e:
408
+ logger.error(f"Fatal error: {e}")
409
+ sys.exit(1)
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.1
2
+ Name: sqlnbfmt
3
+ Version: 0.1.0
4
+ Summary: A tool to format SQL code in Jupyter notebooks.
5
+ Author-email: Qing <qingye779@gmail.com>
6
+ License: MIT
7
+ Requires-Python: >=3.12
8
+ Description-Content-Type: text/markdown
9
+ License-File: LICENSE
10
+ Requires-Dist: astor>=0.8.1
11
+ Requires-Dist: nbformat>=5.10.4
12
+ Requires-Dist: pyyaml>=6.0.2
13
+ Requires-Dist: sqlglot>=25.27.0
14
+
15
+ # sqlnbfmt
16
+ A SQL formatter for Jupyter Notebooks
@@ -0,0 +1,12 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ sqlnbfmt/__init__.py
5
+ sqlnbfmt/formatter.py
6
+ sqlnbfmt.egg-info/PKG-INFO
7
+ sqlnbfmt.egg-info/SOURCES.txt
8
+ sqlnbfmt.egg-info/dependency_links.txt
9
+ sqlnbfmt.egg-info/entry_points.txt
10
+ sqlnbfmt.egg-info/requires.txt
11
+ sqlnbfmt.egg-info/top_level.txt
12
+ tests/test_formatter.py
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ sqlnbfmt = sqlnbfmt.formatter:main
@@ -0,0 +1,4 @@
1
+ astor>=0.8.1
2
+ nbformat>=5.10.4
3
+ pyyaml>=6.0.2
4
+ sqlglot>=25.27.0
@@ -0,0 +1 @@
1
+ sqlnbfmt
@@ -0,0 +1,406 @@
1
+ # tests/test_formatter.py
2
+
3
+ import ast
4
+ import pytest
5
+ import re
6
+ import nbformat
7
+ import logging
8
+ from pathlib import Path
9
+ from typing import Optional
10
+
11
+ from sqlnbfmt.formatter import (
12
+ format_sql_code,
13
+ process_notebook,
14
+ FormattingConfig,
15
+ SQLFormattingError
16
+ )
17
+ from sqlglot import errors
18
+ from sqlglot.errors import TokenError
19
+
20
+ # Helper function to normalize SQL for comparison
21
+ def normalize_sql(sql: str) -> str:
22
+ sql = re.sub(r'\s+', ' ', sql) # Replace multiple spaces with a single space
23
+ sql = re.sub(r'\(\s+', '(', sql) # Remove space after '('
24
+ sql = re.sub(r'\s+\)', ')', sql) # Remove space before ')'
25
+ sql = re.sub(r'\s*,\s*', ', ', sql) # Normalize commas and surrounding spaces
26
+ return sql.strip().upper()
27
+
28
+ # Setup logging for tests
29
+ def setup_logging(level: str) -> logging.Logger:
30
+ logger = logging.getLogger('formatter')
31
+ logger.setLevel(level.upper())
32
+ if not logger.handlers:
33
+ handler = logging.StreamHandler()
34
+ formatter = logging.Formatter('%(levelname)s:%(name)s:%(message)s')
35
+ handler.setFormatter(formatter)
36
+ logger.addHandler(handler)
37
+ return logger
38
+
39
+ # Fixtures
40
+ @pytest.fixture
41
+ def sample_config():
42
+ return FormattingConfig(
43
+ sql_keywords={'SELECT', 'ORDER BY', 'JOIN', 'GROUP BY', 'FROM', 'WHERE'},
44
+ function_names={'execute_sql', 'run_query'},
45
+ sql_decorators={'@sql'},
46
+ single_line_threshold=80,
47
+ indent_width=4
48
+ )
49
+
50
+ # Test 1: Basic Reformatting
51
+ def test_format_sql_code_basic_reformatting(sample_config):
52
+ """
53
+ Test that basic SQL queries are reformatted with proper indentation and uppercase keywords.
54
+ """
55
+ sql = """
56
+ select col1, col2 from table where condition order by col1 desc
57
+ """
58
+ formatted = format_sql_code(sql, None, sample_config)
59
+ expected = "SELECT col1, col2 FROM table WHERE condition ORDER BY col1 DESC"
60
+ assert normalize_sql(expected) == normalize_sql(formatted)
61
+
62
+ # Test 2: Malicious Input Handling
63
+ def test_format_sql_code_malicious_input(sample_config):
64
+ """
65
+ Test that the formatter handles malicious SQL inputs without executing them.
66
+ Specifically, it should raise an error for multiple SQL statements.
67
+ """
68
+ sql = "SELECT * FROM users; DROP TABLE users;"
69
+ # The formatter should raise an error for multiple statements
70
+ with pytest.raises(SQLFormattingError):
71
+ format_sql_code(sql, None, sample_config)
72
+
73
+ # Test 3: Escape Characters in Strings
74
+ def test_format_sql_code_escape_characters(sample_config):
75
+ """
76
+ Test that escape characters in SQL strings are handled correctly.
77
+ """
78
+ # Correct SQL escaping uses two single quotes to escape
79
+ sql = "SELECT * FROM table WHERE name = 'O''Connor'"
80
+ formatted = format_sql_code(sql, None, sample_config)
81
+ expected = "SELECT * FROM table WHERE name = 'O''Connor'"
82
+ assert normalize_sql(expected) == normalize_sql(formatted)
83
+
84
+ # Test 4: Special Characters in Identifiers
85
+ def test_format_sql_code_special_characters(sample_config):
86
+ """
87
+ Test that SQL identifiers containing special characters are properly quoted.
88
+ """
89
+ sql = """
90
+ SELECT @column1, `#column2`, $column3
91
+ FROM table
92
+ WHERE condition = '@value'
93
+ """
94
+ # In MySQL, '#' denotes a comment, so '`#column2`' should be quoted to avoid being treated as a comment
95
+ formatted = format_sql_code(sql, "mysql", sample_config)
96
+ expected = "SELECT @column1, `#column2`, $column3 FROM `table` WHERE `condition` = '@value'"
97
+ assert normalize_sql(expected) == normalize_sql(formatted)
98
+
99
+ # Test 5: Reserved Keywords as Identifiers
100
+ def test_format_sql_code_reserved_keywords(sample_config):
101
+ """
102
+ Test that SQL reserved keywords are correctly quoted when used as identifiers.
103
+ """
104
+ sql = "SELECT `group`, `order` FROM `table` WHERE `select` > 10"
105
+ # Reserved keywords should be quoted to avoid parsing errors
106
+ formatted = format_sql_code(sql, "mysql", sample_config)
107
+ expected = "SELECT `group`, `order` FROM `table` WHERE `select` > 10"
108
+ assert normalize_sql(expected) == normalize_sql(formatted)
109
+
110
+ # Test 6: Process Notebook with Various SQL Scenarios
111
+ def test_process_notebook(tmp_path, sample_config):
112
+ """
113
+ Test processing a notebook containing a single SQL code cell.
114
+ """
115
+ # Create a sample notebook
116
+ nb = nbformat.v4.new_notebook()
117
+ nb.cells.append(nbformat.v4.new_code_cell("execute_sql('SELECT * FROM table')"))
118
+ notebook_path = tmp_path / "sample_notebook.ipynb"
119
+ nbformat.write(nb, notebook_path)
120
+
121
+ logger = setup_logging("DEBUG")
122
+ result = process_notebook(notebook_path, sample_config, None, logger)
123
+ assert isinstance(result, bool)
124
+ assert result is True # Assuming processing succeeds
125
+
126
+ # Read the processed notebook and verify changes
127
+ nb_processed = nbformat.read(notebook_path, as_version=4)
128
+ code_cells = [cell for cell in nb_processed.cells if cell.cell_type == "code"]
129
+
130
+ # Verify first code cell
131
+ first_cell = code_cells[0]
132
+ expected = "execute_sql('SELECT * FROM table')"
133
+ assert normalize_sql(expected) == normalize_sql(first_cell.source.strip())
134
+
135
+ # Test 7: Process Real-World Notebook with Complex SQL Queries
136
+ def test_process_real_world_notebook(tmp_path, sample_config):
137
+ """
138
+ Integration test with a real-world-like notebook containing complex SQL queries.
139
+ """
140
+ # Create a more complex notebook
141
+ nb = nbformat.v4.new_notebook()
142
+ nb.cells.extend([
143
+ nbformat.v4.new_markdown_cell("# Real-World SQL Formatting Test"),
144
+ nbformat.v4.new_code_cell('''
145
+ execute_sql("""
146
+ SELECT a.col1, b.col2
147
+ FROM table_a a
148
+ JOIN table_b b ON a.id = b.a_id
149
+ WHERE a.status = 'active'
150
+ ORDER BY a.created_at DESC
151
+ """)
152
+ '''),
153
+ nbformat.v4.new_code_cell('''
154
+ %%sql
155
+ SELECT
156
+ user_id,
157
+ COUNT(*) AS login_count
158
+ FROM logins
159
+ WHERE login_date > '2023-01-01'
160
+ GROUP BY user_id
161
+ ORDER BY login_count DESC
162
+ '''),
163
+ ])
164
+ notebook_path = tmp_path / "real_world_notebook.ipynb"
165
+ nbformat.write(nb, notebook_path)
166
+
167
+ logger = setup_logging("DEBUG")
168
+ result = process_notebook(notebook_path, sample_config, None, logger)
169
+ assert isinstance(result, bool)
170
+ assert result is True # Assuming processing succeeds
171
+
172
+ # Read the processed notebook and verify changes
173
+ nb_processed = nbformat.read(notebook_path, as_version=4)
174
+ code_cells = [cell for cell in nb_processed.cells if cell.cell_type == "code"]
175
+
176
+ # Helper function to extract SQL from function calls using AST
177
+ def extract_sql_from_execute(cell_source: str) -> Optional[str]:
178
+ try:
179
+ tree = ast.parse(cell_source)
180
+ for node in ast.walk(tree):
181
+ if isinstance(node, ast.Call) and getattr(node.func, 'id', '') == 'execute_sql':
182
+ # Assuming the first argument is the SQL code
183
+ if node.args:
184
+ arg = node.args[0]
185
+ if isinstance(arg, (ast.Constant, ast.Constant)):
186
+ return arg.value.strip()
187
+ elif isinstance(arg, ast.JoinedStr):
188
+ # Handle f-strings
189
+ sql_parts = []
190
+ for value in arg.values:
191
+ if isinstance(value, (ast.Constant, ast.Constant)):
192
+ sql_parts.append(value.s if isinstance(value, ast.Constant) else value.value)
193
+ else:
194
+ sql_parts.append("{expression}")
195
+ return ''.join(sql_parts).strip()
196
+ return None
197
+ except Exception as e:
198
+ print(f"Error parsing cell source: {e}")
199
+ return None
200
+
201
+ # Helper function to extract SQL from magic commands
202
+ def extract_sql_from_magic(cell_source: str) -> Optional[str]:
203
+ match = re.match(r'%%sql\s*\n([\s\S]*)', cell_source, re.DOTALL | re.IGNORECASE)
204
+ if match:
205
+ return match.group(1).strip()
206
+ return None
207
+
208
+ # Verify first code cell (execute_sql)
209
+ first_cell = code_cells[0]
210
+ formatted_sql = extract_sql_from_execute(first_cell.source)
211
+ assert formatted_sql is not None, "SQL not found in execute_sql call."
212
+
213
+ expected = (
214
+ "SELECT a.col1, b.col2 FROM table_a AS a "
215
+ "JOIN table_b AS b ON a.id = b.a_id WHERE a.status = 'active' "
216
+ "ORDER BY a.created_at DESC"
217
+ )
218
+ assert normalize_sql(expected) == normalize_sql(formatted_sql)
219
+
220
+ # Verify second code cell (%%sql magic)
221
+ second_cell = code_cells[1]
222
+ formatted_sql_magic = extract_sql_from_magic(second_cell.source)
223
+ expected_magic = (
224
+ "SELECT user_id, COUNT(*) AS login_count FROM logins "
225
+ "WHERE login_date > '2023-01-01' GROUP BY user_id "
226
+ "ORDER BY login_count DESC"
227
+ )
228
+ assert normalize_sql(expected_magic) == normalize_sql(formatted_sql_magic)
229
+
230
+ # Test 9: Logging Messages at Various Levels
231
+ def test_logging_messages_various_levels(sample_config, caplog):
232
+ """
233
+ Test that logging messages are correctly emitted at different log levels.
234
+
235
+ Verifies that DEBUG messages are captured when the logger is set to DEBUG,
236
+ and not captured when set to INFO.
237
+ """
238
+ sql = "SELECT * FROM table"
239
+ logger = setup_logging("DEBUG")
240
+ with caplog.at_level(logging.DEBUG, logger="formatter"):
241
+ format_sql_code(sql, None, sample_config)
242
+ assert "Original SQL:" in caplog.text
243
+ assert "Formatted SQL:" in caplog.text
244
+
245
+ caplog.clear() # Clear previous logs
246
+
247
+ logger = setup_logging("INFO")
248
+ with caplog.at_level(logging.INFO, logger="formatter"):
249
+ format_sql_code(sql, None, sample_config)
250
+ # DEBUG messages should not appear
251
+ assert "Original SQL:" not in caplog.text
252
+ assert "Formatted SQL:" not in caplog.text
253
+
254
+ # Test 10: Subqueries Formatting
255
+ def test_format_sql_code_with_subqueries(sample_config):
256
+ """
257
+ Test formatting SQL queries that contain subqueries.
258
+
259
+ Ensures that nested SELECT statements are formatted correctly.
260
+ """
261
+ sql = """
262
+ SELECT
263
+ a.id,
264
+ (SELECT COUNT(*) FROM orders o WHERE o.user_id = a.id) AS order_count
265
+ FROM users a
266
+ WHERE a.active = TRUE
267
+ ORDER BY order_count DESC
268
+ """
269
+ formatted = format_sql_code(sql, None, sample_config)
270
+ expected_subquery = "(SELECT COUNT(*) FROM orders AS o WHERE o.user_id = a.id) AS order_count"
271
+ assert normalize_sql(expected_subquery) in normalize_sql(formatted)
272
+
273
+ # Test 11: Various Indentation Widths
274
+ def test_format_sql_code_various_indentation(sample_config):
275
+ """
276
+ Test that the formatter applies different indentation widths correctly.
277
+
278
+ Ensures that the indent_width configuration option is respected.
279
+ """
280
+ # Change indent_width to 2
281
+ sample_config.indent_width = 2
282
+ sql = """
283
+ SELECT
284
+ column1,
285
+ column2
286
+ FROM table
287
+ WHERE condition
288
+ ORDER BY column1 DESC
289
+ """
290
+ formatted = format_sql_code(sql, None, sample_config)
291
+ expected = (
292
+ "SELECT\n"
293
+ " column1,\n"
294
+ " column2\n"
295
+ "FROM table\n"
296
+ "WHERE condition\n"
297
+ "ORDER BY column1 DESC"
298
+ )
299
+ assert normalize_sql(expected) == normalize_sql(formatted)
300
+
301
+ # Test 12: Multiple Formatting Options
302
+ def test_format_sql_code_multiple_options(sample_config):
303
+ """
304
+ Test that multiple formatting options work together correctly.
305
+
306
+ Ensures that both indent_width and single_line_threshold are honored.
307
+ """
308
+ sample_config.indent_width = 2
309
+ sample_config.single_line_threshold = 50
310
+ sql = "SELECT col1, col2 FROM table WHERE condition ORDER BY col1 DESC"
311
+ formatted = format_sql_code(sql, None, sample_config)
312
+ # Should be single line due to low threshold
313
+ expected = "SELECT col1, col2 FROM table WHERE condition ORDER BY col1 DESC"
314
+ assert normalize_sql(expected) == normalize_sql(formatted)
315
+
316
+ # Test 13: Common Table Expressions (CTEs) Formatting
317
+ def test_format_sql_code_with_ctes(sample_config):
318
+ """
319
+ Test that SQL queries with Common Table Expressions (CTEs) are formatted correctly.
320
+ """
321
+ sql = """
322
+ WITH recent_orders AS (
323
+ SELECT user_id, COUNT(*) AS order_count
324
+ FROM orders
325
+ WHERE order_date >= DATE_SUB(CURDATE(), INTERVAL 30 DAY)
326
+ GROUP BY user_id
327
+ ), top_customers AS (
328
+ SELECT u.id, u.name, ro.order_count
329
+ FROM users u
330
+ JOIN recent_orders ro ON u.id = ro.user_id
331
+ WHERE ro.order_count > 5
332
+ )
333
+ SELECT tc.id, tc.name, tc.order_count
334
+ FROM top_customers tc
335
+ ORDER BY tc.order_count DESC
336
+ """
337
+ formatted = format_sql_code(sql, "mysql", sample_config)
338
+ expected = (
339
+ "WITH recent_orders AS ("
340
+ "SELECT user_id, COUNT(*) AS order_count FROM orders "
341
+ "WHERE order_date >= DATE_SUB(CURDATE(), INTERVAL '30' DAY) GROUP BY user_id"
342
+ "), top_customers AS ("
343
+ "SELECT u.id, u.name, ro.order_count FROM users AS u "
344
+ "JOIN recent_orders AS ro ON u.id = ro.user_id WHERE ro.order_count > 5"
345
+ ") "
346
+ "SELECT tc.id, tc.name, tc.order_count FROM top_customers AS tc ORDER BY tc.order_count DESC"
347
+ )
348
+ assert normalize_sql(expected) == normalize_sql(formatted)
349
+
350
+ # Test 14: Different SQL Dialects
351
+ def test_format_sql_code_different_dialects(sample_config):
352
+ """
353
+ Test that SQL formatting respects different SQL dialects.
354
+ """
355
+ sql = 'SELECT * FROM "my_table" WHERE "status" = \'active\''
356
+ formatted_pg = format_sql_code(sql, "postgres", sample_config)
357
+ expected_pg = 'SELECT * FROM "my_table" WHERE "status" = \'active\''
358
+ assert normalize_sql(expected_pg) == normalize_sql(formatted_pg)
359
+
360
+ formatted_sqlite = format_sql_code(sql, "sqlite", sample_config)
361
+ expected_sqlite = 'SELECT * FROM "my_table" WHERE "status" = \'active\''
362
+ assert normalize_sql(expected_sqlite) == normalize_sql(formatted_sqlite)
363
+
364
+ # Test 15: Window Functions Formatting
365
+ def test_format_sql_code_with_window_functions(sample_config):
366
+ """
367
+ Test that SQL queries with window functions are formatted correctly.
368
+ """
369
+ sql = """
370
+ SELECT
371
+ user_id,
372
+ login_date,
373
+ COUNT(*) OVER (PARTITION BY user_id ORDER BY login_date) AS cumulative_logins
374
+ FROM logins
375
+ """
376
+ formatted = format_sql_code(sql, "mysql", sample_config)
377
+ expected = (
378
+ "SELECT user_id, login_date, COUNT(*) OVER (PARTITION BY user_id ORDER BY login_date) AS cumulative_logins FROM logins"
379
+ )
380
+ assert normalize_sql(expected) == normalize_sql(formatted)
381
+
382
+ # Test 16: Placeholder Preservation
383
+ def test_format_sql_code_with_placeholders(sample_config):
384
+ """
385
+ Test that SQL queries with placeholders are preserved correctly.
386
+ """
387
+ sql = "SELECT * FROM users WHERE id = %s AND status = ?"
388
+ formatted = format_sql_code(sql, "mysql", sample_config)
389
+ expected = "SELECT * FROM users WHERE id = %s AND status = ?"
390
+ assert normalize_sql(expected) == normalize_sql(formatted)
391
+
392
+ # Test 17: Handling Comments
393
+ def test_format_sql_code_preserve_comments(sample_config):
394
+ """
395
+ Test that SQL comments are preserved in the formatted output.
396
+ """
397
+ sql = """
398
+ -- This is a comment
399
+ SELECT * FROM users -- Inline comment
400
+ """
401
+ formatted = format_sql_code(sql, None, sample_config)
402
+ expected = """
403
+ /* This is a comment */
404
+ SELECT * FROM users /* Inline comment */
405
+ """
406
+ assert normalize_sql(expected) == normalize_sql(formatted)