SNOW-1709861: table clean up uses its own cursor

src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Dict
 
 from snowflake.snowpark._internal.analyzer.snowflake_plan_node import SnowflakeTable
+from snowflake.snowpark._internal.utils import is_in_stored_procedure
 
 if TYPE_CHECKING:
     from snowflake.snowpark.session import Session  # pragma: no cover
@@ -51,6 +52,20 @@ def _delete_ref_count(self, name: str) -> None:  # pragma: no cover
             self.ref_count_map[name] -= 1
             current_ref_count = self.ref_count_map[name]
         if current_ref_count == 0:
+            if (
+                is_in_stored_procedure()
+                and not self.session._conn._get_client_side_session_parameter(
+                    "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS", False
+                )
+            ):
+                warning_message = "Drop table requires async query which is not supported in stored procedure yet"
+                logging.warning(warning_message)
+                self.session._conn._telemetry_client.send_temp_table_cleanup_abnormal_exception_telemetry(
+                    self.session.session_id,
+                    name,
+                    str(warning_message),
+                )
+                return
             if (
                 self.session.auto_clean_up_temp_table_enabled
                 # if the session is already closed before garbage collection,
@@ -68,13 +83,14 @@ def drop_table(self, name: str) -> None:  # pragma: no cover
         logging.debug(f"Ready to drop {common_log_text}")
         query_id = None
         try:
-            async_job = self.session.sql(
-                f"drop table if exists {name} /* internal query to drop unused temp table */",
-            )._internal_collect_with_tag_no_telemetry(
-                block=False, statement_params={DROP_TABLE_STATEMENT_PARAM_NAME: name}
-            )
-            query_id = async_job.query_id
-            logging.debug(f"Dropping {common_log_text} with query id {query_id}")
+            with self.session.connection.cursor() as cursor:
+                async_job_query_id = cursor.execute_async(
+                    command=f"drop table if exists {name}",
+                    _statement_params={DROP_TABLE_STATEMENT_PARAM_NAME: name},
+                )["queryId"]
+                logging.debug(
+                    f"Dropping {common_log_text} with query id {async_job_query_id}"
+                )
         except Exception as ex:  # pragma: no cover
             warning_message = f"Failed to drop {common_log_text}, exception: {ex}"
             logging.warning(warning_message)

src/snowflake/snowpark/session.py

@@ -585,11 +585,12 @@ def __init__(
                 _PYTHON_SNOWPARK_ELIMINATE_NUMERIC_SQL_VALUE_CAST_ENABLED, False
             )
         )
-        self._auto_clean_up_temp_table_enabled: bool = (
-            self._conn._get_client_side_session_parameter(
-                _PYTHON_SNOWPARK_AUTO_CLEAN_UP_TEMP_TABLE_ENABLED, False
-            )
-        )
+        self._auto_clean_up_temp_table_enabled = True
+        # self._auto_clean_up_temp_table_enabled: bool = (
+        #     self._conn._get_client_side_session_parameter(
+        #         _PYTHON_SNOWPARK_AUTO_CLEAN_UP_TEMP_TABLE_ENABLED, False
+        #     )
+        # )
         self._reduce_describe_query_enabled: bool = (
             self._conn._get_client_side_session_parameter(
                 _PYTHON_SNOWPARK_REDUCE_DESCRIBE_QUERY_ENABLED, False

tests/integ/test_deepcopy.py

@@ -387,9 +387,6 @@ def test_deep_nested_select(session):
     [
         lambda session_: session_.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"]),
         lambda session_: session_.sql("select 1 as a, 2 as b"),
-        lambda session_: session_.table(
-            session_.sql("select 1 as a, 2 as b").cache_result().table_name
-        ),
     ],
 )
 def test_deepcopy_no_duplicate(session, generator):