snowflake-connector-python 5.2.0__cp315-cp315-win_amd64.whl

snowflake/connector/aio/cursor/_dict_cursor.py

@@ -0,0 +1,78 @@
+"""Concrete DictCursor — returns dict rows."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from ..._internal.cursor import DictRow
+from ..._internal.errorhandler import simplified_error_handling
+from ._base import SnowflakeCursorBase
+
+
+class DictCursor(SnowflakeCursorBase):
+    """Cursor returning results as dictionaries with column names as keys.
+
+    Usage::
+
+        async with connection.cursor(DictCursor) as cur:
+            await cur.execute("SELECT 1 AS id, 'hello' AS name")
+            row = await cur.fetchone()
+            # row == {"ID": 1, "NAME": "hello"}
+    """
+
+    @property
+    def _use_dict_result(self) -> bool:
+        return True
+
+    @simplified_error_handling
+    async def fetchone(self) -> DictRow | None:
+        """
+        Fetch the next row of a query result set as a dictionary.
+
+        Returns:
+            dict: Next row as a dictionary with column names as keys,
+                  or None when no more data is available
+        """
+        return cast(DictRow | None, await self._fetchone())
+
+    @simplified_error_handling
+    async def fetchmany(self, size: int | None = None) -> list[DictRow]:
+        """
+        Fetch the next set of rows as dictionaries.
+
+        Args:
+            size (int): Number of rows to fetch (defaults to arraysize)
+
+        Returns:
+            list[dict]: List of rows as dictionaries
+        """
+        return await super().fetchmany(size)
+
+    async def fetchall(self) -> list[DictRow]:
+        """
+        Fetch all (remaining) rows as dictionaries.
+
+        Returns:
+            list[dict]: List of all remaining rows as dictionaries
+        """
+        return await super().fetchall()
+
+    async def __anext__(self) -> DictRow:
+        """
+        Fetch the next row from the currently executed statement as a dictionary.
+
+        Returns:
+            dict: Next row
+
+        Raises:
+            StopAsyncIteration: When no more rows are available
+        """
+        return cast(DictRow, await super().__anext__())
+
+    async def next(self) -> DictRow:
+        """Python 2 compatibility alias for ``__anext__``. Returns a dictionary row.
+
+        Raises:
+            StopAsyncIteration: When no more rows are available
+        """
+        return await self.__anext__()

snowflake/connector/aio/cursor/_result_set_wrapper.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from ..._internal.api_client.client_api import async_core_driver
+from ..._internal.cursor.result_set_wrapper import ResultSetWrapperBase
+from ..._internal.protobuf_gen.database_driver_v1_pb2 import (
+    ResultSetGetChunksResponse,
+)
+from ..._internal.statement_utils import get_stream_ptr
+
+
+class _ResultSetWrapper(ResultSetWrapperBase):
+    """Async result-set handle owner."""
+
+    async def get_arrow_stream_ptr(self) -> int:
+        """Fetch the Arrow stream and return the raw C pointer.
+
+        Each call builds a fresh stream from the stored result set data.
+
+        Raises:
+            ProgrammingError: If no handle is held.
+        """
+        handle = self._require_handle()
+        response = await async_core_driver.result_set_get_stream(result_set_handle=handle)
+        return get_stream_ptr(response)
+
+    async def get_chunks(self) -> ResultSetGetChunksResponse | None:
+        """Call ``result_set_get_chunks`` on the held handle, or return ``None``.
+
+        Safe to call multiple times and independently of :meth:`get_arrow_stream_ptr` —
+        the core layer re-fetches chunk metadata from Snowflake if the cached copy
+        was already consumed.
+        """
+        if self._handle is None:
+            return None
+        return await async_core_driver.result_set_get_chunks(result_set_handle=self._handle)

snowflake/connector/aio/cursor/_snowflake_cursor.py

@@ -0,0 +1,72 @@
+"""Concrete SnowflakeCursor — returns tuple rows."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from ..._internal.cursor import Row
+from ..._internal.errorhandler import simplified_error_handling
+from ._base import SnowflakeCursorBase
+
+
+class SnowflakeCursor(SnowflakeCursorBase):
+    """Cursor returning results as tuples (default).
+
+    This is the standard cursor returned by ``connection.cursor()``.
+    """
+
+    @property
+    def _use_dict_result(self) -> bool:
+        return False
+
+    @simplified_error_handling
+    async def fetchone(self) -> Row | None:
+        """
+        Fetch the next row of a query result set.
+
+        Returns:
+            tuple: Next row, or None when no more data is available
+        """
+        return cast(Row | None, await self._fetchone())
+
+    @simplified_error_handling
+    async def fetchmany(self, size: int | None = None) -> list[Row]:
+        """
+        Fetch the next set of rows of a query result.
+
+        Args:
+            size (int): Number of rows to fetch (defaults to arraysize)
+
+        Returns:
+            list[tuple]: List of rows as tuples
+        """
+        return await super().fetchmany(size)
+
+    async def fetchall(self) -> list[Row]:
+        """
+        Fetch all (remaining) rows of a query result.
+
+        Returns:
+            list[tuple]: List of all remaining rows as tuples
+        """
+        return await super().fetchall()
+
+    async def __anext__(self) -> Row:
+        """
+        Fetch the next row from the currently executed statement.
+
+        Returns:
+            tuple: Next row
+
+        Raises:
+            StopAsyncIteration: When no more rows are available
+        """
+        return cast(Row, await super().__anext__())
+
+    async def next(self) -> Row:
+        """Python 2 compatibility alias for ``__anext__``. Returns a tuple row.
+
+        Raises:
+            StopAsyncIteration: When no more rows are available
+        """
+        return await self.__anext__()

