Fix SQL linting errors

Author: swtwsk

data_pipelines_cli/cli_commands/lint.py

@@ -1,22 +1,26 @@
 import pathlib
 import sys
-from typing import List, Tuple, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, cast
 
 import click
 import sqlfluff
 import yaml
 from sqlfluff.core import SQLFluffUserError
 
 from ..cli_constants import BUILD_DIR
-from ..cli_utils import echo_error, echo_suberror, echo_subinfo, echo_warning
+from ..cli_utils import echo_error, echo_info, echo_suberror, echo_subinfo, echo_warning
 from ..config_generation import read_dictionary_from_config_directory
 from ..errors import SQLLintError
 from ..func_utils import flat_map
 
 if sys.version_info >= (3, 8):
-    from typing import TypedDict  # pylint: disable=no-name-in-module
+    from typing import Protocol, TypedDict  # pylint: disable=no-name-in-module
 else:
-    from typing_extensions import TypedDict
+    from typing_extensions import Protocol, TypedDict
+
+
+T = TypeVar("T", covariant=True)
+S = TypeVar("S")
 
 
 class LintResult(TypedDict):
@@ -26,6 +30,19 @@ class LintResult(TypedDict):
     description: str
 
 
+# Utility class representing sqlfluff `lint` and `fix` types
+class SqlFluffCallable(Protocol[T]):
+    def __call__(
+        self,
+        sql: str,
+        dialect: str = "ansi",
+        rules: Optional[List[str]] = None,
+        exclude_rules: Optional[List[str]] = None,
+        config_path: Optional[str] = None,
+    ) -> T:
+        ...
+
+
 def _get_dialect_or_default() -> str:
     """Read ``dbt.yml`` config file and return its ``target_type`` or just the ``ansi``."""
     env, dbt_filename = "base", "dbt.yml"
@@ -54,8 +71,13 @@ def _get_source_tests_paths() -> List[pathlib.Path]:
     return list(map(lambda dir_name: pathlib.Path.cwd().joinpath(dir_name), dir_names))
 
 
-def _lint_sql_files(dialect: str) -> List[Tuple[pathlib.Path, List[LintResult]]]:
-    lint_results = []
+def _process_sql_files(
+    dialect: str,
+    include_rules: Optional[List[str]],
+    exclude_rules: Optional[List[str]],
+    sqlfluff_fn: SqlFluffCallable[S],
+    result_callback: Callable[[pathlib.Path, S], None],
+) -> None:
     sql_file_paths: List[pathlib.Path] = flat_map(
         lambda dir_path: dir_path.rglob("*.sql"), _get_source_tests_paths()
     )
@@ -64,21 +86,45 @@ def _lint_sql_files(dialect: str) -> List[Tuple[pathlib.Path, List[LintResult]]]
             sql_file_str = sql_file.read()
 
         try:
-            lint_result = sqlfluff.lint(sql_file_str, dialect=dialect)
+            result = sqlfluff_fn(
+                sql_file_str, dialect=dialect, rules=include_rules, exclude_rules=exclude_rules
+            )
         except SQLFluffUserError:  # dialect does not exist, try default instead
             echo_warning(
                 f"Dialect {dialect} did not get recognized. "
                 'Linting using default one ("ansi") instead.'
             )
-            lint_result = sqlfluff.lint(sql_file_str)
+            result = sqlfluff_fn(sql_file_str)
+
+        result_callback(sql_file_path, result)
+
 
