SNOW-3264599: Python - implement cursor reset method

Author: sfc-gh-asolarski

python/src/snowflake/connector/_internal/query_utils.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from .protobuf_gen.database_driver_v1_pb2 import (
+    ExecuteResult,
+    StatementHandle,
+    StatementNewRequest,
+    StatementReleaseRequest,
+    StatementSetSqlQueryRequest,
+)
+
+
+if TYPE_CHECKING:
+    from ..connection import Connection
+
+
+@contextmanager
+def create_statement(connection: Connection, query: str) -> Generator[StatementHandle]:
+    statement_request = StatementNewRequest(conn_handle=connection.conn_handle)
+    statement = connection.db_api.statement_new(request=statement_request)
+    stmt_handle = statement.stmt_handle
+    sql_query_request = StatementSetSqlQueryRequest(stmt_handle=stmt_handle, query=query)
+    try:
+        connection.db_api.statement_set_sql_query(sql_query_request)
+        yield stmt_handle
+    finally:
+        release_request = StatementReleaseRequest(stmt_handle=stmt_handle)
+        connection.db_api.statement_release(release_request)
+
+
+def extract_rowcount(result: ExecuteResult) -> int:
+    """Extract rowcount from execute result."""
+    if result and result.HasField("rows_affected"):
+        return result.rows_affected
+    return -1
+
+
+def extract_sqlstate(result: ExecuteResult | None) -> str | None:
+    """Extract sqlstate from execute result.
+
+    "00000" (successful completion) is treated as None for backwards compatibility with the old connector.
+    """
+    sql_state = result.sql_state if result else None
+    if sql_state and sql_state != "00000":
+        return sql_state
+    return None
+
+
+def get_stream_ptr(result: ExecuteResult | None) -> int:
+    """Get the ArrowArrayStream pointer from execute result.
+
+    Returns:
+        int: The ArrowArrayStream pointer as an integer
+
+    Raises:
+        RuntimeError: If result is invalid or stream pointer is null
+    """
+    if result is None:
+        raise RuntimeError("No query has been executed")
+
+    if not hasattr(result, "stream") or result.stream is None:
+        raise RuntimeError("Execute result does not contain a valid stream")
+
+    if not hasattr(result.stream, "value") or result.stream.value is None:
+        raise RuntimeError("Stream does not contain a valid pointer value")
+
+    stream_value = result.stream.value
+    if len(stream_value) != 8:
+        raise RuntimeError(f"Stream pointer value has wrong length: {len(stream_value)} (expected 8)")
+
+    stream_ptr = int.from_bytes(stream_value, byteorder="little", signed=False)
+
+    if stream_ptr == 0:
+        raise RuntimeError("Stream pointer is null")
+
+    return stream_ptr

python/src/snowflake/connector/cursor.py

@@ -35,10 +35,9 @@
     ExecuteResult,
     QueryBindings,
     StatementExecuteQueryRequest,
-    StatementNewRequest,
-    StatementReleaseRequest,
-    StatementSetSqlQueryRequest,
+    StatementHandle,
 )
+from ._internal.query_utils import create_statement, extract_rowcount, extract_sqlstate, get_stream_ptr
 from ._internal.type_codes import get_type_code
 from .errors import InterfaceError, NotSupportedError, ProgrammingError
 
@@ -93,6 +92,13 @@ def from_column(cls, col: Any) -> ResultMetadata:
             is_nullable=col.nullable,
         )
 
+    @classmethod
+    def create_description(cls, result: ExecuteResult | None) -> list[ResultMetadata] | None:
+        """Extract description from execute result column metadata."""
+        if result and result.columns:
+            return [cls.from_column(col) for col in result.columns]
+        return None
+
 
 # Backward compatibility alias
 ResultMetadataV2 = ResultMetadata
@@ -174,7 +180,7 @@ def __init__(self, connection: Connection) -> None:
         """
         self._connection = connection
         self._description: list[ResultMetadata] | None = None
-        self._rowcount: int | None = None
+        self._rowcount: int = -1
         self._arraysize: int = 1
         self._sqlstate: str | None = None
         self._closed = False
@@ -304,9 +310,20 @@ def callproc(self, procname: str, parameters: Sequence[Any] | None = None) -> Se
         raise NotSupportedError("callproc is not implemented")
 
     @pep249
-    def close(self) -> None:
-        """Close the cursor now (rather than whenever __del__ is called)."""
-        self._closed = True
+    def close(self) -> bool | None:
+        """Close the cursor now (rather than whenever __del__ is called).
+
+        Returns whether the cursor was closed during this call.
+        """
+        try:
+            if self._closed:
+                return False
+            self.reset(closing=True)
+            self._closed = True
+            del self._messages[:]
+            return True
+        except Exception:
+            return None
 
     def _build_query_bindings(self, parameters: Sequence[Any]) -> QueryBindings | None:
         """Serialize parameters and build a QueryBindings protobuf message.
