Repository: TobikoData/sqlmesh
Pull request #3173 (2024-09-24 23:18:37)

Feat: allow formatter to use CAST over :: syntax
Allows users to change the formatter so that it always produces the `CAST(foo AS type)` syntax over `::`. This may come in handy for users that want to easily copy-paste their model queries without having to switch back the `::` manually.

Context: https://tobiko-data.slack.com/archives/C07HJH22U9Z/p1727216482891999


Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index fc4eb8d4..8c0463b3 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/docs/concepts/macros/sqlmesh_macros.md
@@ -923,15 +923,15 @@ select * from raw_data
 would be rendered as:
 
 ```sql linenums="1"
-WITH "raw_data" AS (                                                                           
-  SELECT                                                                                       
-    *                                                                                          
-  FROM "my_table" AS "my_table"                                                                
-  QUALIFY                                                                                      
-    ROW_NUMBER() OVER (PARTITION BY "id", CAST("event_date" AS DATE) ORDER BY "event_date" DESC, "status" ASC) = 1                                                                        
-)                                                                                        
-SELECT                                                                                         
-  *                                                                                            
+WITH "raw_data" AS (
+  SELECT
+    *
+  FROM "my_table" AS "my_table"
+  QUALIFY
+    ROW_NUMBER() OVER (PARTITION BY "id", CAST("event_date" AS DATE) ORDER BY "event_date" DESC, "status" ASC) = 1
+)
+SELECT
+  *
 FROM "raw_data" AS "raw_data"
 ```
 
@@ -957,18 +957,18 @@ SELECT * FROM discount_promotion_dates
 would be rendered as:
 
 ```sql linenums="1"
-WITH "discount_promotion_dates" AS (                                                                             
-  SELECT                                                                                                         
-    "_exploded"."date_day" AS "date_day"                                                                         
-  FROM UNNEST(CAST(GENERATE_SERIES(CAST('2024-01-01' AS DATE), CAST('2024-01-16' AS DATE), INTERVAL '1' DAY) AS  
-DATE[])) AS "_exploded"("date_day")                                                                              
-)                                                                                                                
-SELECT                                                                                                           
-  "discount_promotion_dates"."date_day" AS "date_day"                                                            
+WITH "discount_promotion_dates" AS (
+  SELECT
+    "_exploded"."date_day" AS "date_day"
+  FROM UNNEST(CAST(GENERATE_SERIES(CAST('2024-01-01' AS DATE), CAST('2024-01-16' AS DATE), INTERVAL '1' DAY) AS
+DATE[])) AS "_exploded"("date_day")
+)
+SELECT
+  "discount_promotion_dates"."date_day" AS "date_day"
 FROM "discount_promotion_dates" AS "discount_promotion_dates"
 ```
 
-Note: This is DuckDB SQL and other dialects will be transpiled accordingly. 
+Note: This is DuckDB SQL and other dialects will be transpiled accordingly.
 - Recursive CTEs (common table expressions) will be used for `Redshift / MySQL / MSSQL`.
 - For `MSSQL` in particular, there's a recursion limit of approximately 100. If this becomes a problem, you can add an `OPTION (MAXRECURSION 0)` clause after the date spine macro logic to remove the limit. This applies for long date ranges.
 
@@ -1763,9 +1763,9 @@ SQLMesh supports common Python types for typed macros including:
 - `float`
 - `bool`
 - `SQL` -- When you want the SQL string representation of the argument that's passed in
-- `List[T]` - where `T` is any supported type including sqlglot expressions
-- `Tuple[T]` - where `T` is any supported type including sqlglot expressions
-- `Union[T1, T2, ...]` - where `T1`, `T2`, etc. are any supported types including sqlglot expressions
+- `list[T]` - where `T` is any supported type including sqlglot expressions
+- `tuple[T]` - where `T` is any supported type including sqlglot expressions
+- `T1 | T2 | ...` - where `T1`, `T2`, etc. are any supported types including sqlglot expressions
 
 We also support SQLGlot expressions as type hints, allowing you to ensure inputs are coerced to the desired SQL AST node your intending on working with. Some useful examples include:
 
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index d7a0cc05..bf724df4 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -172,6 +172,8 @@ Options:
   -t, --transpile TEXT        Transpile project models to the specified
                               dialect.
   --append-newline            Include a newline at the end of each file.
