Skip to content

Commit

Permalink
add tupleCursor and type dictCursor
Browse files Browse the repository at this point in the history
  • Loading branch information
Polandia94 authored and Pablo Nicolas Estevez committed Oct 25, 2024
1 parent 9f9c71e commit 7f9d76b
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 4 deletions.
7 changes: 4 additions & 3 deletions src/snowflake/connector/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from threading import Lock
from time import strptime
from types import TracebackType
from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence
from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence, TypeVar
from uuid import UUID

from cryptography.hazmat.backends import default_backend
Expand Down Expand Up @@ -113,6 +113,7 @@
DEFAULT_CLIENT_PREFETCH_THREADS = 4
MAX_CLIENT_PREFETCH_THREADS = 10
DEFAULT_BACKOFF_POLICY = exponential_backoff()
T = TypeVar('T', bound=SnowflakeCursor)


def DefaultConverterClass() -> type:
Expand Down Expand Up @@ -855,8 +856,8 @@ def rollback(self) -> None:
self.cursor().execute("ROLLBACK")

def cursor(
self, cursor_class: type[SnowflakeCursor] = SnowflakeCursor
) -> SnowflakeCursor:
self, cursor_class: type[T] = SnowflakeCursor
) -> T:
"""Creates a cursor object. Each statement will be executed in a new cursor object."""
logger.debug("cursor")
if not self.rest:
Expand Down
27 changes: 26 additions & 1 deletion src/snowflake/connector/cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -1724,13 +1724,38 @@ def get_result_batches(self) -> list[ResultBatch] | None:

class DictCursor(SnowflakeCursor):
"""Cursor returning results in a dictionary."""

def __init__(self, connection) -> None:
super().__init__(
connection,
use_dict_result=True,
)

def fetchone(self) -> dict | None:
return super().fetchone()

def fetchmany(self, size: int | None = None) -> list[dict]:
return super().fetchmany()

def fetchall(self) -> list[dict]:
return super().fetchall()

class TupleCursor(SnowflakeCursor):
"""Cursor returning results in a dictionary."""

def __init__(self, connection) -> None:
super().__init__(
connection,
use_dict_result=False,
)

def fetchone(self) -> tuple | None:
return super().fetchone()

def fetchmany(self, size: int | None = None) -> list[tuple]:
return super().fetchmany()

def fetchall(self) -> list[tuple]:
return super().fetchall()

def __getattr__(name):
if name == "NanoarrowUsage":
Expand Down

0 comments on commit 7f9d76b

Please sign in to comment.