src/sqlfmt/api.py

import asyncio
import codecs
import concurrent.futures
import locale
import sys
from functools import partial
from glob import glob
from itertools import zip_longest
from pathlib import Path, PurePath
from typing import (
    Awaitable,
    Callable,
    Collection,
    Dict,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    TypeVar,
)

from tqdm import tqdm

from sqlfmt.analyzer import Analyzer
from sqlfmt.cache import Cache, check_cache, clear_cache, load_cache, write_cache
from sqlfmt.exception import SqlfmtEquivalenceError, SqlfmtError, SqlfmtUnicodeError
from sqlfmt.formatter import QueryFormatter
from sqlfmt.mode import Mode as Mode
from sqlfmt.query import Query
from sqlfmt.report import STDIN_PATH, Report, SqlFormatResult

T = TypeVar("T")
R = TypeVar("R")


def format_string(source_string: str, mode: Mode) -> str:
    """
    Takes a raw query string and a mode as input, returns the formatted query
    as a string, or raises a SqlfmtError if the string cannot be formatted.

    If mode.fast is False, also performs a safety check to ensure no tokens
    are dropped from the original input.
    """
    analyzer = mode.dialect.initialize_analyzer(line_length=mode.line_length)
    raw_query = analyzer.parse_query(source_string=source_string)
    formatter = QueryFormatter(mode)
    formatted_query = formatter.format(raw_query)
    result = str(formatted_query)

    if not mode.fast and not mode.check and not mode.diff:
        _perform_safety_check(analyzer, raw_query, result)

    return result


def run(
    files: Collection[Path],
    mode: Mode,
    callback: Optional[Callable[[Awaitable[SqlFormatResult]], None]] = None,
) -> Report:
    """
    Runs sqlfmt on all files in Collection of Paths (files), using the specified Mode.

    Modifies sql files in place, by default. Check or diff Mode do not modify files,
    they only create a Report.

    If a callback is provided, will execute the callback after each file is formatted.

    Returns a Report that can be queried or printed.
    """

    if mode.reset_cache:
        clear_cache()
        cache = {}
    else:
        cache = load_cache()

    results = _format_many(files, cache, mode, callback=callback)

    report = Report(results, mode)

    if not (mode.check or mode.diff):
        _update_source_files(results)
    write_cache(cache, results, mode)

    return report


def get_matching_paths(paths: Iterable[Path], mode: Mode) -> Set[Path]:
    """
    Takes a list of paths (files or directories) and a mode as an input, and
    yields paths to individual files that match the input paths (or are contained in
    its directories) and are not excluded by the mode's exclude glob
    """
    include_set = _get_included_paths(paths, mode)
    exclude_set = set()

    if mode.exclude:
        for s in mode.exclude:
            if PurePath(s).is_absolute():
                exclude_set.update([Path(g) for g in glob(s, recursive=True)])
            else:
                exclude_root = mode.exclude_root or Path.cwd()
                try:
                    # Prefer glob.glob() to Path.glob() here, even though the
                    # latter would be more concise, since Path.glob() does not
                    # resolve symlinks properly in Python <3.13
                    exclude_set.update(
                        [Path(g) for g in glob(str(exclude_root / s), recursive=True)]
                    )
                except IndexError:
                    # for some reason Path.glob(".") raises an index error,
                    # although glob.glob(".") returns ["."]
                    pass

    return include_set - exclude_set


def initialize_progress_bar(
    total: int, mode: Mode, force_progress_bar: bool = False
) -> Tuple[tqdm, Callable[[Awaitable[SqlFormatResult]], None]]:
    """
    Return a callable that can be used with api.run to display a progress bar
    that updates after each file is formatted.

    Pass force_progress_bar to enable the progress bar, even on non-TTY
    terminals (this is handy for testing the progress bar).
    """
    if mode.no_progressbar:
        disable = True
    elif force_progress_bar:
        disable = False
    else:
        # will be disabled on non-TTY envs, enabled otherwise
        disable = None
    progress_bar: tqdm = tqdm(
        total=total, leave=False, disable=disable, delay=0.5, unit="file"
    )

    def progress_callback(_: Awaitable[SqlFormatResult]) -> None:
        progress_bar.update()

    return progress_bar, progress_callback


