11#!/usr/bin/env python
22from __future__ import annotations
33
4+ import abc
45import collections
56import logging
67import os
1920 TYPE_CHECKING ,
2021 Any ,
2122 Callable ,
23+ Generic ,
2224 Iterator ,
2325 Literal ,
2426 NamedTuple ,
8688 from .result_batch import ResultBatch
8789
8890T = TypeVar ("T" , bound = collections .abc .Sequence )
91+ FetchRow = TypeVar ("FetchRow" , bound = tuple [Any , ...] | dict [str , Any ])
8992
9093logger = getLogger (__name__ )
9194
@@ -332,29 +335,7 @@ class ResultState(Enum):
332335 RESET = 3
333336
334337
335- class SnowflakeCursor :
336- """Implementation of Cursor object that is returned from Connection.cursor() method.
337-
338- Attributes:
339- description: A list of namedtuples about metadata for all columns.
340- rowcount: The number of records updated or selected. If not clear, -1 is returned.
341- rownumber: The current 0-based index of the cursor in the result set or None if the index cannot be
342- determined.
343- sfqid: Snowflake query id in UUID form. Include this in the problem report to the customer support.
344- sqlstate: Snowflake SQL State code.
345- timestamp_output_format: Snowflake timestamp_output_format for timestamps.
346- timestamp_ltz_output_format: Snowflake output format for LTZ timestamps.
347- timestamp_tz_output_format: Snowflake output format for TZ timestamps.
348- timestamp_ntz_output_format: Snowflake output format for NTZ timestamps.
349- date_output_format: Snowflake output format for dates.
350- time_output_format: Snowflake output format for times.
351- timezone: Snowflake timezone.
352- binary_output_format: Snowflake output format for binary fields.
353- arraysize: The default number of rows fetched by fetchmany.
354- connection: The connection object by which the cursor was created.
355- errorhandle: The class that handles error handling.
356- is_file_transfer: Whether, or not the current command is a put, or get.
357- """
338+ class SnowflakeCursorBase (abc .ABC , Generic [FetchRow ]):
358339
359340 # TODO:
360341 # Most of these attributes have no reason to be properties, we could just store them in public variables.
@@ -1514,8 +1495,17 @@ def executemany(
15141495
15151496 return self
15161497
1517- def fetchone (self ) -> dict | tuple | None :
1518- """Fetches one row."""
1498+ @abc .abstractmethod
1499+ def fetchone (self ) -> FetchRow :
1500+ pass
1501+
1502+ def _fetchone (self ) -> dict [str , Any ] | tuple [Any , ...] | None :
1503+ """
1504+ Fetches one row.
1505+
1506+ Returns a dict if self._use_dict_result is True, otherwise
1507+ returns tuple.
1508+ """
15191509 if self ._prefetch_hook is not None :
15201510 self ._prefetch_hook ()
15211511 if self ._result is None and self ._result_set is not None :
@@ -1539,7 +1529,7 @@ def fetchone(self) -> dict | tuple | None:
15391529 else :
15401530 return None
15411531
1542- def fetchmany (self , size : int | None = None ) -> list [tuple ] | list [ dict ]:
1532+ def fetchmany (self , size : int | None = None ) -> list [FetchRow ]:
15431533 """Fetches the number of specified rows."""
15441534 if size is None :
15451535 size = self .arraysize
@@ -1565,7 +1555,7 @@ def fetchmany(self, size: int | None = None) -> list[tuple] | list[dict]:
15651555
15661556 return ret
15671557
1568- def fetchall (self ) -> list [tuple ] | list [ dict ]:
1558+ def fetchall (self ) -> list [FetchRow ]:
15691559 """Fetches all of the results."""
15701560 ret = []
15711561 while True :
@@ -1925,7 +1915,37 @@ def _create_file_transfer_agent(
19251915 )
19261916
19271917
1928- class DictCursor (SnowflakeCursor ):
1918+ class SnowflakeCursor (SnowflakeCursorBase [tuple [Any , ...]]):
1919+ """Implementation of Cursor object that is returned from Connection.cursor() method.
1920+
1921+ Attributes:
1922+ description: A list of namedtuples about metadata for all columns.
1923+ rowcount: The number of records updated or selected. If not clear, -1 is returned.
1924+ rownumber: The current 0-based index of the cursor in the result set or None if the index cannot be
1925+ determined.
1926+ sfqid: Snowflake query id in UUID form. Include this in the problem report to the customer support.
1927+ sqlstate: Snowflake SQL State code.
1928+ timestamp_output_format: Snowflake timestamp_output_format for timestamps.
1929+ timestamp_ltz_output_format: Snowflake output format for LTZ timestamps.
1930+ timestamp_tz_output_format: Snowflake output format for TZ timestamps.
1931+ timestamp_ntz_output_format: Snowflake output format for NTZ timestamps.
1932+ date_output_format: Snowflake output format for dates.
1933+ time_output_format: Snowflake output format for times.
1934+ timezone: Snowflake timezone.
1935+ binary_output_format: Snowflake output format for binary fields.
1936+ arraysize: The default number of rows fetched by fetchmany.
1937+ connection: The connection object by which the cursor was created.
1938+ errorhandle: The class that handles error handling.
1939+ is_file_transfer: Whether, or not the current command is a put, or get.
1940+ """
1941+
1942+ def fetchone (self ) -> tuple [Any , ...] | None :
1943+ row = self ._fetchone ()
1944+ assert row is None or isinstance (row , tuple )
1945+ return row
1946+
1947+
1948+ class DictCursor (SnowflakeCursorBase [dict [str , Any ]]):
19291949 """Cursor returning results in a dictionary."""
19301950
19311951 def __init__ (self , connection ) -> None :
@@ -1934,6 +1954,11 @@ def __init__(self, connection) -> None:
19341954 use_dict_result = True ,
19351955 )
19361956
1957+ def fetchone (self ) -> dict [str , Any ] | None :
1958+ row = self ._fetchone ()
1959+ assert row is None or isinstance (row , dict )
1960+ return row
1961+
19371962
19381963def __getattr__ (name ):
19391964 if name == "NanoarrowUsage" :
0 commit comments