Skip to content

Commit 96d8944

Browse files
Polandia94sfc-gh-jszczerbinski
authored andcommitted
add tupleCursor and type dictCursor
1 parent 598ab4c commit 96d8944

2 files changed

Lines changed: 30 additions & 4 deletions

File tree

src/snowflake/connector/connection.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from threading import Lock
2626
from time import strptime
2727
from types import TracebackType
28-
from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence
28+
from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence, TypeVar
2929
from uuid import UUID
3030

3131
from cryptography.hazmat.backends import default_backend
@@ -115,6 +115,7 @@
115115
DEFAULT_CLIENT_PREFETCH_THREADS = 4
116116
MAX_CLIENT_PREFETCH_THREADS = 10
117117
DEFAULT_BACKOFF_POLICY = exponential_backoff()
118+
T = TypeVar('T', bound=SnowflakeCursor)
118119

119120

120121
def DefaultConverterClass() -> type:
@@ -869,8 +870,8 @@ def rollback(self) -> None:
869870
self.cursor().execute("ROLLBACK")
870871

871872
def cursor(
872-
self, cursor_class: type[SnowflakeCursor] = SnowflakeCursor
873-
) -> SnowflakeCursor:
873+
self, cursor_class: type[T] = SnowflakeCursor
874+
) -> T:
874875
"""Creates a cursor object. Each statement will be executed in a new cursor object."""
875876
logger.debug("cursor")
876877
if not self.rest:

src/snowflake/connector/cursor.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1728,13 +1728,38 @@ def get_result_batches(self) -> list[ResultBatch] | None:
17281728

17291729
class DictCursor(SnowflakeCursor):
17301730
"""Cursor returning results in a dictionary."""
1731-
17321731
def __init__(self, connection) -> None:
17331732
super().__init__(
17341733
connection,
17351734
use_dict_result=True,
17361735
)
17371736

1737+
def fetchone(self) -> dict | None:
1738+
return super().fetchone()
1739+
1740+
def fetchmany(self, size: int | None = None) -> list[dict]:
1741+
return super().fetchmany()
1742+
1743+
def fetchall(self) -> list[dict]:
1744+
return super().fetchall()
1745+
1746+
class TupleCursor(SnowflakeCursor):
1747+
"""Cursor returning results in a dictionary."""
1748+
1749+
def __init__(self, connection) -> None:
1750+
super().__init__(
1751+
connection,
1752+
use_dict_result=False,
1753+
)
1754+
1755+
def fetchone(self) -> tuple | None:
1756+
return super().fetchone()
1757+
1758+
def fetchmany(self, size: int | None = None) -> list[tuple]:
1759+
return super().fetchmany()
1760+
1761+
def fetchall(self) -> list[tuple]:
1762+
return super().fetchall()
17381763

17391764
def __getattr__(name):
17401765
if name == "NanoarrowUsage":

0 commit comments

Comments
 (0)