Repository: TobikoData/sqlmesh
Pull request #2361 (2024-03-28 12:23:12)

Feat: allow UNION macros as model queries


Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 4f646924..03a7bcb1 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -341,6 +341,18 @@ Options:
   --help           Show this message and exit.
 ```
 
+## table_name
+```
+Usage: sqlmesh table_name [OPTIONS] MODEL_NAME
+
+  Prints the name of the physical table for the given model.
+
+Options:
+  --dev   Print the name of the snapshot table used for previews in
+          development environments.
+  --help  Show this message and exit.
+```
+
 ## test
 ```
 Usage: sqlmesh test [OPTIONS] [TESTS]...
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 6c845c03..2e4eaeeb 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -739,3 +739,18 @@ def prompt(
 def clean(obj: Context) -> None:
     """Clears the SQLMesh cache and any build artifacts."""
     obj.clear_caches()
+
+
+@cli.command("table_name")
+@click.argument("model_name", required=True)
+@click.option(
+    "--dev",
+    is_flag=True,
+    help="Print the name of the snapshot table used for previews in development environments.",
+    default=False,
+)
+@click.pass_obj
+@error_handler
+def table_name(obj: Context, model_name: str, dev: bool) -> None:
+    """Prints the name of the physical table for the given model."""
+    print(obj.table_name(model_name, dev))
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 0e8e57d9..42df4d8b 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1583,6 +1583,26 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return success
 
+    def table_name(self, model_name: str, dev: bool) -> str:
+        """Returns the name of the pysical table for the given model name.
+
+        Args:
+            model_name: The name of the model.
+            dev: Whether to use the deployability index for the table name.
+
+        Returns:
+            The name of the physical table.
+        """
+        deployability_index = (
+            DeployabilityIndex.create(self.snapshots.values())
+            if dev
+            else DeployabilityIndex.all_deployable()
+        )
+        snapshot = self.get_snapshot(model_name)
+        if not snapshot:
+            raise SQLMeshError(f"Model '{model_name}' was not found.")
+        return snapshot.table_name(is_deployable=deployability_index.is_deployable(snapshot))
+
     def clear_caches(self) -> None:
         for path in self.configs:
             rmtree(path / c.CACHE)
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 56e2b1f6..accbe982 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1477,8 +1477,12 @@ def load_sql_based_model(
         **meta_fields,
     )
 
-    if query_or_seed_insert is not None and isinstance(
-        query_or_seed_insert, (exp.Query, d.JinjaQuery)
+    if query_or_seed_insert is not None and (
+        isinstance(query_or_seed_insert, (exp.Query, d.JinjaQuery))
+        or (
+            isinstance(query_or_seed_insert, d.MacroFunc)
+            and query_or_seed_insert.this.name.lower() == "union"
+        )
     ):
         jinja_macro_references.update(
             extract_macro_references(query_or_seed_insert.sql(dialect=dialect))
@@ -1792,6 +1796,7 @@ def _split_sql_model_statements(
         if (
             isinstance(expression, (exp.Query, d.JinjaQuery))
             or expression == INSERT_SEED_MACRO_CALL
+            or (isinstance(expression, d.MacroFunc) and expression.this.name.lower() == "union")
         ):
             query_positions.append((expression, idx))
 
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 9952ec4d..1b98d4c1 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -597,6 +597,25 @@ class SQLMeshMagics(Magics):
             show_sample=args.show_sample,
         )
 
+    @magic_arguments()
+    @argument(
+        "model_name",
+        nargs="?",
+        type=str,
+        help="The name of the model to get the table name for.",
+    )
+    @argument(
+        "--dev",
+        action="store_true",
+        help="Print the name of the snapshot table used for previews in development environments.",
+    )
+    @line_magic
+    @pass_sqlmesh_context
+    def table_name(self, context: Context, line: str) -> None:
+        """Prints the name of the physical table for the given model."""
+        args = parse_argstring(self.table_name, line)
+        context.console.log_status_update(context.table_name(args.model_name, args.dev))
+
     @magic_arguments()
     @argument(
         "--read",


Test patch:
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index b416f193..497a32cd 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -610,3 +610,22 @@ def test_clean(runner, tmp_path):
     # Confirm cache was cleared
     assert result.exit_code == 0
     assert not cache_path.exists()
+
+
+def test_table_name(runner, tmp_path):
+    # Create and backfill `prod` environment
+    create_example_project(tmp_path)
+    init_prod_and_backfill(runner, tmp_path)
+    result = runner.invoke(
+        cli,
+        [
+            "--log-file-dir",
+            tmp_path,
+            "--paths",
+            tmp_path,
+            "table_name",
+            "sqlmesh_example.full_model",
+        ],
+    )
+    assert result.exit_code == 0
+    assert result.output.startswith("db.sqlmesh__sqlmesh_example.sqlmesh_example__full_model__")
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 67ebe5cd..f1f9ccab 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -190,7 +190,7 @@ def test_model_validation(query, error):
     assert error in str(ex.value)
 
 
-def test_model_union_query():
+def test_model_union_query(sushi_context, assert_exp_eq):
     expressions = d.parse(
         """
         MODEL (
@@ -204,6 +204,37 @@ def test_model_union_query():
 
     load_sql_based_model(expressions)
 
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test,
+            kind FULL,
+        );
+
+        @union('all', sushi.marketing, sushi.marketing)
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+    assert_exp_eq(
+        sushi_context.get_model("sushi.test").render_query(),
+        """SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMP) AS "updated_at",
+  CAST("marketing"."valid_from" AS TIMESTAMP) AS "valid_from",
+  CAST("marketing"."valid_to" AS TIMESTAMP) AS "valid_to"
+FROM "memory"."sushi"."marketing" AS "marketing"
+UNION ALL
+SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMP) AS "updated_at",
+  CAST("marketing"."valid_from" AS TIMESTAMP) AS "valid_from",
+  CAST("marketing"."valid_to" AS TIMESTAMP) AS "valid_to"
+FROM "memory"."sushi"."marketing" AS "marketing"
+        """,
+    )
+
 
 def test_model_validation_union_query():
     expressions = d.parse(
diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py
index 88401375..d2f33aae 100644
--- a/tests/integrations/jupyter/test_magics.py
+++ b/tests/integrations/jupyter/test_magics.py
@@ -671,4 +671,15 @@ revenue      100.0""",
     ]
 
 
-# TODO: Add rewrite test
+@pytest.mark.slow
+@freeze_time(FREEZE_TIME)
+def test_table_name(notebook, loaded_sushi_context, convert_all_html_output_to_text):
+    with capture_output() as output:
+        notebook.run_line_magic(magic_name="table_name", line="sushi.orders")
+
+    assert not output.stdout
+    assert not output.stderr
+    assert len(output.outputs) == 1
+    assert convert_all_html_output_to_text(output)[0].startswith(
+        "memory.sqlmesh__sushi.sushi__orders__"
+    )