Fix: Improve the table_name command

Author: izeigerman

sqlmesh/cli/main.py

@@ -29,6 +29,7 @@
     "run",
     "environments",
     "invalidate",
+    "table_name",
 )
 SKIP_CONTEXT_COMMANDS = ("init", "ui")
 
@@ -988,17 +989,16 @@ def clean(obj: Context) -> None:
 @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,
+    "--environment",
+    "--env",
+    help="The environment to source the model version from.",
 )
 @click.pass_obj
 @error_handler
 @cli_analytics
-def table_name(obj: Context, model_name: str, dev: bool) -> None:
+def table_name(obj: Context, model_name: str, environment: t.Optional[str] = None) -> None:
     """Prints the name of the physical table for the given model."""
-    print(obj.table_name(model_name, dev))
+    print(obj.table_name(model_name, environment))
 
 
 @cli.command("dlt_refresh")

sqlmesh/core/context.py

@@ -2135,25 +2135,41 @@ def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]])
         )
 
     @python_api_analytics
-    def table_name(self, model_name: str, dev: bool) -> str:
-        """Returns the name of the pysical table for the given model name.
+    def table_name(self, model_name: str, environment: t.Optional[str] = None) -> str:
+        """Returns the name of the pysical table for the given model name in the target environment.
 
         Args:
             model_name: The name of the model.
-            dev: Whether to use the deployability index for the table name.
+            environment: The environment to source the model version from.
 
         Returns:
             The name of the physical table.
         """
-        deployability_index = (
-            DeployabilityIndex.create(self.snapshots.values())
-            if dev
-            else DeployabilityIndex.all_deployable()
+        environment = environment or self.config.default_target_environment
+        fqn = self._node_or_snapshot_to_fqn(model_name)
+        target_env = self.state_reader.get_environment(environment)
+        if not target_env:
+            raise SQLMeshError(f"Environment '{environment}' was not found.")
+
+        snapshot_info = None
+        for s in target_env.snapshots:
+            if s.name == fqn:
+                snapshot_info = s
+                break
+        if not snapshot_info:
+            raise SQLMeshError(
+                f"Model '{model_name}' was not found in environment '{environment}'."
+            )
+
+        if target_env.name == c.PROD:
+            return snapshot_info.table_name()
+
+        snapshots = self.state_reader.get_snapshots(target_env.snapshots)
+        deployability_index = DeployabilityIndex.create(snapshots)
+
+        return snapshot_info.table_name(
+            is_deployable=deployability_index.is_representative(snapshot_info.snapshot_id)
         )
-        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:

tests/core/test_integration.py

@@ -61,7 +61,7 @@
     SnapshotTableInfo,
 )
 from sqlmesh.utils.date import TimeLike, now, to_date, to_datetime, to_timestamp
-from sqlmesh.utils.errors import NoChangesPlanError
+from sqlmesh.utils.errors import NoChangesPlanError, SQLMeshError
 from sqlmesh.utils.pydantic import validate_string
 from tests.conftest import DuckDBMetadata, SushiDataValidator
 from tests.utils.test_helpers import use_terminal_console
@@ -3034,7 +3034,7 @@ def _dates_in_table(table_name: str) -> t.List[str]:
     ]
 
     # mess with A independently of SQLMesh to prove a whole day gets restated for B instead of just 1hr
-    snapshot_table_name = ctx.table_name("test.a", False)
+    snapshot_table_name = ctx.table_name("test.a", "dev")
     engine_adapter.execute(
         f"delete from {snapshot_table_name} where cast(ts as date) == '2024-01-01'"
     )
@@ -4094,6 +4094,55 @@ def test_evaluate_uncategorized_snapshot(init_and_plan_context: t.Callable):
     assert set(df["one"].tolist()) == {1}
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_table_name(init_and_plan_context: t.Callable):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    snapshot = context.get_snapshot("sushi.waiter_revenue_by_day")
+    assert snapshot
+    assert (
+        context.table_name("sushi.waiter_revenue_by_day", "prod")
+        == f"memory.sqlmesh__sushi.sushi__waiter_revenue_by_day__{snapshot.version}"
+    )
+
+    with pytest.raises(SQLMeshError, match="Environment 'dev' was not found."):
+        context.table_name("sushi.waiter_revenue_by_day", "dev")
+
+    with pytest.raises(
+        SQLMeshError, match="Model 'sushi.missing' was not found in environment 'prod'."
+    ):
+        context.table_name("sushi.missing", "prod")
+
+    # Add a new projection
+    model = context.get_model("sushi.waiter_revenue_by_day")
+    context.upsert_model(add_projection_to_model(t.cast(SqlModel, model)))
+
+    context.plan("dev_a", auto_apply=True, no_prompts=True, skip_tests=True)
+
+    new_snapshot = context.get_snapshot("sushi.waiter_revenue_by_day")
+    assert new_snapshot.version != snapshot.version
+
+    assert (
+        context.table_name("sushi.waiter_revenue_by_day", "dev_a")
+        == f"memory.sqlmesh__sushi.sushi__waiter_revenue_by_day__{new_snapshot.version}"
+    )
+
+    # Make a forward-only change
+    context.upsert_model(model, stamp="forward_only")
+
+    context.plan("dev_b", auto_apply=True, no_prompts=True, skip_tests=True, forward_only=True)
+
+    forward_only_snapshot = context.get_snapshot("sushi.waiter_revenue_by_day")
+    assert forward_only_snapshot.version == snapshot.version
+    assert forward_only_snapshot.dev_version != snapshot.version
+
+    assert (
+        context.table_name("sushi.waiter_revenue_by_day", "dev_b")
+        == f"memory.sqlmesh__sushi.sushi__waiter_revenue_by_day__{forward_only_snapshot.dev_version}__dev"
+    )
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_dbt_requirements(sushi_dbt_context: Context):
     assert set(sushi_dbt_context.requirements) == {"dbt-core", "dbt-duckdb"}