SNOW-2437690 make sproc cancel API work with public connector (#3914)

Author: sfc-gh-zyao

src/snowflake/snowpark/async_job.py

@@ -284,10 +284,34 @@ def cancel(self) -> None:
                 "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS", False
             )
         ):
-            cancel_resp = self._session._conn._conn.cancel_query(self.query_id)
-            if not cancel_resp.get("success", False):
+            import _snowflake
+            import json
+            import uuid
+
+            try:
+                uuid.UUID(self.query_id)
+            except ValueError:
+                raise ValueError(f"Invalid UUID: '{self.query_id}'")
+
+            raw_cancel_resp = _snowflake.cancel_query(self.query_id)
+
+            # Set failure_response when
+            #   - success != True in the response or
+            #   - cannot parse the response at all.
+            failure_response = None
+            try:
+                parsed_cancel_resp = json.loads(raw_cancel_resp)
+                if not parsed_cancel_resp.get("success", False):
+                    failure_response = parsed_cancel_resp
+            except (TypeError, json.JSONDecodeError) as e:
+                failure_response = {
+                    "success": False,
+                    "error": f"Error parsing response: {e}",
+                }
+
+            if failure_response:
                 raise DatabaseError(
-                    f"Failed to cancel query. Returned response: {cancel_resp}"
+                    f"Failed to cancel query. Returned response: {failure_response}"
                 )
         else:
             self._cursor.execute(f"select SYSTEM$CANCEL_QUERY('{self.query_id}')")

tests/integ/scala/test_async_job_suite.py

@@ -369,10 +369,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():

tests/unit/test_async_job.py

@@ -0,0 +1,137 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import sys
+import uuid
+from types import SimpleNamespace
+from unittest import mock
+from unittest.mock import MagicMock
+from enum import Enum
+
+import pytest
+
+from snowflake.connector.errors import DatabaseError
+from snowflake.snowpark.async_job import AsyncJob
+from snowflake.snowpark.session import Session
+
+
+class CancelQueryErrorType(Enum):
+    """This is test-only enum to indicate what error we should mock for."""
+
+    VALID_JSON_BUT_FAIL = 0
+    INVALID_JSON = 1
+
+
+def _make_session_with_cursor() -> tuple[Session, mock.MagicMock]:
+    server_conn = MagicMock()
+    inner_conn = mock.MagicMock()
+    mock_cursor = mock.MagicMock()
+    inner_conn.cursor.return_value = mock_cursor
+    server_conn._conn = inner_conn
+
+    # Return True for async feature flag.
+    def _param_side_effect(name, default_value):
+        if name and name.upper() == "ENABLE_ASYNC_QUERY_IN_PYTHON_STORED_PROCS":
+            return True
+        return default_value
+
+    server_conn._get_client_side_session_parameter.side_effect = _param_side_effect
+    session = Session(server_conn)
+    return session, mock_cursor
+
+
+def test_async_job_cancel_executes_sys_func_in_regular_client(monkeypatch):
+    session, mock_cursor = _make_session_with_cursor()
+
+    # Test with regular client / non-sproc path.
+    monkeypatch.setattr(
+        "snowflake.snowpark.async_job.is_in_stored_procedure", lambda: False
+    )
+
+    qid = str(uuid.uuid4())
+
+    job = AsyncJob(query_id=qid, query=None, session=session)
+    job.cancel()
+
+    mock_cursor.execute.assert_called_once_with(f"select SYSTEM$CANCEL_QUERY('{qid}')")
+
+
+def test_async_job_cancel_in_sproc_success(monkeypatch):
+    session, mock_cursor = _make_session_with_cursor()
+
+    # Test with sproc path.
+    monkeypatch.setattr(
+        "snowflake.snowpark.async_job.is_in_stored_procedure", lambda: True
+    )
+
+    # Inject a fake _snowflake module with a mock cancel_query.
+    cancel_mock = mock.MagicMock(return_value='{"success": true}')
+    fake_mod = SimpleNamespace(cancel_query=cancel_mock)
+    monkeypatch.setitem(sys.modules, "_snowflake", fake_mod)
+
+    qid = str(uuid.uuid4())
+    job = AsyncJob(query_id=qid, query=None, session=session)
+    job.cancel()  # should not raise
+
+    # Verify that _snowflake.cancel_query is called with given query id.
+    cancel_mock.assert_called_once_with(qid)
+    # Verify that the system function is not called.
+    mock_cursor.execute.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "fake_response_and_type",
+    [
+        # Valid JSON, but success=false.
+        (
+            '{"success": false, "error": "boom"}',
+            CancelQueryErrorType.VALID_JSON_BUT_FAIL,
+        ),
+        # Invalid JSON.
+        ("{", CancelQueryErrorType.INVALID_JSON),
+    ],
+)
+def test_async_job_cancel_in_sproc_failure_in_cancel_query(
+    monkeypatch, fake_response_and_type
+):
+    session, _ = _make_session_with_cursor()
+
+    # Test with sproc path.
+    monkeypatch.setattr(
+        "snowflake.snowpark.async_job.is_in_stored_procedure", lambda: True
+    )
+
+    # Mock _snowflake.cancel_query to return error response as instructed to.
+    fake_response, error_type = fake_response_and_type
+    fake_mod = SimpleNamespace(cancel_query=lambda qid: fake_response)
+    monkeypatch.setitem(sys.modules, "_snowflake", fake_mod)
+
+    job = AsyncJob(query_id=str(uuid.uuid4()), query=None, session=session)
+    with pytest.raises(DatabaseError, match="Failed to cancel query") as exc_info:
+        job.cancel()
+    if error_type == CancelQueryErrorType.VALID_JSON_BUT_FAIL:
+        assert "Error parsing response" not in str(exc_info.value)
+    elif error_type == CancelQueryErrorType.INVALID_JSON:
+        assert "Error parsing response" in str(exc_info.value)
+    else:
+        raise ValueError(f"Invalid test case: {fake_response_and_type}")
+
+
+def test_async_job_cancel_in_sproc_failure_in_uuid_validation(monkeypatch):
+    session, _ = _make_session_with_cursor()
+
+    # Test with sproc path.
+    monkeypatch.setattr(
+        "snowflake.snowpark.async_job.is_in_stored_procedure", lambda: True
+    )
+
+    # Inject a fake _snowflake module with a mock cancel_query.
+    cancel_mock = mock.MagicMock(return_value='{"success": true}')
+    fake_mod = SimpleNamespace(cancel_query=cancel_mock)
+    monkeypatch.setitem(sys.modules, "_snowflake", fake_mod)
+
+    qid = "qid-invalid-123"
+    job = AsyncJob(query_id=qid, query=None, session=session)
+    with pytest.raises(ValueError, match=f"Invalid UUID: '{qid}'"):
+        job.cancel()