def _get_included_paths(paths: Iterable[Path], mode: Mode) -> Set[Path]:
    """
    Takes a list of absolute paths (files or directories) and a mode as an input, and
    yields paths to individual files that match the input paths (or are contained in
    its directories)
    """
    include_set = set()
    for p in paths:
        if p == STDIN_PATH:
            include_set.add(p)
        elif p.is_file() and str(p).endswith(tuple(mode.SQL_EXTENSIONS)):
            include_set.add(p)
        elif p.is_dir():
            include_set |= _get_included_paths(p.iterdir(), mode)

    return include_set


def _format_many(
    paths: Collection[Path],
    cache: Cache,
    mode: Mode,
    callback: Optional[Callable[[Awaitable[SqlFormatResult]], None]] = None,
) -> List[SqlFormatResult]:
    """
    Runs sqlfmt on all files in a collection of paths, using the specified mode.

    If there are multiple paths and the mode allows it, uses asyncio's implementation
    of multiprocessing. Otherwise, reverts to single-processing behavior

    Returns a list of SqlFormatResults. Does not write formatted Queries back to disk
    """
    results: List[SqlFormatResult] = []
    cache_misses: List[Path] = []
    for path in paths:
        cached = check_cache(cache=cache, p=path)
        if cached:
            results.append(
                SqlFormatResult(
                    source_path=path,
                    source_string="",
                    formatted_string="",
                    encoding="",
                    utf_bom="",
                    from_cache=True,
                )
            )
        else:
            cache_misses.append(path)

    format_func = partial(_format_one, mode=mode)
    if len(cache_misses) > 1 and not mode.single_process:
        results.extend(
            asyncio.run(_multiprocess_map(format_func, paths, callback=callback))
        )
    else:
        results.extend((map(format_func, cache_misses)))

    return results


async def _multiprocess_map(
    func: Callable[[T], R],
    seq: Iterable[T],
    callback: Optional[Callable[[Awaitable[R]], None]] = None,
) -> List[R]:
    """
    Using multiple processes, creates a Future for each application of func to
    an item in seq, then gathers all Futures and returns the result.

    Provides a similar interface to the map() built-in, but executes in multiple
    processes.
    """
    loop = asyncio.get_event_loop()
    with concurrent.futures.ProcessPoolExecutor() as pool:
        tasks = []
        for item in seq:
            future = loop.run_in_executor(pool, func, item)
            if callback:
                future.add_done_callback(callback)
            tasks.append(future)
        results: List[R] = await asyncio.gather(*tasks)
    return results


def _format_one(path: Path, mode: Mode) -> SqlFormatResult:
    """
    Runs format_string on the contents of a single file (found at path). Handles
    potential user errors in formatted code, and returns a SqlfmtResult
    """
    source, encoding, utf_bom = _read_path_or_stdin(path, mode)
    try:
        formatted = format_string(source, mode)
        return SqlFormatResult(
            source_path=path,
            source_string=source,
            formatted_string=formatted,
            encoding=encoding,
            utf_bom=utf_bom,
        )
    except SqlfmtError as e:
        return SqlFormatResult(
            source_path=path,
            source_string=source,
            formatted_string="",
            encoding=encoding,
            utf_bom=utf_bom,
            exception=e,
        )


def _update_source_files(results: Iterable[SqlFormatResult]) -> None:
    """
    Overwrites file contents at result.source_path with result.formatted_string.

    No-ops for unchanged files, results without a source path, and empty files
    """
    for res in results:
        if res.has_changed and res.source_path != STDIN_PATH and res.formatted_string:
            with open(res.source_path, "w", encoding=res.encoding) as f:
                f.write(res.formatted_string)


