From be8ebcdbad653eab408cb446ebf7f6965f6a1ec5 Mon Sep 17 00:00:00 2001 From: Sophie Tan Date: Wed, 21 Feb 2024 13:52:00 -0800 Subject: [PATCH] SNOW-1058663 Fix `ServerConnection. _to_data_or_iter` to return the original query id (#1260) * Return original query ID * Address comments --- src/snowflake/snowpark/_internal/server_connection.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/snowflake/snowpark/_internal/server_connection.py b/src/snowflake/snowpark/_internal/server_connection.py index e07a9ac50ec..a650087d6a7 100644 --- a/src/snowflake/snowpark/_internal/server_connection.py +++ b/src/snowflake/snowpark/_internal/server_connection.py @@ -442,13 +442,12 @@ def _to_data_or_iter( to_pandas: bool = False, to_iter: bool = False, ) -> Dict[str, Any]: + qid = results_cursor.sfqid if ( to_iter and not to_pandas ): # Fix for SNOW-869536, to_pandas doesn't have this issue, SnowflakeCursor.fetch_pandas_batches already handles the isolation. new_cursor = results_cursor.connection.cursor() - new_cursor.execute( - f"SELECT * FROM TABLE(RESULT_SCAN('{results_cursor.sfqid}'))" - ) + new_cursor.execute(f"SELECT * FROM TABLE(RESULT_SCAN('{qid}'))") results_cursor = new_cursor if to_pandas: @@ -481,7 +480,7 @@ def _to_data_or_iter( iter(results_cursor) if to_iter else results_cursor.fetchall() ) - return {"data": data_or_iter, "sfqid": results_cursor.sfqid} + return {"data": data_or_iter, "sfqid": qid} def execute( self,