129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
import json
|
|
import re
|
|
import subprocess
|
|
import os
|
|
|
|
def format_sql(sql_str):
|
|
"""
|
|
Given a single-line SQL string from the test runner,
|
|
formats it with beautiful indentation according to the
|
|
rules seen in the jspg project fixtures.
|
|
"""
|
|
|
|
# 1. First, let's normalize spaces around operators to make splitting easier.
|
|
# We'll use a simple regex tokenizer.
|
|
# The actual SQL doesn't have spaces around =, >, <, etc.
|
|
sql_str = re.sub(r'([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)=([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)', r'\1.\2 = \3.\4', sql_str)
|
|
sql_str = re.sub(r"([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)='([a-zA-Z0-9_]+)'", r"\1.\2 = '\3'", sql_str)
|
|
sql_str = re.sub(r"([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)>([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)", r"\1.\2 > \3.\4", sql_str)
|
|
sql_str = sql_str.replace("AND ", " AND ")
|
|
sql_str = sql_str.replace("WHERE NOT", "WHERE NOT")
|
|
|
|
# We'll just run a basic custom state-machine formatter
|
|
# Let's clean up tokens to preserve spaces.
|
|
|
|
# We will build the string by adding newlines and indentation where appropriate.
|
|
out = []
|
|
indent = 0
|
|
i = 0
|
|
|
|
# A quick helper to match and consume
|
|
def match(prefix):
|
|
if sql_str[i:].startswith(prefix):
|
|
return True
|
|
return False
|
|
|
|
in_build_object = []
|
|
|
|
# Let's just use a simpler replacement strategy for line breaks,
|
|
# then iterate over lines to fix indentation.
|
|
|
|
# Pre-process for line breaks:
|
|
s = sql_str
|
|
|
|
# Break before certain keywords
|
|
s = s.replace("(SELECT COALESCE", "\n(SELECT COALESCE")
|
|
s = s.replace("FROM ", "\nFROM ")
|
|
s = s.replace("JOIN ", "\nJOIN ")
|
|
s = s.replace("WHERE ", "\nWHERE\n ")
|
|
s = s.replace(" AND ", "\n AND ")
|
|
|
|
# Break before keys in jsonb_build_object, but only if they are followed by a subquery
|
|
# We'll do this by matching: ,'key_name',(SELECT
|
|
s = re.sub(r",('([^']+)')\s*,\s*\(SELECT", r",\n\1,\n(SELECT", s)
|
|
|
|
# Also break scalar keys in jsonb_build_object
|
|
s = re.sub(r",('([^']+)')\s*,", r",\n\1, ", s)
|
|
s = s.replace("jsonb_build_object('", "jsonb_build_object(\n'")
|
|
|
|
# CASE statements
|
|
s = s.replace("CASE WHEN", "CASE\nWHEN")
|
|
s = s.replace("THEN(", "THEN\n(")
|
|
s = s.replace("ELSE NULL END", "\nELSE NULL END")
|
|
s = s.replace(" WHEN ", "\nWHEN ")
|
|
|
|
lines = [l.strip() for l in s.split('\n') if l.strip()]
|
|
|
|
# Now we do a pass to compute indentations based on parenthesis matching and keywords.
|
|
formatted_lines = []
|
|
current_indent = 0
|
|
|
|
for idx, line in enumerate(lines):
|
|
# Calculate indent delta before
|
|
close_paren_count = 0
|
|
while line.startswith(')'):
|
|
close_paren_count += 1
|
|
line = line[1:]
|
|
|
|
if close_paren_count > 0:
|
|
current_indent = max(0, current_indent - 2 * close_paren_count)
|
|
# Prepend the closed parens to the line properly if there's text left,
|
|
# or just emit them if it's just parens.
|
|
if line:
|
|
pass # We handle adding them back later
|
|
else:
|
|
formatted_lines.append(" " * current_indent + ")" * close_paren_count)
|
|
continue
|
|
|
|
# Handle specific keywords
|
|
if line.startswith("FROM ") or line.startswith("JOIN ") or line.startswith("WHERE"):
|
|
pass # Keep parent indent
|
|
elif line.startswith("AND "):
|
|
line = " " + line
|
|
elif line.startswith("WHEN "):
|
|
line = " " + line
|
|
elif line.startswith("ELSE "):
|
|
line = " " + line
|
|
|
|
# If it's a key value pair in build_object, we indent
|
|
if line.startswith("'") and "jsonb_build_object" not in line:
|
|
# We add 2 extra spaces for the items inside build_object
|
|
line = " " + line
|
|
|
|
if line.startswith("(SELECT jsonb_build_object"):
|
|
line = " " + line
|
|
|
|
formatted_line = (" " * current_indent) + (")" * close_paren_count) + line
|
|
|
|
# Calculate indent delta after
|
|
open_paren_count = line.count('(') - line.count(')')
|
|
current_indent += max(0, open_paren_count * 2)
|
|
|
|
formatted_lines.append(formatted_line)
|
|
|
|
return formatted_lines
|
|
|
|
def format_sql_regex(sql_str):
|
|
# The actual jspg parser output might be tricky, let's use a simpler heuristic formatting
|
|
# based exactly on the user's provided output format.
|
|
# It requires custom tokenizing because of nested SELECTs.
|
|
|
|
# Let's try to tokenise
|
|
tokens = re.split(r"(\(SELECT COALESCE|\(SELECT jsonb_build_object|FROM|JOIN|WHERE|AND|CASE|WHEN|THEN|ELSE NULL END|\n|,\s*')", sql_str)
|
|
|
|
pass
|
|
|
|
# We will actually just run `cargo test -- --nocapture` to grab the actual SQLs
|
|
# and do some string replacements.
|
|
# Given the complexity, let's build a dedicated node-based formatter in python.
|