_check_missing_sqlglot.py

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Mapping


def check_missing_sqlglot(output: Path) -> int:
    txt = _header(_run_qry())
    output.touch()
    return output.write_text(txt, encoding="utf-8")


def _run_qry() -> str:
    import polars as pl
    from pyochain import Dict
    from sqlglot.parsers.duckdb import DuckDBParser

    original = Dict(DuckDBParser.FUNCTIONS)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportUnknownVariableType]
    _check_no_overlap(original)  # pyright: ignore[reportUnknownArgumentType]
    import belugas as bl
    from belugas._sqlglot_patch import DUCKDB_FUNCTIONS  # noqa: PLC2701

    from ._utils import set_pl_config
    from .fn_generator._query import (
        DuckCols,  # pyright: ignore[reportPrivateLocalImportUsage]
        _filters,  # pyright: ignore[reportPrivateUsage]
    )

    set_pl_config()
    function_name = pl.col("function_name")
    alias_of = pl.col("alias_of")
    alias_root = pl.col("alias_root")
    all_aliases = pl.col("all_aliases")
    other_aliases = pl.col("other_aliases").list
    known_function_names = pl.col("known_function_names")
    patched = pl.col("patched")

    def _format(expr: pl.Expr) -> pl.Expr:
        return expr.list.unique().list.sort().list.join(", ").pipe(_emphase)

    def _emphase(expr: pl.Expr) -> pl.Expr:
        return (
            pl.lit("`").add(expr).add("`").str.replace_all("``", "null", literal=True)
        )

    return (
        original
        .iter()
        .collect()
        .into(pl.Series)
        .str.to_uppercase()
        .alias("glot_name")
        .to_frame()
        .lazy()
        .pipe(
            lambda fn_keys: (
                bl.meta
                .functions()
                .collect()
                .lazy()
                .pipe(_filters, DuckCols())
                .select(
                    function_name.str.to_uppercase(),
                    pl
                    .coalesce(alias_of, function_name)
                    .str.to_uppercase()
                    .alias("alias_root"),
                )
                .group_by(alias_root)
                .agg(function_name.unique().sort().alias("all_aliases"))
                .with_columns(all_aliases.alias("function_name"))
                .explode("function_name")
                .join(fn_keys, left_on=function_name, right_on="glot_name", how="anti")
                .join(
                    Dict
                    .from_ref(DUCKDB_FUNCTIONS)
                    .iter()
                    .collect()
                    .into(pl.Series)
                    .str.to_uppercase()
                    .alias("function_name")
                    .to_frame()
                    .lazy()
                    .with_columns(pl.lit(value=True).alias("patched")),
                    on="function_name",
                    how="left",
                )
                .with_columns(patched.fill_null(value=False))
            )
        )
        .drop("alias_root")
        .with_columns(
            all_aliases.list.set_difference(pl.concat_list(function_name)).alias(
                "other_aliases"
            )
        )
        .pipe(
            lambda lf: lf.join(
                lf.select(
                    function_name
                    .unique()
                    .sort()
                    .implode()
                    .alias("known_function_names")
                ),
                how="cross",
            )
        )
        .select(
            function_name.pipe(_emphase).alias("function_name"),
            patched,
            other_aliases
            .set_intersection(known_function_names)
            .pipe(_format)
            .alias("absent_aliases"),
            other_aliases
            .set_difference(known_function_names)
            .pipe(_format)
            .alias("present_aliases"),
        )
        .sort(patched, function_name, descending=[True, False])
        .with_row_index("idx")
        .collect()
        .pipe(repr)
    )


def _header(content: str) -> str:
    return f"""
# Missing SQLGlot Functions

auto-generated by [this script]({Path(__file__).relative_to(Path.cwd()).as_posix()}). Do not edit manually.

## Schema

- `function_name` is the name of the function extracted from the `duckdb_functions` table.
- `patched` is `true` when the function is added by the `belugas` sqlglot patch and absent from the original `sqlglot` mapping.
- `present_aliases` column contains the aliases that are present in the sqlglot `DuckDB` parser `FUNCTIONS` mapping.
- `absent_aliases` column contains the aliases that are absent in the sqlglot `DuckDB` parser `FUNCTIONS` mapping.

## Data

{content}
"""


def _check_no_overlap(original: Mapping[str, object]) -> None:

    from pyochain import Iter, Some

    from belugas._sqlglot_patch import (  # noqa: PLC2701
        _missing_from_glot,  # pyright: ignore[reportPrivateUsage]
    )

    res = (
        Iter(_missing_from_glot()).filter(lambda k: k in original).collect().then_some()
    )
    match res:
        case Some(overlaps):
            msg = f"Some of the functions in _MISSING_FROM_GLOT are already defined in DuckDBParser.FUNCTIONS: {overlaps}"
            raise ValueError(msg)
        case _:
            pass