diff --git a/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py b/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py index 69453d48596..ce5ff638fd1 100644 --- a/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py +++ b/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 @@ -68,13 +69,28 @@ 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}") + 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 + 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) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index b1920eb7e8d..d2bc87bd2a8 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -582,11 +582,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._query_compilation_stage_enabled: bool = ( self._conn._get_client_side_session_parameter(