Fix!: Improve the table_name command (#4299)

Author: izeigerman

docs/guides/table_migration.md

@@ -131,7 +131,7 @@ Consider an existing table named `my_schema.existing_table`. Migrating this tabl
 
     c. Create the model in the SQLMesh project without backfilling any data by running `sqlmesh plan [environment name] --empty-backfill --start 2024-01-01`, replacing "[environment name]" with an environment name other than `prod` and using the same start date from the `MODEL` DDL in step 3b.
 
-4. Determine the name of the model's snapshot physical table by running `sqlmesh table_name my_schema.existing_table`. For example, it might return `sqlmesh__my_schema.existing_table_123456`.
+4. Determine the name of the model's snapshot physical table by running `sqlmesh table_name --env [environment name] my_schema.existing_table`. For example, it might return `sqlmesh__my_schema.existing_table_123456`.
 5. Rename the original table `my_schema.existing_table_temp` to `sqlmesh__my_schema.existing_table_123456`
 
 The model would have code similar to:

docs/reference/cli.md

@@ -559,9 +559,11 @@ 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.
+  --environment, --env TEXT  The environment to source the model version from.
+  --prod                     If set, return the name of the physical table
+                             that will be used in production for the model
+                             version promoted in the target environment.
+  --help                     Show this message and exit.
 ```
 
 ## test

sqlmesh/cli/main.py

@@ -29,6 +29,7 @@
     "run",
     "environments",
     "invalidate",
+    "table_name",
 )
 SKIP_CONTEXT_COMMANDS = ("init", "ui")
 
@@ -988,17 +989,24 @@ def clean(obj: Context) -> None:
 @cli.command("table_name")
 @click.argument("model_name", required=True)
 @click.option(
-    "--dev",
+    "--environment",
+    "--env",
+    help="The environment to source the model version from.",
+)
+@click.option(
+    "--prod",
     is_flag=True,
-    help="Print the name of the snapshot table used for previews in development environments.",
     default=False,
+    help="If set, return the name of the physical table that will be used in production for the model version promoted in the target environment.",
 )
 @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, prod: bool = False
+) -> 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, prod))
 
 
 @cli.command("dlt_refresh")

sqlmesh/core/context.py

@@ -2135,25 +2135,45 @@ 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, prod: bool = False
+    ) -> 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.
+            prod: If True, return the name of the physical table that will be used in production for the model version
+                promoted in the target environment.
 
         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 or 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_deployable(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:

sqlmesh/magics.py

@@ -721,16 +721,23 @@ def table_diff(self, context: Context, line: str) -> None:
         help="The name of the model to get the table name for.",
     )
     @argument(
-        "--dev",
+        "--environment",
+        type=str,
+        help="The environment to source the model version from.",
+    )
+    @argument(
+        "--prod",
         action="store_true",
-        help="Print the name of the snapshot table used for previews in development environments.",
+        help="If set, return the name of the physical table that will be used in production for the model version promoted in the target environment.",
     )
     @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))
+        context.console.log_status_update(
+            context.table_name(args.model_name, args.environment, args.prod)
+        )
 
     @magic_arguments()
     @argument(

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,60 @@ 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"
+    )
+
+    assert (
+        context.table_name("sushi.waiter_revenue_by_day", "dev_b", prod=True)
+        == f"memory.sqlmesh__sushi.sushi__waiter_revenue_by_day__{snapshot.version}"
+    )
+
+
 @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"}