snowflake/connector/aio/pandas_tools.py

@@ -0,0 +1,279 @@
+"""Async write_pandas for snowflake.connector.aio.
+
+The sync write_pandas in snowflake.connector.pandas_tools is unchanged.
+This module adds an aio-native write_pandas that runs the full
+DDL → PUT → INFER_SCHEMA → COPY INTO pipeline on AsyncConnection /
+AsyncSnowflakeCursor.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+from .._common.extras import pandas, requires_dependency
+from .._internal.decorators import api_telemetry
+from .._internal.errorhandler import route_exception
+from .._internal.logging import get_logger
+from .._internal.write_pandas_operation import (
+    WritePandasConfig,
+    WritePandasMixin,
+    WritePandasResult,
+    generate_temp_name,
+)
+from ..errors import Error, ProgrammingError
+from .cursor import SnowflakeCursor as AsyncSnowflakeCursor
+
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+    from .connection import Connection
+
+logger = get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# _AsyncWritePandasOperation
+# ---------------------------------------------------------------------------
+
+
+class _AsyncWritePandasOperation(WritePandasMixin):
+    """Async-native write_pandas pipeline from a validated WritePandasConfig.
+
+    execute() runs the full pipeline without mutating self — all intermediate
+    state is threaded through local variables and return values.
+    """
+
+    def __init__(self, cfg: WritePandasConfig) -> None:
+        self._cfg = cfg
+
+    # -- Object lifetime -----------------------------------------------------
+
+    async def _drop_object(self, cur: AsyncSnowflakeCursor, name: str, object_type: str) -> None:
+        await cur.execute(**self._build_drop_object_sql(name, object_type))
+
+    async def _create_temp_object(
+        self,
+        cur: AsyncSnowflakeCursor,
+        sql_builder: Callable[[str], dict[str, Any]],
+        qualified_name: str,
+        bare_name: str,
+    ) -> str:
+        """Try creating in target schema; fall back to current schema on privilege error."""
+        try:
+            await cur.execute(**sql_builder(qualified_name))
+            return qualified_name
+        except ProgrammingError:
+            await cur.execute(**sql_builder(bare_name))
+            return bare_name
+
+    # -- Main pipeline -------------------------------------------------------
+
+    async def execute(self) -> WritePandasResult:
+        cfg = self._cfg
+        cfg.emit_warnings()
+
+        # This pipeline runs on the async connection. ``cfg.conn`` also allows the sync one.
+        cur = cast("Connection", cfg.conn).cursor(AsyncSnowflakeCursor)
+        target_location: str | None = None
+        try:
+            stage_location = await self._create_stage(cur)
+            nchunks, nrows = await self._upload_to_stage(cur, stage_location)
+
+            column_type_map: dict[str, str] | None = None
+            if cfg.needs_inference:
+                file_format_location = await self._create_file_format(cur)
+                column_type_map = await self._infer_column_types(cur, stage_location, file_format_location)
+
+            target_location = self._resolve_target_table()
+
+            if cfg.needs_table_creation:
+                await self._create_table(cur, target_location, column_type_map)
+
+            if cfg.needs_truncate:
+                await self._truncate_table(cur, target_location)
+
+            copy_results = await self._copy_into(cur, stage_location, target_location, column_type_map)
+
+            if cfg.needs_swap:
+                await self._swap_tables(cur, target_location)
+
+        except Exception:
+            if cfg.needs_swap and target_location:
+                try:
+                    await self._drop_object(cur, target_location, "TABLE")
+                except Exception:
+                    logger.warning(
+                        "aio.write_pandas failed and could not drop temporary "
+                        "staging table %s created for overwrite swap — "
+                        "it may need to be dropped manually",
+                        target_location,
+                    )
+            raise
+        finally:
+            cur.close()
+
+        success = all(row[1] == "LOADED" for row in copy_results)
+        return WritePandasResult(success, nchunks, nrows, copy_results)
+
+    # -- Stage & upload ------------------------------------------------------
+
+    async def _create_stage(self, cur: AsyncSnowflakeCursor) -> str:
+        cfg = self._cfg
+        name = generate_temp_name("STAGE")
+        qualified = cfg.qualify(name)
+        return await self._create_temp_object(cur, self._build_create_stage_sql, qualified, name)
+
+    async def _upload_to_stage(self, cur: AsyncSnowflakeCursor, stage_location: str) -> tuple[int, int]:
+        cfg = self._cfg
+        nchunks = 0
+        nrows = 0
+        with TemporaryDirectory() as tmp_dir:
+            for idx, chunk in self._iter_chunks():
+                chunk_path = Path(tmp_dir) / f"file{idx}.txt"
+                chunk.to_parquet(path=chunk_path, compression=cfg.compression, **cfg.parquet_kwargs)
+                nchunks += 1
+                nrows += len(chunk)
+
+                if not cfg.bulk_upload:
+                    await self._put_file(cur, stage_location, chunk_path)
+                    chunk_path.unlink()
+
+            if cfg.bulk_upload:
+                await self._put_directory(cur, stage_location, tmp_dir)
+
+        return nchunks, nrows
+
+    async def _put_file(self, cur: AsyncSnowflakeCursor, stage_location: str, path: Path) -> None:
+        await cur.execute(**self._build_put_file_sql(stage_location, path))
+
+    async def _put_directory(self, cur: AsyncSnowflakeCursor, stage_location: str, directory: str) -> None:
+        await cur.execute(**self._build_put_directory_sql(stage_location, directory))
+
+    # -- Schema inference ----------------------------------------------------
+
+    async def _create_file_format(self, cur: AsyncSnowflakeCursor) -> str:
+        cfg = self._cfg
+        name = generate_temp_name("FILE_FORMAT")
+        qualified = cfg.qualify(name)
+        return await self._create_temp_object(cur, self._build_create_file_format_sql, qualified, name)
+
+    async def _infer_column_types(
+        self,
+        cur: AsyncSnowflakeCursor,
+        stage_location: str,
+        file_format_location: str,
+    ) -> dict[str, str]:
+        """Run INFER_SCHEMA and return {UPPER_COL_NAME: SQL_TYPE} mapping."""
+        await cur.execute(**self._build_infer_column_types_sql(stage_location, file_format_location))
+        rows = await cur.fetchall()
+        return {row[0].upper(): row[1] for row in rows}
+
+    # -- Table management ----------------------------------------------------
+
+    async def _create_table(
+        self,
+        cur: AsyncSnowflakeCursor,
+        target_location: str,
+        column_type_map: dict[str, str] | None,
+    ) -> None:
+        await cur.execute(**self._build_create_table_sql(target_location, column_type_map))
+
+    async def _truncate_table(self, cur: AsyncSnowflakeCursor, target_location: str) -> None:
+        await cur.execute(**self._build_truncate_table_sql(target_location))
+
+    async def _swap_tables(self, cur: AsyncSnowflakeCursor, target_location: str) -> None:
+        """Replace original table with temp target via DROP + RENAME.
+
+        ALTER TABLE SWAP WITH would be atomic but doesn't work when the
+        original table doesn't exist yet or tables are different types
+        (TEMPORARY vs permanent).
+        """
+        original = self._cfg.qualify(self._cfg.table_name)
+        await self._drop_object(cur, original, "TABLE")
+        await cur.execute(**self._build_rename_table_sql(target_location))
+
+    # -- COPY INTO -----------------------------------------------------------
+
+    async def _copy_into(
+        self,
+        cur: AsyncSnowflakeCursor,
+        stage_location: str,
+        target_location: str,
+        column_type_map: dict[str, str] | None,
+    ) -> list:
+        await cur.execute(**self._build_copy_into_sql(stage_location, target_location, column_type_map))
+        return await cur.fetchall()
+
+
+# ---------------------------------------------------------------------------
+# Public async entry point
+# ---------------------------------------------------------------------------
+
+
+@requires_dependency(pandas)
+@api_telemetry
+async def write_pandas(
+    conn: Connection,
+    df: DataFrame,
+    table_name: str,
+    database: str | None = None,
+    schema: str | None = None,
+    chunk_size: int | None = None,
+    compression: Literal["gzip", "snappy"] = "gzip",
+    on_error: str = "abort_statement",
+    parallel: int = 4,
+    quote_identifiers: bool = True,
+    infer_schema: bool = False,
+    auto_create_table: bool = False,
+    overwrite: bool = False,
+    table_type: Literal["", "temp", "temporary", "transient"] = "",
+    use_logical_type: bool | None = None,
+    iceberg_config: dict[str, str] | None = None,
+    bulk_upload_chunks: bool = False,
+    use_vectorized_scanner: bool = False,
+    **kwargs: Any,
+) -> WritePandasResult:
+    """Write a pandas DataFrame to a Snowflake table via Parquet stage upload (async).
+
+    Async-native counterpart of :func:`snowflake.connector.pandas_tools.write_pandas`.
+    Requires an :class:`~snowflake.connector.aio.Connection`.
+
+    Returns a WritePandasResult named tuple (success, nchunks, nrows, copy_results).
+    Backward-compatible with plain tuple unpacking and indexing.
+
+    Note:
+        Unlike the sync module, this module does not expose ``pd_writer`` or
+        ``make_pd_writer``.  ``DataFrame.to_sql(method=...)`` invokes the method
+        callback synchronously and cannot await a coroutine, so an async adapter
+        for ``to_sql`` is not possible.  Use ``pd_writer`` from
+        :mod:`snowflake.connector.pandas_tools` if you need ``to_sql`` integration.
+    """
+    try:
+        cfg = WritePandasConfig(
+            conn,
+            df,
+            table_name,
+            database=database,
+            schema=schema,
+            chunk_size=chunk_size,
+            compression=compression,
+            on_error=on_error,
+            parallel=parallel,
+            quote_identifiers=quote_identifiers,
+            infer_schema=infer_schema,
+            auto_create_table=auto_create_table,
+            overwrite=overwrite,
+            table_type=table_type,
+            use_logical_type=use_logical_type,
+            iceberg_config=iceberg_config,
+            bulk_upload_chunks=bulk_upload_chunks,
+            use_vectorized_scanner=use_vectorized_scanner,
+            **kwargs,
+        )
+        return await _AsyncWritePandasOperation(cfg).execute()
+    except Error as exc:
+        route_exception(conn, None, exc)

snowflake/connector/aio/result_batch.py

@@ -0,0 +1,165 @@
+"""Async result batch classes for distributed fetch."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any, cast
+
+from .._common.extras import pandas, pyarrow, requires_dependency
+from .._internal.api_client.client_api import async_core_driver
+from .._internal.arrow_stream_async import (
+    collect_arrow_table_async,
+    create_async_row_iterator_from_stream_ptr,
+    to_pandas_async,
+)
+from .._internal.arrow_stream_utils import (
+    create_table_iterator,
+)
+from .._internal.protobuf_gen.database_driver_v1_pb2 import ColumnMetadata, ResultChunk
+from .._internal.result_batch import IterTableStructure, IterUnit, ResultBatchMixin
+from .._internal.statement_utils import get_stream_ptr
+
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+    from pyarrow import Table
+
+    from ..cursor import ResultMetadata
+    from .connection import Connection
+
+
+class ResultBatch(ResultBatchMixin):
+    """Async counterpart of :class:`~snowflake.connector.result_batch.ResultBatch`."""
+
+    _connection: Connection | None
+
+    def __init__(
+        self,
+        chunk: ResultChunk,
+        description: list[ResultMetadata],
+        connection: Connection | None,
+        columns: list[ColumnMetadata] | None = None,
+    ) -> None:
+        super().__init__(chunk, description, connection, columns)
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def connection(self) -> Connection | None:
+        return self._connection
+
+    @connection.setter
+    def connection(self, value: Connection | None) -> None:
+        self._connection = value
+
+    # ------------------------------------------------------------------
+    # Connection resolution
+    # ------------------------------------------------------------------
+
+    def _resolve_connection(self, connection: Connection | None = None) -> Connection | None:
+        return connection or self._connection
+
+    # ------------------------------------------------------------------
+    # Data fetching
+    # ------------------------------------------------------------------
+
+    async def _fetch_arrow_stream_ptr(self, connection: Connection | None = None) -> int:
+        conn_handle = connection.conn_handle if connection is not None else None
+        response = await async_core_driver.database_fetch_chunk(
+            conn_handle=conn_handle,
+            chunk=self._chunk,
+            columns=self._columns,
+        )
+        return get_stream_ptr(response)
+
+    async def _take_arrow_stream_ptr(self, connection: Connection | None = None) -> int:
+        if self._arrow_stream_ptr is None:
+            await self.populate_data(connection=connection)
+        stream_ptr = cast(int, self._arrow_stream_ptr)
+        self._arrow_stream_ptr = None
+        return stream_ptr
+
+    async def populate_data(self, connection: Connection | None = None, **kwargs: Any) -> ResultBatch:
+        conn = self._resolve_connection(connection)
+        self._arrow_stream_ptr = await self._fetch_arrow_stream_ptr(conn)
+        return self
+
+    def __aiter__(self) -> AsyncIterator[tuple | dict | Exception]:
+        return self.create_iter()
+
+    async def create_iter(
+        self,
+        connection: Connection | None = None,
+        iter_unit: IterUnit | str = IterUnit.ROW_UNIT,
+        structure: IterTableStructure | str = IterTableStructure.PANDAS,
+        use_dict_result: bool = False,
+        number_to_decimal: bool = False,
+        force_microsecond_precision: bool = False,
+    ) -> AsyncIterator[tuple | dict | Exception] | AsyncIterator[Table] | AsyncIterator[DataFrame]:
+        iter_unit = IterUnit.of(iter_unit)
+        structure = IterTableStructure.of(structure)
+
+        conn = self._resolve_connection(connection)
+        if iter_unit == IterUnit.TABLE_UNIT:
+            if structure == IterTableStructure.PANDAS:
+                yield await self.to_pandas(
+                    connection=conn,
+                    number_to_decimal=number_to_decimal,
+                    force_microsecond_precision=force_microsecond_precision,
+                )
+                return
+            yield await self.to_arrow(
+                connection=conn,
+                number_to_decimal=number_to_decimal,
+                force_microsecond_precision=force_microsecond_precision,
+            )
+            return
+
+        iterator = create_async_row_iterator_from_stream_ptr(
+            await self._take_arrow_stream_ptr(conn),
+            context=self._arrow_context,
+            use_dict_result=use_dict_result,
+            use_numpy=self._numpy,
+            use_core_arrow=self._use_core_arrow,
+        )
+        async for row in iterator:
+            yield row
+
+    @requires_dependency(pyarrow)
+    async def to_arrow(
+        self,
+        connection: Connection | None = None,
+        number_to_decimal: bool = False,
+        force_microsecond_precision: bool = False,
+    ) -> Table:
+        conn = self._resolve_connection(connection)
+        stream_ptr = await self._take_arrow_stream_ptr(conn)
+        return await collect_arrow_table_async(
+            create_table_iterator(
+                stream_ptr,
+                context=self._arrow_context,
+                number_to_decimal=number_to_decimal,
+                force_microsecond_precision=force_microsecond_precision,
+                use_core_arrow=self._use_core_arrow,
+            ),
+            self._description,
+        )
+
+    @requires_dependency(pandas)
+    async def to_pandas(
+        self,
+        connection: Connection | None = None,
+        number_to_decimal: bool = False,
+        force_microsecond_precision: bool = False,
+    ) -> DataFrame:
+        table = await self.to_arrow(
+            connection=connection,
+            number_to_decimal=number_to_decimal,
+            force_microsecond_precision=force_microsecond_precision,
+        )
+        return await to_pandas_async(table)
+
+
+__all__ = ["IterUnit", "IterTableStructure", "ResultBatch"]

snowflake/connector/auth/workload_identity.py

@@ -0,0 +1,5 @@
+"""BACKWARD COMPATIBILITY MODULE ONLY"""
+
+
+class ApiFederatedAuthenticationType:
+    pass

snowflake/connector/compat.py

@@ -0,0 +1,22 @@
+"""BACKWARD COMPATIBILITY MODULE ONLY"""
+
+import platform
+
+from typing import Any
+
+from ._internal.decorators import backward_compatibility
+
+
+IS_LINUX = platform.system() == "Linux"
+IS_WINDOWS = platform.system() == "Windows"
+IS_MACOS = platform.system() == "Darwin"
+
+
+@backward_compatibility
+def IS_UNICODE(v: Any) -> bool:
+    """To check whether v is a Unicode string."""
+    return isinstance(v, str)
+
+
+# backward compatibility
+IS_STR = IS_UNICODE