+  --no-rewrite-casts          Preserve the existing casts, without rewriting
+                              them to use the :: syntax.
   --normalize                 Whether or not to normalize identifiers to
                               lowercase.
   --pad INTEGER               Determines the pad size in a formatted string.
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index ba4fdb33..12e16151 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -93,6 +93,7 @@ Formatting settings for the `sqlmesh format` command and UI.
 | `leading_comma`       | Whether to use leading commas (Default: False)                                                 | boolean |    N     |
 | `max_text_width`      | The maximum text width in a segment before creating new lines (Default: 80)                    |   int   |    N     |
 | `append_newline`      | Whether to append a newline to the end of the file (Default: False)                            | boolean |    N     |
+| `no_rewrite_casts`    | Preserve the existing casts, without rewriting them to use the :: syntax. (Default: False)                  | boolean |    N     |
 
 ## UI
 
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index b000d1d5..df6b1434 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -441,8 +441,8 @@ options:
 
 #### format
 ```
-%format [--transpile TRANSPILE] [--append-newline] [--normalize]
-              [--pad PAD] [--indent INDENT]
+%format [--transpile TRANSPILE] [--append-newline] [--no-rewrite-casts]
+              [--normalize] [--pad PAD] [--indent INDENT]
               [--normalize-functions NORMALIZE_FUNCTIONS] [--leading-comma]
               [--max-text-width MAX_TEXT_WIDTH] [--check]
 
@@ -453,6 +453,8 @@ options:
                         Transpile project models to the specified dialect.
   --append-newline      Whether or not to append a newline to the end of the
                         file.
+  --no-rewrite-casts    Preserve the existing casts, without rewriting them
+                        to use the :: syntax.
   --normalize           Whether or not to normalize identifiers to lowercase.
   --pad PAD             Determines the pad size in a formatted string.
   --indent INDENT       Determines the indentation size in a formatted string.
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 14bc164b..4daaf566 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -223,6 +223,12 @@ def evaluate(
     help="Include a newline at the end of each file.",
     default=None,
 )
+@click.option(
+    "--no-rewrite-casts",
+    is_flag=True,
+    help="Preserve the existing casts, without rewriting them to use the :: syntax.",
+    default=None,
+)
 @click.option(
     "--normalize",
     is_flag=True,
@@ -266,6 +272,9 @@ def evaluate(
 @cli_analytics
 def format(ctx: click.Context, **kwargs: t.Any) -> None:
     """Format all SQL models and audits."""
+    if kwargs.pop("no_rewrite_casts", None):
+        kwargs["rewrite_casts"] = False
+
     if not ctx.obj.format(**{k: v for k, v in kwargs.items() if v is not None}):
         ctx.exit(1)
 
diff --git a/sqlmesh/core/config/format.py b/sqlmesh/core/config/format.py
index 2b8fd215..8730425d 100644
--- a/sqlmesh/core/config/format.py
+++ b/sqlmesh/core/config/format.py
@@ -16,6 +16,7 @@ class FormatConfig(BaseConfig):
         leading_comma: Whether to use leading commas or not.
         max_text_width: The maximum text width in a segment before creating new lines.
         append_newline: Whether to append a newline to the end of the file or not.
+        no_rewrite_casts: Preserve the existing casts, without rewriting them to use the :: syntax.
     """
 
     normalize: bool = False
@@ -25,6 +26,7 @@ class FormatConfig(BaseConfig):
     leading_comma: bool = False
     max_text_width: int = 80
     append_newline: bool = False
+    no_rewrite_casts: bool = False
 
     @property
     def generator_options(self) -> t.Dict[str, t.Any]:
@@ -33,4 +35,4 @@ class FormatConfig(BaseConfig):
         Returns:
             The generator options.
         """
-        return self.dict(exclude={"append_newline"})
+        return self.dict(exclude={"append_newline", "no_rewrite_casts"})
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 622c5eb4..3fcf9b4b 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -900,6 +900,7 @@ class GenericContext(BaseContext, t.Generic[C]):
     def format(
         self,
         transpile: t.Optional[str] = None,
+        rewrite_casts: t.Optional[bool] = None,
         append_newline: t.Optional[bool] = None,
         *,
         check: t.Optional[bool] = None,
@@ -910,6 +911,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         for target in format_targets.values():
             if target._path is None or target._path.suffix != ".sql":
                 continue
+
             with open(target._path, "r+", encoding="utf-8") as file:
                 before = file.read()
                 expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
@@ -922,13 +924,24 @@ class GenericContext(BaseContext, t.Generic[C]):
                                     value=exp.Literal.string(transpile or target.dialect),
                                 )
                             )
-                format = self.config_for_node(target).format
-                opts = {**format.generator_options, **kwargs}
-                after = format_model_expressions(expressions, transpile or target.dialect, **opts)
+
+                format_config = self.config_for_node(target).format
+                after = format_model_expressions(
+                    expressions,
+                    transpile or target.dialect,
+                    rewrite_casts=(
+                        rewrite_casts
+                        if rewrite_casts is not None
+                        else not format_config.no_rewrite_casts
+                    ),
+                    **{**format_config.generator_options, **kwargs},
+                )
+
                 if append_newline is None:
-                    append_newline = format.append_newline
+                    append_newline = format_config.append_newline
                 if append_newline:
                     after += "\n"
+
                 if not check:
                     file.seek(0)
                     file.write(after)
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index e67d480e..de895b69 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -643,13 +643,17 @@ def _override(klass: t.Type[Tokenizer | Parser], func: t.Callable) -> None:
 
 
 def format_model_expressions(
-    expressions: t.List[exp.Expression], dialect: t.Optional[str] = None, **kwargs: t.Any
+    expressions: t.List[exp.Expression],
+    dialect: t.Optional[str] = None,
+    rewrite_casts: bool = True,
+    **kwargs: t.Any,
 ) -> str:
     """Format a model's expressions into a standardized format.
 
     Args:
         expressions: The model's expressions, must be at least model def + query.
         dialect: The dialect to render the expressions as.
+        rewrite_casts: Whether to rewrite all casts to use the :: syntax.
         **kwargs: Additional keyword arguments to pass to the sql generator.
 
     Returns:
@@ -660,26 +664,28 @@ def format_model_expressions(
 
     *statements, query = expressions
 
-    def cast_to_colon(node: exp.Expression) -> exp.Expression:
-        if isinstance(node, exp.Cast) and not any(
-            # Only convert CAST into :: if it doesn't have additional args set, otherwise this
-            # conversion could alter the semantics (eg. changing SAFE_CAST in BigQuery to CAST)
-            arg
-            for name, arg in node.args.items()
-            if name not in ("this", "to")
-        ):
-            this = node.this
+    if rewrite_casts:
 
-            if not isinstance(this, (exp.Binary, exp.Unary)) or isinstance(this, exp.Paren):
-                cast = DColonCast(this=this, to=node.to)
-                cast.comments = node.comments
-                node = cast
+        def cast_to_colon(node: exp.Expression) -> exp.Expression:
+            if isinstance(node, exp.Cast) and not any(
+                # Only convert CAST into :: if it doesn't have additional args set, otherwise this
+                # conversion could alter the semantics (eg. changing SAFE_CAST in BigQuery to CAST)
+                arg
+                for name, arg in node.args.items()
+                if name not in ("this", "to")
+            ):
+                this = node.this
 
-        exp.replace_children(node, cast_to_colon)
-        return node
+                if not isinstance(this, (exp.Binary, exp.Unary)) or isinstance(this, exp.Paren):
+                    cast = DColonCast(this=this, to=node.to)
+                    cast.comments = node.comments
+                    node = cast
+
+            exp.replace_children(node, cast_to_colon)
+            return node
 
-    query = query.copy()
-    exp.replace_children(query, cast_to_colon)
+        query = query.copy()
+        exp.replace_children(query, cast_to_colon)
 
     return ";\n\n".join(
         [
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index d4971205..ea3b5557 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -72,6 +72,14 @@ class MacroStrTemplate(Template):
 EXPRESSIONS_NAME_MAP = {}
 SQL = t.NewType("SQL", str)
 
+SUPPORTED_TYPES = {
+    "t": t,
+    "typing": t,
+    "List": t.List,
+    "Tuple": t.Tuple,
+    "Union": t.Union,
+}
+
 for klass in sqlglot.Parser.EXPRESSION_PARSERS:
     name = klass if isinstance(klass, str) else klass.__name__  # type: ignore
     EXPRESSIONS_NAME_MAP[name.lower()] = name
@@ -195,7 +203,7 @@ class MacroEvaluator:
             raise MacroEvalError("Error trying to eval macro.") from e
 
         try:
-            annotations = t.get_type_hints(func)
+            annotations = t.get_type_hints(func, localns=SUPPORTED_TYPES)
         except NameError:  # forward references aren't handled
             annotations = {}
 
diff --git a/sqlmesh/core/snapshot/categorizer.py b/sqlmesh/core/snapshot/categorizer.py
index 09d83784..88a1ef37 100644
--- a/sqlmesh/core/snapshot/categorizer.py
+++ b/sqlmesh/core/snapshot/categorizer.py
@@ -8,7 +8,11 @@ from sqlmesh.utils.errors import SQLMeshError
 
 
 def categorize_change(
-    new: Snapshot, old: Snapshot, config: t.Optional[CategorizerConfig] = None
+    new: Snapshot,
+    old: Snapshot,
+    config: t.Optional[CategorizerConfig] = None,
+    is_breaking_change: t.Optional[t.Callable[..., t.Optional[bool]]] = None,
+    **kwargs: t.Any,
 ) -> t.Optional[SnapshotChangeCategory]:
     """Attempts to automatically categorize a change between two snapshots.
 
@@ -19,6 +23,10 @@ def categorize_change(
     Args:
         new: The new snapshot.
         old: The old snapshot.
+        config: Configuration for the automatic categorizer of snapshot changes.
+        is_breaking_change: Callable that compares two models (new, old) and determines
+            whether there is a breaking change between them.
+        kwargs: Additional arguments to pass to is_breaking_change.
 
     Returns:
         The change category or None if the category can't be determined automatically.
@@ -48,11 +56,14 @@ def categorize_change(
             return SnapshotChangeCategory.NON_BREAKING
         return None
 
-    is_breaking_change = new_model.is_breaking_change(old_model)
-    if is_breaking_change is None:
+    breaking_change = (
+        is_breaking_change(new_model, old_model, **kwargs)
+        if is_breaking_change
+        else new_model.is_breaking_change(old_model)
+    )
+    if breaking_change is None:
         return default_category
+
     return (
-        SnapshotChangeCategory.BREAKING
-        if is_breaking_change
-        else SnapshotChangeCategory.NON_BREAKING
+        SnapshotChangeCategory.BREAKING if breaking_change else SnapshotChangeCategory.NON_BREAKING
     )
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 839288da..01bf09d0 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -199,7 +199,10 @@ class SQLMeshMagics(Magics):
                 expressions = parse(file.read(), default_dialect=config.dialect)
 
         formatted = format_model_expressions(
-            expressions, model.dialect, **config.format.generator_options
+            expressions,
+            model.dialect,
+            rewrite_casts=not config.format.no_rewrite_casts,
+            **config.format.generator_options,
         )
 
         self._shell.set_next_input(
@@ -703,6 +706,12 @@ class SQLMeshMagics(Magics):
         help="Whether or not to append a newline to the end of the file.",
         default=None,
     )
+    @argument(
+        "--no-rewrite-casts",
+        action="store_true",
+        help="Whether or not to preserve the existing casts, without rewriting them to use the :: syntax.",
+        default=None,
+    )
     @argument(
         "--normalize",
         action="store_true",
@@ -745,8 +754,11 @@ class SQLMeshMagics(Magics):
     @pass_sqlmesh_context
     def format(self, context: Context, line: str) -> bool:
         """Format all SQL models and audits."""
-        args = parse_argstring(self.format, line)
-        return context.format(**{k: v for k, v in vars(args).items() if v is not None})
+        format_opts = vars(parse_argstring(self.format, line))
+        if format_opts.pop("no_rewrite_casts", None):
+            format_opts["rewrite_casts"] = False
+
+        return context.format(**{k: v for k, v in format_opts.items() if v is not None})
 
     @magic_arguments()
     @argument("environment", type=str, help="The environment to diff local state against.")


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index edbfdd68..195c5b6a 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -194,6 +194,25 @@ SELECT
   SAFE_CAST('bla' AS INT64) AS FOO"""
     )
 
+    x = format_model_expressions(
+        parse(
+            """
+            MODEL(name foo);
+            SELECT 1::INT AS bla
+            """
+        ),
+        rewrite_casts=False,
+    )
+    assert (
+        x
+        == """MODEL (
+  name foo
+);
+
+SELECT
+  CAST(1 AS INT) AS bla"""
+    )
+
 
 def test_macro_format():
     assert parse_one("@EACH(ARRAY(1,2), x -> x)").sql() == "@EACH(ARRAY(1, 2), x -> x)"