Smirzaei snow 932731 snowpark sp cancel query (#1084)

* Cancel query API through new end-point

src/snowflake/snowpark/async_job.py

@@ -7,11 +7,13 @@
 from typing import TYPE_CHECKING, Iterator, List, Literal, Optional, Union
 
 import snowflake.snowpark
+from snowflake.connector.errors import DatabaseError
 from snowflake.connector.options import pandas
 from snowflake.snowpark._internal.analyzer.analyzer_utils import result_scan_statement
 from snowflake.snowpark._internal.analyzer.snowflake_plan import Query
 from snowflake.snowpark._internal.utils import (
     check_is_pandas_dataframe_in_to_pandas,
+    is_in_stored_procedure,
     result_set_to_iter,
     result_set_to_rows,
 )
@@ -242,7 +244,19 @@ def is_done(self) -> bool:
     def cancel(self) -> None:
         """Cancels the query associated with this instance."""
         # stop and cancel current query id
-        self._cursor.execute(f"select SYSTEM$CANCEL_QUERY('{self.query_id}')")
+        if (
+            is_in_stored_procedure()
+            and self._session._conn._get_client_side_session_parameter(
+                "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS", False
+            )
+        ):
+            cancel_resp = self._session._conn._conn.cancel_query(self.query_id)
+            if cancel_resp.get("success", False):
+                raise DatabaseError(
+                    f"Failed to cancel query. Returned response: {cancel_resp}"
+                )
+        else:
+            self._cursor.execute(f"select SYSTEM$CANCEL_QUERY('{self.query_id}')")
 
     def _table_result(
         self,

tests/integ/scala/test_async_job_suite.py

@@ -340,10 +340,6 @@ def test_async_batch_insert(session):
         analyzer.ARRAY_BIND_THRESHOLD = original_value
 
 
-@pytest.mark.skipif(
-    IS_IN_STORED_PROC,
-    reason="TODO(SNOW-932722): Cancel query is not allowed in stored proc",
-)
 def test_async_is_running_and_cancel(session):
     async_job = session.sql("select SYSTEM$WAIT(3)").collect_nowait()
     while not async_job.is_done():