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 +21 -0
- sqlnbfmt-0.1.0/PKG-INFO +16 -0
- sqlnbfmt-0.1.0/README.md +2 -0
- sqlnbfmt-0.1.0/pyproject.toml +34 -0
- sqlnbfmt-0.1.0/setup.cfg +4 -0
- sqlnbfmt-0.1.0/sqlnbfmt/__init__.py +0 -0
- sqlnbfmt-0.1.0/sqlnbfmt/formatter.py +409 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/PKG-INFO +16 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/SOURCES.txt +12 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/dependency_links.txt +1 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/entry_points.txt +2 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/requires.txt +4 -0
- sqlnbfmt-0.1.0/sqlnbfmt.egg-info/top_level.txt +1 -0
- sqlnbfmt-0.1.0/tests/test_formatter.py +406 -0
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.
|
sqlnbfmt-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
sqlnbfmt-0.1.0/README.md
ADDED
|
@@ -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
|
+
]
|
sqlnbfmt-0.1.0/setup.cfg
ADDED
|
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 @@
|
|
|
1
|
+
|
|
@@ -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)
|