def _read_path_or_stdin(path: Path, mode: Mode) -> Tuple[str, str, str]:
    """
    If passed a Path, calls open() and read() and returns contents as
    a tuple of strings. The first element is the contents of the file; the
    second element is the encoding used to read the file; the third
    element is either the utf BOM or an empty string.

    If passed Path("-"), calls sys.stdin.read()
    """
    encoding = (
        (
            locale.getpreferredencoding()
            if mode.encoding.lower() == "inherit"
            else mode.encoding
        )
        .lower()
        .replace("-", "_")
    )
    bom_map: Dict[str, List[bytes]] = {
        "utf": [codecs.BOM_UTF8],
        "utf8": [codecs.BOM_UTF8],
        "u8": [codecs.BOM_UTF8],
        "utf16": [codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE],
        "u16": [codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE],
        "utf16le": [codecs.BOM_UTF16_LE],
        "utf16be": [codecs.BOM_UTF16_BE],
        "utf32": [codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE],
        "u32": [codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE],
        "utf32le": [codecs.BOM_UTF32_LE],
        "utf32be": [codecs.BOM_UTF32_BE],
    }
    detected_bom = ""
    if path == STDIN_PATH:
        # todo: customize encoding of stdin
        source = sys.stdin.read()
    else:
        try:
            with open(path, "r", encoding=encoding) as f:
                source = f.read()
            if encoding.startswith("utf") and encoding != "utf_8_sig":
                for b in [
                    bom.decode(encoding)
                    for bom in bom_map.get(encoding.replace("_", ""), [])
                ]:
                    if source.startswith(b):
                        detected_bom = b
                        source = source[len(b) :]
                        break

        except UnicodeDecodeError as e:
            raise SqlfmtUnicodeError(
                f"Error reading file {path}\n"
                f"File could not be decoded using {encoding}. "
                f"Specifically, {repr(e.object)} at position {e.start} failed "
                f"with: {e.reason}.\n"
                "You can specify a different encoding by running sqlfmt "
                "with the --encoding option. Or set --encoding to 'none' to "
                "use the system default encoding. We suggest always using "
                "utf-8 for all files."
            ) from e
    return source, encoding, detected_bom


def _perform_safety_check(analyzer: Analyzer, raw_query: Query, result: str) -> None:
    """
    Raises a SqlfmtEquivalenceError if re-lexing
    the result produces a different set of tokens than
    the original.
    """
    result_query = analyzer.parse_query(source_string=result)
    filtered_raw_tokens = [
        token.type for token in raw_query.tokens if token.type.is_equivalent_in_output
    ]
    filtered_result_tokens = [
        token.type
        for token in result_query.tokens
        if token.type.is_equivalent_in_output
    ]

    try:
        assert filtered_raw_tokens == filtered_result_tokens
    except AssertionError as e:
        raw_len = len(filtered_raw_tokens)
        result_len = len(filtered_result_tokens)
        mismatch_pos = 0
        mismatch_raw = ""
        mismatch_res = ""

        for i, (raw, res) in enumerate(
            zip_longest(filtered_raw_tokens, filtered_result_tokens)
        ):
            if raw is not res:
                mismatch_pos = i
                mismatch_raw = str(raw)
                mismatch_res = str(res)
                break

        raise SqlfmtEquivalenceError(
            "There was a problem formatting your query that "
            "caused the safety check to fail. Please open an "
            f"issue. Raw query was {raw_len} tokens; formatted "
            f"query was {result_len} tokens. First mismatching "
            f"token at position {mismatch_pos}: raw: {mismatch_raw}; "
            f"result: {mismatch_res}."
        ) from e

    raw_comments = [
        comment.body for line in raw_query.lines for comment in line.comments
    ]
    result_comments = [
        comment.body for line in result_query.lines for comment in line.comments
    ]
    stripped_raw = "".join(["".join(c.split()) for c in raw_comments])
    stripped_res = "".join(["".join(c.split()) for c in result_comments])
    try:
        assert stripped_raw == stripped_res
    except AssertionError as e:
        raw_len = len(stripped_raw)
        result_len = len(stripped_res)
        mismatch_pos = 0
        mismatch_raw = ""
        mismatch_res = ""

        for i, (raw, res) in enumerate(zip_longest(stripped_raw, stripped_res)):
            if raw is not res:
                mismatch_pos = i
                mismatch_raw = stripped_raw[i : i + 25]
                mismatch_res = stripped_res[i : i + 25]
                break

        raise SqlfmtEquivalenceError(
            "There was a problem formatting your query that "
            "caused the safety check to fail. Please open an "
            f"issue. Raw query had {raw_len} comment characters; formatted "
            f"query had {result_len} comment characters. First mismatching "
            f"character at position {mismatch_pos}: raw: {mismatch_raw}; "
            f"result: {mismatch_res}."
        ) from e