SNOW-981562: Expose is_temp_table_for_cleanup to Session.table (#2784)

Author: sfc-gh-aalam

src/snowflake/snowpark/session.py

@@ -2268,7 +2268,12 @@ def update_query_tag(self, tag: dict) -> None:
                 )
 
     @publicapi
-    def table(self, name: Union[str, Iterable[str]], _emit_ast: bool = True) -> Table:
+    def table(
+        self,
+        name: Union[str, Iterable[str]],
+        is_temp_table_for_cleanup: bool = False,
+        _emit_ast: bool = True,
+    ) -> Table:
         """
         Returns a Table that points the specified table.
 
@@ -2302,13 +2307,20 @@ def table(self, name: Union[str, Iterable[str]], _emit_ast: bool = True) -> Tabl
             elif isinstance(name, Iterable):
                 ast.name.sp_table_name_structured.name.extend(name)
             ast.variant.sp_session_table = True
+            ast.is_temp_table_for_cleanup = is_temp_table_for_cleanup
         else:
             stmt = None
 
         if not isinstance(name, str) and isinstance(name, Iterable):
             name = ".".join(name)
         validate_object_name(name)
-        t = Table(name, session=self, _ast_stmt=stmt, _emit_ast=_emit_ast)
+        t = Table(
+            name,
+            session=self,
+            is_temp_table_for_cleanup=is_temp_table_for_cleanup,
+            _ast_stmt=stmt,
+            _emit_ast=_emit_ast,
+        )
         # Replace API call origin for table
         set_api_call_source(t, "Session.table")
         return t

tests/ast/data/session_table_temp_table_cleanup.test

@@ -0,0 +1,174 @@
+## TEST CASE
+
+df1 = session.table(f"mock_schema.{tables.table1}", is_temp_table_for_cleanup=True)
+df2 = session.table(f"mock_schema.{tables.table1}", is_temp_table_for_cleanup=False)
+df = df1.union_all(df2).select("num")
+
+## EXPECTED UNPARSER OUTPUT
+
+df1 = session.table("mock_schema.table1", is_temp_table_for_cleanup=True)
+
+df2 = session.table("mock_schema.table1")
+
+df = df1.union_all(df2)
+
+df = df.select(col("num"))
+
+## EXPECTED ENCODED AST
+
+body {
+  assign {
+    expr {
+      sp_table {
+        is_temp_table_for_cleanup: true
+        name {
+          sp_table_name_flat {
+            name: "mock_schema.table1"
+          }
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 25
+        }
+        variant {
+          sp_session_table: true
+        }
+      }
+    }
+    symbol {
+      value: "df1"
+    }
+    uid: 1
+    var_id {
+      bitfield1: 1
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_table {
+        name {
+          sp_table_name_flat {
+            name: "mock_schema.table1"
+          }
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 26
+        }
+        variant {
+          sp_session_table: true
+        }
+      }
+    }
+    symbol {
+      value: "df2"
+    }
+    uid: 2
+    var_id {
+      bitfield1: 2
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_union_all {
+        df {
+          sp_dataframe_ref {
+            id {
+              bitfield1: 1
+            }
+          }
+        }
+        other {
+          sp_dataframe_ref {
+            id {
+              bitfield1: 2
+            }
+          }
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 27
+        }
+      }
+    }
+    symbol {
+      value: "df"
+    }
+    uid: 3
+    var_id {
+      bitfield1: 3
+    }
+  }
+}
+body {
+  assign {
+    expr {
+      sp_dataframe_select__columns {
+        cols {
+          apply_expr {
+            fn {
+              builtin_fn {
+                name {
+                  fn_name_flat {
+                    name: "col"
+                  }
+                }
+              }
+            }
+            pos_args {
+              string_val {
+                src {
+                  file: "SRC_POSITION_TEST_MODE"
+                  start_line: 27
+                }
+                v: "num"
+              }
+            }
+            src {
+              file: "SRC_POSITION_TEST_MODE"
+              start_line: 27
+            }
+          }
+        }
+        df {
+          sp_dataframe_ref {
+            id {
+              bitfield1: 3
+            }
+          }
+        }
+        src {
+          file: "SRC_POSITION_TEST_MODE"
+          start_line: 27
+        }
+        variadic: true
+      }
+    }
+    symbol {
+      value: "df"
+    }
+    uid: 4
+    var_id {
+      bitfield1: 4
+    }
+  }
+}
+client_ast_version: 1
+client_language {
+  python_language {
+    version {
+      label: "final"
+      major: 3
+      minor: 9
+      patch: 1
+    }
+  }
+}
+client_version {
+  major: 1
+  minor: 26
+}