@@ -397,6 +414,7 @@ def execute(
     ) -> SnowflakeCursorBase:
         """
         Execute a database operation (query or command).
+        Resets the cursor state before the execution.
 
         Args:
             operation (str): SQL statement to execute
@@ -405,52 +423,43 @@ def execute(
                 For pyformat paramstyle: sequence (%s) or dict (%(name)s)
                 For format paramstyle: sequence (%s)
         """
+        self.reset()
+        return self._execute(operation, parameters, _is_put_get, **kwargs)
+
+    def _execute(
+        self,
+        operation: str,
+        parameters: Sequence[Any] | dict[str, Any] | None = None,
+        _is_put_get: bool | None = None,
+        **kwargs: Any,
+    ) -> SnowflakeCursorBase:
+        """Execute query logic."""
         query, bindings = self._prepare_query(operation, parameters)
-        stmt_handle = self._connection.db_api.statement_new(
-            StatementNewRequest(conn_handle=self._connection.conn_handle)
-        ).stmt_handle
-        self._connection.db_api.statement_set_sql_query(
-            StatementSetSqlQueryRequest(stmt_handle=stmt_handle, query=query)
-        )
 
-        request = StatementExecuteQueryRequest(stmt_handle=stmt_handle, bindings=bindings)
+        result: ExecuteResult | None = None
+        with create_statement(self.connection, query) as stmt_handle:
+            result = self._execute_query(stmt_handle, bindings)
 
+        # populate description, rowcount, and sqlstate
+        self._description = ResultMetadata.create_description(result)
+        self._rowcount = extract_rowcount(result)
+        self._sqlstate = extract_sqlstate(result)
+        # save execute result
+        self.execute_result = result
+
+        return self
+
+    def _execute_query(self, stmt_handle: StatementHandle, bindings: QueryBindings | None) -> ExecuteResult:
         try:
-            self.execute_result = self._connection.db_api.statement_execute_query(request).result
+            request = StatementExecuteQueryRequest(stmt_handle=stmt_handle, bindings=bindings)
+            return self._connection.db_api.statement_execute_query(request).result
         except ProgrammingError as exc:
             self._sqlstate = exc.sqlstate or None
             raise
-        finally:
-            self._connection.db_api.statement_release(StatementReleaseRequest(stmt_handle=stmt_handle))
-
-        # Reset streaming state for a new result
-        self._binding_data = None
-        self._iterator = None
-        self._fetch_mode = None
-        self._rownumber = -1
-
-        # Populate description, rowcount, and sqlstate
-        self._populate_description()
-        self._populate_rowcount()
-        self._populate_sqlstate()
-        return self
-
-    def _populate_rowcount(self) -> None:
-        if self.execute_result and self.execute_result.HasField("rows_affected"):
-            self._rowcount = self.execute_result.rows_affected
-        else:
-            self._rowcount = None
-
-    def _populate_sqlstate(self) -> None:
-        # "00000" (successful completion) is treated as None for
-        # backwards compatibility with the old connector.
-        sql_state = self.execute_result.sql_state if self.execute_result else None
-        if sql_state and sql_state != "00000":
-            self._sqlstate = sql_state
-        else:
-            self._sqlstate = None
 
     @pep249
