Skip to content

Commit

Permalink
add string replacment to account for sql variables
Browse files Browse the repository at this point in the history
  • Loading branch information
Camyll committed Jan 22, 2025
1 parent d913008 commit 8803f97
Show file tree
Hide file tree
Showing 2 changed files with 158 additions and 2 deletions.
4 changes: 2 additions & 2 deletions .lintrunner.toml
Original file line number Diff line number Diff line change
Expand Up @@ -339,13 +339,13 @@ init_command = [
is_formatter = true

[[linter]]
code = 'CLICKHOUSE'
code = 'SQLFLUFF'
include_patterns = ['torchci/clickhouse_queries/**/*.sql']
exclude_patterns = [
]
command = [
'python3',
'tools/linter/adapters/clickhouse_sql_linter.py',
'tools/linter/adapters/sqlfluff_linter.py',
'@{{PATHSFILE}}',
]
init_command = [
Expand Down
156 changes: 156 additions & 0 deletions tools/linter/adapters/sqlfluff_linter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
import argparse
import concurrent.futures
import json
import logging
import os
import re
import subprocess
import time
from enum import Enum
from typing import List, NamedTuple, Optional, Pattern


LINTER_CODE = "SQLFLUFF"


class LintSeverity(str, Enum):
ERROR = "error"
WARNING = "warning"
ADVICE = "advice"
DISABLED = "disabled"


class LintMessage(NamedTuple):
path: Optional[str]
line: Optional[int]
char: Optional[int]
code: str
severity: LintSeverity
name: str
original: Optional[str]
replacement: Optional[str]
description: Optional[str]


RESULTS_RE: Pattern[str] = re.compile(
r"""(?mx)
^
(?P<file>.*?):
(?P<line>\d+):
(?P<char>\d+):
\s(?P<message>.*)
\s(?P<code>\[.*\])
$
"""
)


def run_command(
args: List[str],
) -> "subprocess.CompletedProcess[bytes]":
logging.debug("$ %s", " ".join(args))
start_time = time.monotonic()
try:
return subprocess.run(
args,
capture_output=True,
)
finally:
end_time = time.monotonic()
logging.debug("took %dms", (end_time - start_time) * 1000)


def check_file(
filename: str,
) -> List[LintMessage]:
with open(filename, 'r') as f:
original = f.read()
original = original.replace('{', '\'{').replace('}', '}\'')
with open(filename, 'w') as f:
f.write(original)

try:
# proc.run_command(sed -i -e "s/'{/{/g" -e "s/}'/}/g")
proc = run_command(
[
"sqlfluff",
"format",
"--dialect",
"clickhouse",
filename,
]
)
except OSError as err:
return [
LintMessage(
path=None,
line=None,
char=None,
code=LINTER_CODE,
severity=LintSeverity.ERROR,
name="command-failed",
original=None,
replacement=None,
description=(f"Failed due to {err.__class__.__name__}:\n{err}"),
)
]

with open(filename, 'r') as f:
final = f.read()
final = final.replace('\'{', '{').replace('}\'', '}')
with open(filename, 'w') as f:
f.write(final)

lint_message = proc.stdout


return [
LintMessage(
path=filename,
line=None,
char=None,
code=LINTER_CODE,
severity=LintSeverity.WARNING,
name="format",
original=None,
replacement=None,
description=lint_message.decode("utf-8"),
)
]


def main() -> None:
parser = argparse.ArgumentParser(
description=f"sqlfluff format linter for sql queries.",
fromfile_prefix_chars="@",
)
parser.add_argument(
"filenames",
nargs="+",
help="paths to lint",
)

args = parser.parse_args()

with concurrent.futures.ThreadPoolExecutor(
max_workers=os.cpu_count(),
thread_name_prefix="Thread",
) as executor:
futures = {
executor.submit(
check_file,
filename,
): filename
for filename in args.filenames
}
for future in concurrent.futures.as_completed(futures):
try:
for lint_message in future.result():
print(json.dumps(lint_message._asdict()), flush=True)
except Exception:
logging.critical('Failed at "%s".', futures[future])
raise


if __name__ == "__main__":
main()

0 comments on commit 8803f97

Please sign in to comment.