SNOW-3473261: Add Iceberg tag time travel via version_tag (#4211)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: sfc-gh-igarish

src/snowflake/snowpark/_internal/utils.py

@@ -1956,6 +1956,7 @@ class TimeTravelConfig(NamedTuple):
     timestamp_type: Optional[str] = None
     stream: Optional[str] = None
     version: Optional[int] = None
+    version_tag: Optional[str] = None
 
     @staticmethod
     def validate_and_normalize_params(
@@ -1966,6 +1967,7 @@ def validate_and_normalize_params(
         timestamp_type: Optional[Union[str, "TimestampTimeZone"]] = None,
         stream: Optional[str] = None,
         version: Optional[int] = None,
+        version_tag: Optional[str] = None,
     ) -> Optional["TimeTravelConfig"]:
         """
         Validates and normalizes time travel parameters.
@@ -1988,7 +1990,8 @@ def validate_and_normalize_params(
             ValueError: If parameters are invalid.
         """
         time_travel_arg_count = sum(
-            arg is not None for arg in (statement, offset, timestamp, stream, version)
+            arg is not None
+            for arg in (statement, offset, timestamp, stream, version, version_tag)
         )
 
         # Validate mode
@@ -2023,10 +2026,32 @@ def validate_and_normalize_params(
                 f"'version' must be an int Iceberg snapshot id, got {type(version).__name__}."
             )
 
+        # version_tag (Iceberg tag name, mapped to Snowflake's
+        # ``AT(VERSION_TAG => '<name>')`` grammar) only works with 'at' mode —
+        # Iceberg tag reads are positional (bound to a specific snapshot),
+        # not range-of-time, so ``BEFORE`` has no meaning.
+        if version_tag is not None and time_travel_mode.lower() != "at":
+            raise ValueError(
+                "Iceberg version_tag time travel can only be used with "
+                "time_travel_mode='at', not 'before'."
+            )
+
+        # Validate version_tag type — Iceberg tag names are strings. Empty
+        # strings are invalid.
+        if version_tag is not None:
+            if not isinstance(version_tag, str):
+                raise ValueError(
+                    f"'version_tag' must be a string Iceberg tag name, "
+                    f"got {type(version_tag).__name__}."
+                )
+            if not version_tag:
+                raise ValueError("'version_tag' must be a non-empty Iceberg tag name.")
+
         # Validate exactly one parameter is provided
         if time_travel_arg_count != 1:
             raise ValueError(
-                "Exactly one of 'statement', 'offset', 'timestamp', 'stream', or 'version' must be provided."
+                "Exactly one of 'statement', 'offset', 'timestamp', 'stream', "
+                "'version', or 'version_tag' must be provided."
             )
 
         # Normalize timestamp
@@ -2061,6 +2086,7 @@ def validate_and_normalize_params(
             timestamp_type=timestamp_type,
             stream=stream,
             version=version,
+            version_tag=version_tag,
         )
 
     def generate_sql_clause(self) -> str:
@@ -2069,19 +2095,41 @@ def generate_sql_clause(self) -> str:
         Args:
             config: Time travel configuration.
         Returns:
-            SQL clause like " AT (TIMESTAMP => TO_TIMESTAMP_NTZ('...'))" or
-            " AT (VERSION => 1234567890)" for Iceberg snapshot id time travel.
+            SQL clause like " AT (TIMESTAMP => TO_TIMESTAMP_NTZ('...'))",
+            " AT (VERSION => 1234567890)" for Iceberg snapshot id time travel,
+            or " AT (VERSION_TAG => 'release_v1')" for Iceberg tag time
+            travel.
+
+        Note on escaping: string-valued parameters (``statement``,
+        ``stream``, ``version_tag``, ``timestamp``) are embedded inside
+        single-quoted SQL literals via the existing ``str_to_sql``
+        helper in ``analyzer.datatype_mapper`` so embedded ``'``, ``\\``
+        and newline characters are properly escaped. This keeps the
+        emission consistent with the rest of Snowpark Python's SQL
+        generation and avoids both broken SQL and injection surface
+        (e.g. ``x'); DROP TABLE foo; --``). Numeric parameters
+        (``offset``, ``version``) are not quoted and don't need
+        escaping.
+
+        ``str_to_sql`` is imported lazily here because
+        ``analyzer.datatype_mapper`` imports from this module at top
+        level — same pattern used elsewhere in this file for analyzer
+        cross-references.
         """
+        from snowflake.snowpark._internal.analyzer.datatype_mapper import str_to_sql
+
         clause = f" {self.time_travel_mode.upper()} "
 
         if self.statement is not None:
-            clause += f"(STATEMENT => '{self.statement}')"
+            clause += f"(STATEMENT => {str_to_sql(self.statement)})"
         elif self.offset is not None:
             clause += f"(OFFSET => {self.offset})"
         elif self.stream is not None:
-            clause += f"(STREAM => '{self.stream}')"
+            clause += f"(STREAM => {str_to_sql(self.stream)})"
         elif self.version is not None:
             clause += f"(VERSION => {self.version})"
+        elif self.version_tag is not None:
+            clause += f"(VERSION_TAG => {str_to_sql(self.version_tag)})"
         elif self.timestamp is not None:
             if self.timestamp_type is not None:
                 timestamp_type = self.timestamp_type.upper()
@@ -2093,9 +2141,9 @@ def generate_sql_clause(self) -> str:
                     func_name = "TO_TIMESTAMP_TZ"
                 else:
                     func_name = "TO_TIMESTAMP"
-                clause += f"(TIMESTAMP => {func_name}('{self.timestamp}'))"
+                clause += f"(TIMESTAMP => {func_name}({str_to_sql(self.timestamp)}))"
             else:
-                clause += f"(TIMESTAMP => '{self.timestamp}')"
+                clause += f"(TIMESTAMP => {str_to_sql(self.timestamp)})"
 
         return clause
 

src/snowflake/snowpark/dataframe_reader.py

@@ -161,6 +161,15 @@ def _extract_time_travel_from_options(options: dict) -> dict:
     - Automatically set time_travel_mode to 'at'
       (Iceberg snapshot ids only support ``AT(VERSION => N)``, not ``BEFORE``)
     - Cannot be used with time_travel_mode='before' (raises error)
+
+    Special handling for 'VERSION_TAG' / 'VERSION-TAG' (Iceberg tag name) —
+    both aliases map to the internal ``version_tag`` time travel parameter
+    and emit ``AT(VERSION_TAG => '<name>')`` on the Snowflake side (see
+    Spark Iceberg's ``VERSION AS OF '<tag_name>'`` reader path):
+    - Automatically set time_travel_mode to 'at'
+      (tag reads are positional — bound to a specific snapshot — not
+      range-of-time)
+    - Cannot be used with time_travel_mode='before' (raises error)
     """
     result = {}
     excluded_keys = set()
@@ -211,6 +220,27 @@ def _extract_time_travel_from_options(options: dict) -> dict:
             )
         result["time_travel_mode"] = "at"
 
+    # Handle Iceberg tag (``version_tag`` / ``version-tag``). Both aliases
+    # route to the internal ``version_tag`` parameter and emit
+    # ``AT(VERSION_TAG => '<name>')`` server-side. Auto-sets mode='at'.
+    version_tag_value = options.get("VERSION_TAG")
+    version_tag_source = "version_tag"
+    if version_tag_value is None:
+        version_tag_value = options.get("VERSION-TAG")
+        version_tag_source = "version-tag"
+    if version_tag_value is not None:
+        if (
+            "TIME_TRAVEL_MODE" in options
+            and options["TIME_TRAVEL_MODE"].lower() == "before"
+        ):
+            raise ValueError(
+                f"Cannot use '{version_tag_source}' option with "
+                "time_travel_mode='before'. Iceberg tag time travel only "
+                "supports time_travel_mode='at'."
+            )
+        result["version_tag"] = str(version_tag_value)
+        result["time_travel_mode"] = "at"
+
     for option_key, param_name in _TIME_TRAVEL_OPTIONS_PARAMS_MAP.items():
         if option_key in options and option_key not in excluded_keys:
             result[param_name] = options[option_key]
@@ -634,11 +664,13 @@ def table(
             ...     .option("offset", -60)                 # This will be IGNORED
             ...     .table("my_table", time_travel_mode="at", offset=-3600))  # Only this is used
         """
-        # ``version`` (Iceberg snapshot id) is intentionally not in the public
-        # signature — it's consumed by Snowpark Connect and may be removed
-        # once a first-class API lands. Accept it through **kwargs so direct
-        # callers can still pass it without us advertising it.
+        # ``version`` (Iceberg snapshot id) and ``version_tag`` (Iceberg tag
+        # name) are intentionally not in the public signature — they are
+        # consumed by Snowpark Connect and may be removed once a first-class
+        # API lands. Accept them through **kwargs so direct callers can
+        # still pass them without us advertising the surface.
         version = kwargs.pop("version", None)
+        version_tag = kwargs.pop("version_tag", None)
         if kwargs:
             raise TypeError(
                 f"table() got unexpected keyword arguments: {sorted(kwargs)}"
@@ -664,13 +696,20 @@ def table(
             if stream is not None:
                 ast.stream.value = stream
 
-        if time_travel_mode is not None or version is not None:
-            # If version is provided without mode, default to 'at' (snapshot ids
-            # only make sense with AT — symmetric with iceberg_tag handling).
+        if (
+            time_travel_mode is not None
+            or version is not None
+            or version_tag is not None
+        ):
+            # If version / version_tag is provided without mode, default to
+            # 'at' — snapshot ids and tag reads only make sense with AT
+            # (symmetric with the as-of-timestamp option handling).
             effective_mode = (
                 time_travel_mode
                 if time_travel_mode
-                else ("at" if version is not None else None)
+                else (
+                    "at" if (version is not None or version_tag is not None) else None
+                )
             )
             time_travel_params = {
                 "time_travel_mode": effective_mode,
@@ -680,6 +719,7 @@ def table(
                 "timestamp_type": timestamp_type,
                 "stream": stream,
                 "version": version,
+                "version_tag": version_tag,
             }
         else:
             # if time_travel_mode is not provided, extract time travel config from options

src/snowflake/snowpark/session.py

@@ -2773,11 +2773,13 @@ def table(
             # timestamp_type remains "NTZ" (user's explicit choice respected)
             >>> table2 = session.read.table("my_table", time_travel_mode="at", timestamp=tz_aware, timestamp_type="NTZ")  # doctest: +SKIP
         """
-        # ``version`` (Iceberg snapshot id) is intentionally not in the public
-        # signature — it's consumed by Snowpark Connect and may be removed
-        # once a first-class API lands. Accept it through **kwargs so direct
-        # callers can still pass it without us advertising it.
+        # ``version`` (Iceberg snapshot id) and ``version_tag`` (Iceberg tag
+        # name) are intentionally not in the public signature — they are
+        # consumed by Snowpark Connect and may be removed once a first-class
+        # API lands. Accept them through **kwargs so direct callers can
+        # still pass them without us advertising the surface.
         version = kwargs.pop("version", None)
+        version_tag = kwargs.pop("version_tag", None)
         if kwargs:
             raise TypeError(
                 f"table() got unexpected keyword arguments: {sorted(kwargs)}"
@@ -2820,6 +2822,7 @@ def table(
             timestamp_type=timestamp_type,
             stream=stream,
             version=version,
+            version_tag=version_tag,
         )
         # Replace API call origin for table
         set_api_call_source(t, "Session.table")

src/snowflake/snowpark/table.py

@@ -296,11 +296,13 @@ def __init__(
         stream: Optional[str] = None,
         **kwargs,
     ) -> None:
-        # ``version`` (Iceberg snapshot id) is intentionally not in the public
-        # signature — it's consumed by Snowpark Connect and may be removed
-        # once a first-class API lands. Accept it through **kwargs so direct
-        # callers can still pass it without us advertising it.
+        # ``version`` (Iceberg snapshot id) and ``version_tag`` (Iceberg tag
+        # name) are intentionally not in the public signature — they are
+        # consumed by Snowpark Connect and may be removed once a first-class
+        # API lands. Accept them through **kwargs so direct callers can
+        # still pass them without us advertising the surface.
         version = kwargs.pop("version", None)
+        version_tag = kwargs.pop("version_tag", None)
         if kwargs:
             raise TypeError(
                 f"Table() got unexpected keyword arguments: {sorted(kwargs)}"
@@ -333,6 +335,7 @@ def __init__(
             timestamp_type=timestamp_type,
             stream=stream,
             version=version,
+            version_tag=version_tag,
         )
 
         snowflake_table_plan = SnowflakeTable(

tests/integ/test_dataframe.py

@@ -8423,3 +8423,70 @@ def test_iceberg_snapshot_id_time_travel_dataframe_reader_option(session):
         session.read.option("snapshot-id", snapshot_id).table(table_fqn).collect()
     )
     assert via_kwarg == via_option
+
+
+# ----------------------------------------------------------------------
+# Iceberg tag (``version_tag=``) time travel.
+#
+# TODO(SNOW-3473261): Wire these up to a CI test account that has:
+#   * a Catalog-Linked Database (CLD) such as cldUnity / cldglue, AND
+#   * an unmanaged Iceberg table inside it that exposes a named tag
+#     (e.g. created via ``ALTER TABLE ... CREATE TAG <name>`` on the OSS
+#     Iceberg side, or via the catalog's tag API).
+#
+# Snowflake's ``AT(VERSION_TAG => '<name>')`` is the released tag-only
+# Iceberg time-travel clause (Spark Iceberg's
+# ``VERSION AS OF '<tag_name>'`` for tag reads). Like the snapshot-id
+# surface, it currently requires ``FEATURE_ICEBERG_TIME_TRAVEL`` on the
+# account and is scoped to unmanaged Iceberg tables in CLDs, so these
+# tests are skipped by default and exercised manually against
+# ``sfctest0`` (see ``tests/sas_tests/test_iceberg_version_tag_sample.py``
+# in snowflake-eng/sas for the manual reproducer).
+# ----------------------------------------------------------------------
+@pytest.mark.skip(
+    reason=(
+        "Requires a CLD-linked unmanaged Iceberg table with at least one "
+        "named tag and FEATURE_ICEBERG_TIME_TRAVEL enabled on the account. "
+        "Tested manually; see TODO above."
+    )
+)
+def test_iceberg_version_tag_time_travel_session_table(session):
+    """End-to-end: ``Session.table(..., version_tag='<name>')`` returns the
+    table state at the requested Iceberg tag."""
+    table_fqn = "CLDUNITY.scosschema.snapshot_demo"
+    # Demo table is set up with a tag ``first_load`` pointing at the
+    # earliest snapshot (see sas-side reproducer for the setup script).
+    tag_name = "first_load"
+
+    tagged = session.table(
+        table_fqn, time_travel_mode="at", version_tag=tag_name
+    ).collect()
+    latest = session.table(table_fqn).collect()
+    # Tag reads use the snapshot's schema as it existed at the tag —
+    # row count at an earlier tag should be ≤ current row count, since the
+    # tag is bound to a specific (earlier) snapshot.
+    assert len(tagged) <= len(latest)
+
+
+@pytest.mark.skip(
+    reason=(
+        "Requires a CLD-linked unmanaged Iceberg table with at least one "
+        "named tag and FEATURE_ICEBERG_TIME_TRAVEL enabled on the account. "
+        "Tested manually; see TODO above."
+    )
+)
+def test_iceberg_version_tag_time_travel_dataframe_reader_option(session):
+    """End-to-end: ``session.read.option('version_tag', 'name').table(...)``
+    routes through the Iceberg-compat option alias and produces the same
+    result as the explicit ``version_tag=`` kwarg."""
+    table_fqn = "CLDUNITY.scosschema.snapshot_demo"
+    tag_name = "first_load"
+
+    via_kwarg = session.read.table(
+        table_fqn, time_travel_mode="at", version_tag=tag_name
+    ).collect()
+    via_option = session.read.option("version_tag", tag_name).table(table_fqn).collect()
+    via_hyphen_option = (
+        session.read.option("version-tag", tag_name).table(table_fqn).collect()
+    )
+    assert via_kwarg == via_option == via_hyphen_option