+    @_requires_not_closed
+    @_requires_open_connection
     def executemany(self, operation: str, seq_of_parameters: Sequence[Sequence[Any] | dict[str, Any]]) -> None:
         """
         Execute a database operation repeatedly for each element in seq_of_parameters.
@@ -476,10 +485,11 @@ def executemany(self, operation: str, seq_of_parameters: Sequence[Sequence[Any]
         # - Client-side binding (pyformat/format)
         # - Dict parameters (server-side doesn't support named binding)
         if paramstyle.is_client_side() or isinstance(first_params, dict):
+            self.reset()
             total_rowcount = 0
             unknown_rowcount = False
             for params in seq_of_parameters:
-                self.execute(operation, params)
+                self._execute(operation, params)  # no reset between calls
                 rc = self._rowcount
                 if rc is None or rc == -1:
                     unknown_rowcount = True
@@ -515,50 +525,8 @@ def executemany(self, operation: str, seq_of_parameters: Sequence[Sequence[Any]
     # Arrow stream helpers
     # ------------------------------------------------------------------
 
-    def _get_stream_ptr(self) -> int:
-        """Get the ArrowArrayStream pointer from execute result.
-
-        Returns:
-            int: The ArrowArrayStream pointer as an integer
-
-        Raises:
-            RuntimeError: If execute_result is invalid or stream pointer is null
-        """
-        if self.execute_result is None:
-            raise RuntimeError("No query has been executed")
-
-        if not hasattr(self.execute_result, "stream") or self.execute_result.stream is None:
-            raise RuntimeError("Execute result does not contain a valid stream")
-
-        if not hasattr(self.execute_result.stream, "value") or self.execute_result.stream.value is None:
-            raise RuntimeError("Stream does not contain a valid pointer value")
-
-        stream_value = self.execute_result.stream.value
-        if len(stream_value) != 8:
-            raise RuntimeError(f"Stream pointer value has wrong length: {len(stream_value)} (expected 8)")
-
-        stream_ptr = int.from_bytes(stream_value, byteorder="little", signed=False)
-
-        if stream_ptr == 0:
-            raise RuntimeError("Stream pointer is null")
-
-        return stream_ptr
-
-    def _populate_description(self) -> None:
-        """Populate cursor description from execute result column metadata."""
-        if self.execute_result is None:
-            self._description = None
-            return
-
-        columns = self.execute_result.columns
-        if not columns:
-            self._description = None
-            return
-
-        self._description = [ResultMetadata.from_column(col) for col in columns]
-
     def _get_iterator(self) -> ArrowStreamIterator:
-        stream_ptr = self._get_stream_ptr()
+        stream_ptr = get_stream_ptr(self.execute_result)
         arrow_context = ArrowConverterContext()
         return ArrowStreamIterator(
             stream_ptr,
@@ -572,7 +540,7 @@ def _get_table_iterator(
         self,
         force_microsecond_precision: bool = False,
     ) -> ArrowStreamTableIterator:
-        stream_ptr = self._get_stream_ptr()
+        stream_ptr = get_stream_ptr(self.execute_result)
         arrow_context = ArrowConverterContext()
         return ArrowStreamTableIterator(
             stream_ptr,
@@ -849,8 +817,27 @@ def scroll(self, value: int, mode: str = "relative") -> None:
         raise NotSupportedError("scroll is not supported")
 
     def reset(self, closing: bool = False) -> None:
-        """Reset the result set."""
-        raise NotImplementedError("reset is not yet implemented")
+        """Reset the result set.
+
+        Clears all result-related state, preparing the cursor for a new execution.
+        This is called automatically by execute() before each query, and by close().
+
+        Args:
+            closing: If True, do not reset rowcount,
+                     see: SNOW-647539: Do not erase the rowcount information when closing the cursor.
+                     If False, reset rowcount to None.
+        """
+        self._description = None
+        if not closing:
+            self._rowcount = -1
+        self._sqlstate = None
+
+        self.execute_result = None
+        self._iterator = None
+        self._fetch_mode = None
+
+        self._binding_data = None
+        self._rownumber = -1
 
     def query_result(self, qid: str) -> SnowflakeCursorBase:
         """Query the result of a previously executed query."""

python/tests/integ/test_cursor.py

@@ -427,16 +427,16 @@ def test_description_numeric_precision_and_scale(self, cursor):
 class TestCursorRowcount:
     """Integration tests for Cursor.rowcount property."""
 
-    def test_rowcount_is_none_before_execute(self, connection):
-        """Test that rowcount returns None before any query is executed."""
+    def test_rowcount_is_minus_one_before_execute(self, connection):
+        """Test that rowcount returns -1 before any query is executed."""
         # Given a new cursor
         cursor = connection.cursor()
 
         # When accessing rowcount before execute
         result = cursor.rowcount
 
-        # Then it should be None (per PEP 249)
-        assert result is None
+        # Then it should be -1 (per PEP 249)
+        assert result == -1
 
     def test_rowcount_after_select_single_row(self, cursor):
         """Test rowcount after a SELECT query returning single row."""