+def _fix_sql_files(
+    dialect: str, include_rules: Optional[List[str]], exclude_rules: Optional[List[str]]
+) -> None:
+    def result_callback(sql_file_path: pathlib.Path, fix_result: str) -> None:
+        with open(sql_file_path, "w") as sql_file:
+            sql_file.write(fix_result)
+
+    _process_sql_files(dialect, include_rules, exclude_rules, sqlfluff.fix, result_callback)
+
+
+def _lint_sql_files(
+    dialect: str, include_rules: Optional[List[str]], exclude_rules: Optional[List[str]]
+) -> List[Tuple[pathlib.Path, List[LintResult]]]:
+    lint_results = []
+
+    def result_callback(sql_file_path: pathlib.Path, lint_result: List[Dict[str, Any]]) -> None:
+        nonlocal lint_results
         if len(lint_result) > 0:
             lint_results.append(
                 (
                     sql_file_path.relative_to(pathlib.Path.cwd()),
                     cast(List[LintResult], lint_result),
                 )
             )
+
+    _process_sql_files(dialect, include_rules, exclude_rules, sqlfluff.lint, result_callback)
     return lint_results
 
 
@@ -102,17 +148,58 @@ def _print_lint_results(lint_results: List[Tuple[pathlib.Path, List[LintResult]]
         click.echo("")
 
 
-def lint() -> None:
+def lint(fix: bool, include_rules: Optional[List[str]], exclude_rules: Optional[List[str]]) -> None:
     """
     Lint and format SQL.
+
+    :param fix: Whether to lint and fix linting errors, or just lint.
+    :type fix: bool
+    :param include_rules: A subset of rules to lint with.
+    :type include_rules: Optional[List[str]]
+    :param exclude_rules: A subset of rules not to lint with.
+    :type exclude_rules: Optional[List[str]]
     """
+    echo_info("Linting SQLs:")
     dialect = _get_dialect_or_default()
-    lint_results = _lint_sql_files(dialect)
+    if fix:
+        echo_subinfo("Attempting to fix SQLs. Not every error can be automatically fixed.")
+        _fix_sql_files(dialect, include_rules, exclude_rules)
+    echo_subinfo("Linting SQLs.")
+    lint_results = _lint_sql_files(dialect, include_rules, exclude_rules)
     _print_lint_results(lint_results)
     if len(lint_results) > 0:
         raise SQLLintError(list(map(lambda tup: str(tup[0]), lint_results)))
 
 
-@click.command(name="lint", help="Lint and format SQL")
-def lint_command() -> None:
-    lint()
+@click.command(
+    name="lint",
+    short_help="Lint and format SQL",
+    help="Lint and format SQL using SQLFluff.\n\n"
+    "For more information on rules and the workings of SQLFluff, "
+    "refer to https://docs.sqlfluff.com/",
+)
+@click.option(
+    "--no-fix",
+    is_flag=True,
+    default=False,
+    type=bool,
+    help="Whether to lint and fix linting errors, or just lint.",
+)
+@click.option(
+    "--rules",
+    required=False,
+    type=str,
+    help="A subset of rules to lint with, as string of rules separated by a comma.",
+)
+@click.option(
+    "--exclude-rules",
+    required=False,
+    type=str,
+    help="A subset of rules not to lint with, as string of rules separated by a comma.",
+)
+def lint_command(no_fix: bool, rules: Optional[str], exclude_rules: Optional[str]) -> None:
+    lint(
+        not no_fix,
+        rules.split(",") if rules else None,
+        exclude_rules.split(",") if exclude_rules else None,
+    )

tests/cli_commands/test_lint.py

@@ -63,6 +63,10 @@ def _lint_without_errors(self, sql_content, *_args, **_kwargs):
         self.linted_sqls.append(sql_content)
         return []
 
+    @staticmethod
+    def _fix(sql_content: str, *_args, **_kwargs) -> str:
+        return str(len(sql_content))
+
     @patch("data_pipelines_cli.cli_commands.lint.BUILD_DIR", pathlib.Path("/a/b/c/d/e/f"))
     def test_lint_no_sqls(self):
         with tempfile.TemporaryDirectory() as tmp_dir, patch(
@@ -74,46 +78,71 @@ def test_lint_no_sqls(self):
             )
 
             runner = CliRunner()
-            result = runner.invoke(_cli, ["lint"])
+            result = runner.invoke(_cli, ["lint", "--no-fix"])
             self.assertEqual(
                 0, result.exit_code, msg="\n".join([str(result.exception), str(result.output)])
             )
 
     @patch("data_pipelines_cli.cli_commands.lint.BUILD_DIR", pathlib.Path("/a/b/c/d/e/f"))
     def test_lint_sqls_with_errors(self):
+        fix_mock = MagicMock()
         with patch("pathlib.Path.cwd", lambda: pathlib.Path(self.dbt_project_tmp_dir)), patch(
             "sqlfluff.lint", self._lint_with_errors
-        ):
+        ), patch("data_pipelines_cli.cli_commands.lint._fix_sql_files", fix_mock):
             runner = CliRunner()
-            result = runner.invoke(_cli, ["lint"])
+            result = runner.invoke(_cli, ["lint", "--no-fix"])
             self.assertEqual(
                 1, result.exit_code, msg="\n".join([str(result.exception), str(result.output)])
             )
             self.assertIsInstance(result.exception, SQLLintError)
             self.assertSetEqual(set(map(lambda t: t[1], self.sqls)), set(self.linted_sqls))
             self.assertIn("L:   1 | P:   2 | C:L213 | description\n", result.output)
             self.assertIn("L:   2 | P:   4 | C:L137 | some other description\n", result.output)
+            self.assertFalse(fix_mock.called)
 
     @patch("data_pipelines_cli.cli_commands.lint.BUILD_DIR", pathlib.Path("/a/b/c/d/e/f"))
     def test_lint_sqls_with_errors_raises(self):
         with patch("pathlib.Path.cwd", lambda: pathlib.Path(self.dbt_project_tmp_dir)), patch(
             "sqlfluff.lint", self._lint_with_errors
         ):
             with self.assertRaises(SQLLintError):
-                lint()
+                lint(False, None, None)
 
     @patch("data_pipelines_cli.cli_commands.lint.BUILD_DIR", pathlib.Path("/a/b/c/d/e/f"))
     def test_lint_sqls_without_errors(self):
         with patch("pathlib.Path.cwd", lambda: pathlib.Path(self.dbt_project_tmp_dir)), patch(
             "sqlfluff.lint", self._lint_without_errors
         ):
             runner = CliRunner()
-            result = runner.invoke(_cli, ["lint"])
+            result = runner.invoke(_cli, ["lint", "--no-fix"])
             self.assertEqual(
                 0, result.exit_code, msg="\n".join([str(result.exception), str(result.output)])
             )
             self.assertSetEqual(set(map(lambda t: t[1], self.sqls)), set(self.linted_sqls))
 
+    @patch("data_pipelines_cli.cli_commands.lint.BUILD_DIR", pathlib.Path("/a/b/c/d/e/f"))
+    def test_fix_sqls(self):
+        with patch("pathlib.Path.cwd", lambda: pathlib.Path(self.dbt_project_tmp_dir)), patch(
+            "sqlfluff.lint", lambda _sql_content, *_args, **_kwargs: []
+        ), patch("sqlfluff.fix", self._fix):
+            runner = CliRunner()
+            result = runner.invoke(_cli, ["lint"])
+            self.assertEqual(
+                0, result.exit_code, msg="\n".join([str(result.exception), str(result.output)])
+            )
+
+            for sql_name, sql_content in self.sqls:
+                with open(
+                    pathlib.Path(self.dbt_project_tmp_dir).joinpath("models", sql_name), "r"
+                ) as fixed_file:
+                    fixed_content = fixed_file.read()
+                self.assertEqual(
+                    str(len(sql_content)),
+                    fixed_content,
+                    msg="Contents of an expected value and "
+                    f"the fixed {sql_name} file are different",
+                )
+
 
 class LintHelperFunctionsTestCase(unittest.TestCase):
     def test_get_dialect(self):