Repository: TobikoData/sqlmesh
Pull request #2518 (2024-04-25 19:54:40)

feat!: make star macro more flexible
Prior to this change Start macro required the types to be known since it always did a cast of the selected columns to their desired type. This change makes it so if we don't know the columns to types then we just put in the columns otherwise we will do the cast. Therefore now to use the macro we just have to know the columns of the model being referenced.

An alternative implementation would have been to check if each column has a type and if so cast, otherwise select but decided to just do an all or nothing approach. 

Consider this breaking because if someone starts using this feature then their models would not be compatible with someone using an earlier release. 


Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index 501fbe44..2bc18701 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/docs/concepts/macros/sqlmesh_macros.md
@@ -592,7 +592,9 @@ WHERE
 
 `@STAR` is named after SQL's star operator `*`, but it allows you to programmatically generate a set of column selections and aliases instead of just selecting all available columns. A query may use more than one `@STAR` and may also include explicit column selections.
 
-`@STAR` uses SQLMesh's knowledge of each table's columns and data types to generate the appropriate column list. The resulting query always `CAST`s columns to their data type in the source table.
+`@STAR` uses SQLMesh's knowledge of each table's columns and data types to generate the appropriate column list. 
+If the column data types are known then the resulting query `CAST`s columns to their data type in the source table.
+Otherwise, the columns will be listed without any casting. 
 
 `@STAR` supports the following arguments, in this order:
 
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 0b52569a..2a46bc7f 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -25,7 +25,12 @@ from sqlmesh.core.dialect import (
     StagedFilePath,
     normalize_model_name,
 )
-from sqlmesh.utils import DECORATOR_RETURN_TYPE, UniqueKeyDict, registry_decorator
+from sqlmesh.utils import (
+    DECORATOR_RETURN_TYPE,
+    UniqueKeyDict,
+    columns_to_types_all_known,
+    registry_decorator,
+)
 from sqlmesh.utils.errors import MacroEvalError, SQLMeshError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, has_jinja
 from sqlmesh.utils.metaprogramming import Executable, prepare_env, print_exception
@@ -662,11 +667,11 @@ def star(
         An array of columns.
 
     Example:
-        >>> from sqlglot import parse_one
+        >>> from sqlglot import parse_one, exp
         >>> from sqlglot.schema import MappingSchema
         >>> from sqlmesh.core.macros import MacroEvaluator
         >>> sql = "SELECT @STAR(foo, bar, [c], 'baz_') FROM foo AS bar"
-        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "string", "b": "string", "c": "string", "d": "int"}})).transform(parse_one(sql)).sql()
+        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": exp.DataType.build("string"), "b": exp.DataType.build("string"), "c": exp.DataType.build("string"), "d": exp.DataType.build("int")}})).transform(parse_one(sql)).sql()
         'SELECT CAST("bar"."a" AS TEXT) AS "baz_a", CAST("bar"."b" AS TEXT) AS "baz_b", CAST("bar"."d" AS INT) AS "baz_d" FROM foo AS bar'
     """
     if alias and not isinstance(alias, (exp.Identifier, exp.Column)):
@@ -684,14 +689,23 @@ def star(
     quoted = quote_identifiers.this
     table_identifier = alias.name or relation.name
 
+    columns_to_types = {
+        k: v for k, v in evaluator.columns_to_types(relation).items() if k not in exclude
+    }
+    if columns_to_types_all_known(columns_to_types):
+        return [
+            exp.cast(
+                exp.column(column, table=table_identifier, quoted=quoted),
+                dtype,
+                dialect=evaluator.dialect,
+            ).as_(f"{prefix.this}{column}{suffix.this}", quoted=quoted)
+            for column, dtype in columns_to_types.items()
+        ]
     return [
-        exp.cast(
-            exp.column(column, table=table_identifier, quoted=quoted),
-            type_,
-            dialect=evaluator.dialect,
-        ).as_(f"{prefix.this}{column}{suffix.this}", quoted=quoted)
+        exp.column(column, table=table_identifier, quoted=quoted).as_(
+            f"{prefix.this}{column}{suffix.this}", quoted=quoted
+        )
         for column, type_ in evaluator.columns_to_types(relation).items()
-        if column not in exclude
     ]
 
 


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 32b7a31f..df73591f 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -36,7 +36,17 @@ def macro_evaluator() -> MacroEvaluator:
 def test_star(assert_exp_eq) -> None:
     sql = """SELECT @STAR(foo) FROM foo"""
     expected_sql = """SELECT CAST([foo].[a] AS DATETIMEOFFSET) AS [a] FROM foo"""
-    schema = MappingSchema({"foo": {"a": "datetimeoffset"}}, dialect="tsql")
+    schema = MappingSchema(
+        {"foo": {"a": exp.DataType.build("datetimeoffset", dialect="tsql")}}, dialect="tsql"
+    )
+    evaluator = MacroEvaluator(schema=schema, dialect="tsql")
+    assert_exp_eq(evaluator.transform(parse_one(sql, read="tsql")), expected_sql, dialect="tsql")
+
+
+def test_start_no_column_types(assert_exp_eq) -> None:
+    sql = """SELECT @STAR(foo) FROM foo"""
+    expected_sql = """SELECT [foo].[a] AS [a] FROM foo"""
+    schema = MappingSchema({"foo": {"a": exp.DataType.build("UNKNOWN")}}, dialect="tsql")
     evaluator = MacroEvaluator(schema=schema, dialect="tsql")
     assert_exp_eq(evaluator.transform(parse_one(sql, read="tsql")), expected_sql, dialect="tsql")