snowflake-connector-python 4.4.0__cp314-cp314-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- snowflake/connector/__init__.py +118 -0
- snowflake/connector/_query_context_cache.py +276 -0
- snowflake/connector/_sql_util.py +48 -0
- snowflake/connector/_utils.py +353 -0
- snowflake/connector/aio/README.md +88 -0
- snowflake/connector/aio/__init__.py +163 -0
- snowflake/connector/aio/_azure_storage_client.py +227 -0
- snowflake/connector/aio/_bind_upload_agent.py +69 -0
- snowflake/connector/aio/_connection.py +1216 -0
- snowflake/connector/aio/_cursor.py +1402 -0
- snowflake/connector/aio/_description.py +5 -0
- snowflake/connector/aio/_direct_file_operation_utils.py +88 -0
- snowflake/connector/aio/_file_transfer_agent.py +316 -0
- snowflake/connector/aio/_gcs_storage_client.py +360 -0
- snowflake/connector/aio/_network.py +860 -0
- snowflake/connector/aio/_ocsp_asn1crypto.py +51 -0
- snowflake/connector/aio/_ocsp_snowflake.py +608 -0
- snowflake/connector/aio/_pandas_tools.py +561 -0
- snowflake/connector/aio/_result_batch.py +442 -0
- snowflake/connector/aio/_result_set.py +309 -0
- snowflake/connector/aio/_s3_storage_client.py +460 -0
- snowflake/connector/aio/_session_manager.py +588 -0
- snowflake/connector/aio/_storage_client.py +331 -0
- snowflake/connector/aio/_telemetry.py +97 -0
- snowflake/connector/aio/_time_util.py +57 -0
- snowflake/connector/aio/_wif_util.py +324 -0
- snowflake/connector/aio/auth/__init__.py +52 -0
- snowflake/connector/aio/auth/_auth.py +406 -0
- snowflake/connector/aio/auth/_by_plugin.py +131 -0
- snowflake/connector/aio/auth/_default.py +28 -0
- snowflake/connector/aio/auth/_idtoken.py +54 -0
- snowflake/connector/aio/auth/_keypair.py +60 -0
- snowflake/connector/aio/auth/_no_auth.py +31 -0
- snowflake/connector/aio/auth/_oauth.py +27 -0
- snowflake/connector/aio/auth/_oauth_code.py +117 -0
- snowflake/connector/aio/auth/_oauth_credentials.py +100 -0
- snowflake/connector/aio/auth/_okta.py +247 -0
- snowflake/connector/aio/auth/_pat.py +27 -0
- snowflake/connector/aio/auth/_usrpwdmfa.py +30 -0
- snowflake/connector/aio/auth/_webbrowser.py +405 -0
- snowflake/connector/aio/auth/_workload_identity.py +60 -0
- snowflake/connector/arrow_context.py +211 -0
- snowflake/connector/auth/__init__.py +53 -0
- snowflake/connector/auth/_auth.py +669 -0
- snowflake/connector/auth/_http_server.py +257 -0
- snowflake/connector/auth/_oauth_base.py +472 -0
- snowflake/connector/auth/by_plugin.py +219 -0
- snowflake/connector/auth/default.py +36 -0
- snowflake/connector/auth/idtoken.py +72 -0
- snowflake/connector/auth/keypair.py +225 -0
- snowflake/connector/auth/no_auth.py +39 -0
- snowflake/connector/auth/oauth.py +49 -0
- snowflake/connector/auth/oauth_code.py +490 -0
- snowflake/connector/auth/oauth_credentials.py +71 -0
- snowflake/connector/auth/okta.py +339 -0
- snowflake/connector/auth/pat.py +40 -0
- snowflake/connector/auth/usrpwdmfa.py +69 -0
- snowflake/connector/auth/webbrowser.py +514 -0
- snowflake/connector/auth/workload_identity.py +110 -0
- snowflake/connector/azure_storage_client.py +278 -0
- snowflake/connector/backoff_policies.py +141 -0
- snowflake/connector/bind_upload_agent.py +89 -0
- snowflake/connector/cache.py +696 -0
- snowflake/connector/compat.py +127 -0
- snowflake/connector/config_manager.py +553 -0
- snowflake/connector/connection.py +2671 -0
- snowflake/connector/connection_diagnostic.py +793 -0
- snowflake/connector/constants.py +460 -0
- snowflake/connector/converter.py +809 -0
- snowflake/connector/converter_issue23517.py +87 -0
- snowflake/connector/converter_null.py +14 -0
- snowflake/connector/converter_snowsql.py +205 -0
- snowflake/connector/crl.py +919 -0
- snowflake/connector/crl_cache.py +734 -0
- snowflake/connector/cursor.py +2099 -0
- snowflake/connector/dbapi.py +53 -0
- snowflake/connector/description.py +21 -0
- snowflake/connector/direct_file_operation_utils.py +88 -0
- snowflake/connector/encryption_util.py +220 -0
- snowflake/connector/errorcode.py +94 -0
- snowflake/connector/errors.py +666 -0
- snowflake/connector/externals_utils/__init__.py +0 -0
- snowflake/connector/externals_utils/externals_setup.py +30 -0
- snowflake/connector/feature.py +4 -0
- snowflake/connector/file_compression_type.py +118 -0
- snowflake/connector/file_lock.py +72 -0
- snowflake/connector/file_transfer_agent.py +1244 -0
- snowflake/connector/file_util.py +153 -0
- snowflake/connector/gcs_storage_client.py +476 -0
- snowflake/connector/gzip_decoder.py +85 -0
- snowflake/connector/interval_util.py +26 -0
- snowflake/connector/local_storage_client.py +90 -0
- snowflake/connector/log_configuration.py +62 -0
- snowflake/connector/logging_utils/__init__.py +0 -0
- snowflake/connector/logging_utils/filters.py +72 -0
- snowflake/connector/minicore/__init__.py +0 -0
- snowflake/connector/minicore/aix_ppc64/libsf_mini_core.a +0 -0
- snowflake/connector/minicore/aix_ppc64/libsf_mini_core_static.a +0 -0
- snowflake/connector/minicore/linux_aarch64_glibc/libsf_mini_core.so +0 -0
- snowflake/connector/minicore/linux_aarch64_musl/libsf_mini_core.so +0 -0
- snowflake/connector/minicore/linux_x86_64_glibc/libsf_mini_core.so +0 -0
- snowflake/connector/minicore/linux_x86_64_musl/libsf_mini_core.so +0 -0
- snowflake/connector/minicore/macos_aarch64/libsf_mini_core.dylib +0 -0
- snowflake/connector/minicore/macos_x86_64/libsf_mini_core.dylib +0 -0
- snowflake/connector/minicore/windows_x86_64/sf_mini_core.dll +0 -0
- snowflake/connector/minicore/windows_x86_64/sf_mini_core_static.lib +0 -0
- snowflake/connector/nanoarrow_arrow_iterator.cp314-win_amd64.pyd +0 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/ArrayConverter.cpp +60 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/ArrayConverter.hpp +29 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/BinaryConverter.cpp +19 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/BinaryConverter.hpp +26 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/BooleanConverter.cpp +21 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/BooleanConverter.hpp +23 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.cpp +570 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.hpp +98 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowIterator.cpp +125 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowIterator.hpp +115 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.cpp +1140 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.hpp +138 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DateConverter.cpp +47 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DateConverter.hpp +46 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DecFloatConverter.cpp +83 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DecFloatConverter.hpp +35 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DecimalConverter.cpp +97 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/DecimalConverter.hpp +72 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/FixedSizeListConverter.cpp +73 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/FixedSizeListConverter.hpp +28 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/FloatConverter.cpp +30 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/FloatConverter.hpp +35 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/IColumnConverter.hpp +17 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/IntConverter.cpp +23 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/IntConverter.hpp +45 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/IntervalConverter.cpp +75 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/IntervalConverter.hpp +57 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/LICENSE.txt +209 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/MapConverter.cpp +75 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/MapConverter.hpp +30 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/ObjectConverter.cpp +46 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/ObjectConverter.hpp +29 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Python/Common.cpp +8 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Python/Common.hpp +95 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Python/Helpers.cpp +57 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Python/Helpers.hpp +36 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.cpp +34 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.hpp +48 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/StringConverter.cpp +19 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/StringConverter.hpp +26 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/TimeConverter.cpp +36 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/TimeConverter.hpp +31 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/TimeStampConverter.cpp +346 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/TimeStampConverter.hpp +145 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Util/macros.hpp +14 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Util/time.cpp +67 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/Util/time.hpp +68 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_accessors.h +101 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_alloc.h +127 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_assert.h +45 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_builder.h +1908 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_emitter.h +215 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_endian.h +125 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_epilogue.h +7 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_flatbuffers.h +55 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_identifier.h +148 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_iov.h +31 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_prologue.h +8 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_refmap.h +144 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_rtconfig.h +162 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_types.h +97 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/flatcc_verifier.h +239 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/flatcc_portable.h +14 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/paligned_alloc.h +210 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pattributes.h +84 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pdiagnostic.h +84 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pdiagnostic_pop.h +20 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pdiagnostic_push.h +51 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pendian.h +206 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pendian_detect.h +118 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pinline.h +19 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pinttypes.h +52 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/portable.h +2 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/portable_basic.h +25 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pstatic_assert.h +67 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pstdalign.h +162 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pstdint.h +898 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/punaligned.h +190 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pversion.h +6 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc/portable/pwarnings.h +52 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/flatcc.c +3204 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow.c +3217 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow.h +3618 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow.hpp +379 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_arrow_iterator.pyx +262 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_device.c +512 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_device.h +350 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_ipc.c +33273 -0
- snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_ipc.h +397 -0
- snowflake/connector/nanoarrow_cpp/Logging/logging.cpp +100 -0
- snowflake/connector/nanoarrow_cpp/Logging/logging.hpp +49 -0
- snowflake/connector/network.py +1227 -0
- snowflake/connector/ocsp_asn1crypto.py +447 -0
- snowflake/connector/ocsp_snowflake.py +1939 -0
- snowflake/connector/options.py +179 -0
- snowflake/connector/os_details.py +88 -0
- snowflake/connector/pandas_tools.py +734 -0
- snowflake/connector/platform_detection.py +555 -0
- snowflake/connector/proxy.py +28 -0
- snowflake/connector/py.typed +0 -0
- snowflake/connector/result_batch.py +917 -0
- snowflake/connector/result_set.py +340 -0
- snowflake/connector/s3_storage_client.py +614 -0
- snowflake/connector/secret_detector.py +181 -0
- snowflake/connector/session_manager.py +650 -0
- snowflake/connector/sf_dirs.py +64 -0
- snowflake/connector/sfbinaryformat.py +35 -0
- snowflake/connector/sfdatetime.py +345 -0
- snowflake/connector/snow_logging.py +144 -0
- snowflake/connector/sqlstate.py +9 -0
- snowflake/connector/ssd_internal_keys.py +30 -0
- snowflake/connector/ssl_wrap_socket.py +281 -0
- snowflake/connector/storage_client.py +486 -0
- snowflake/connector/telemetry.py +254 -0
- snowflake/connector/telemetry_oob.py +544 -0
- snowflake/connector/test_util.py +30 -0
- snowflake/connector/time_util.py +159 -0
- snowflake/connector/token_cache.py +457 -0
- snowflake/connector/tool/__init__.py +0 -0
- snowflake/connector/tool/dump_certs.py +57 -0
- snowflake/connector/tool/dump_ocsp_response.py +139 -0
- snowflake/connector/tool/dump_ocsp_response_cache.py +194 -0
- snowflake/connector/tool/probe_connection.py +69 -0
- snowflake/connector/url_util.py +55 -0
- snowflake/connector/util_text.py +320 -0
- snowflake/connector/vendored/__init__.py +3 -0
- snowflake/connector/vendored/requests/LICENSE +175 -0
- snowflake/connector/vendored/requests/__init__.py +184 -0
- snowflake/connector/vendored/requests/__version__.py +14 -0
- snowflake/connector/vendored/requests/_internal_utils.py +50 -0
- snowflake/connector/vendored/requests/adapters.py +696 -0
- snowflake/connector/vendored/requests/api.py +157 -0
- snowflake/connector/vendored/requests/auth.py +314 -0
- snowflake/connector/vendored/requests/certs.py +17 -0
- snowflake/connector/vendored/requests/compat.py +106 -0
- snowflake/connector/vendored/requests/cookies.py +561 -0
- snowflake/connector/vendored/requests/exceptions.py +151 -0
- snowflake/connector/vendored/requests/help.py +134 -0
- snowflake/connector/vendored/requests/hooks.py +33 -0
- snowflake/connector/vendored/requests/models.py +1039 -0
- snowflake/connector/vendored/requests/packages.py +23 -0
- snowflake/connector/vendored/requests/sessions.py +831 -0
- snowflake/connector/vendored/requests/status_codes.py +128 -0
- snowflake/connector/vendored/requests/structures.py +99 -0
- snowflake/connector/vendored/requests/utils.py +1086 -0
- snowflake/connector/vendored/urllib3/LICENSE.txt +21 -0
- snowflake/connector/vendored/urllib3/__init__.py +211 -0
- snowflake/connector/vendored/urllib3/_base_connection.py +165 -0
- snowflake/connector/vendored/urllib3/_collections.py +487 -0
- snowflake/connector/vendored/urllib3/_request_methods.py +278 -0
- snowflake/connector/vendored/urllib3/_version.py +34 -0
- snowflake/connector/vendored/urllib3/connection.py +1108 -0
- snowflake/connector/vendored/urllib3/connectionpool.py +1178 -0
- snowflake/connector/vendored/urllib3/contrib/__init__.py +0 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/__init__.py +17 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/connection.py +260 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/emscripten_fetch_worker.js +110 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/fetch.py +726 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/request.py +22 -0
- snowflake/connector/vendored/urllib3/contrib/emscripten/response.py +277 -0
- snowflake/connector/vendored/urllib3/contrib/pyopenssl.py +564 -0
- snowflake/connector/vendored/urllib3/contrib/socks.py +228 -0
- snowflake/connector/vendored/urllib3/exceptions.py +335 -0
- snowflake/connector/vendored/urllib3/fields.py +341 -0
- snowflake/connector/vendored/urllib3/filepost.py +89 -0
- snowflake/connector/vendored/urllib3/http2/__init__.py +53 -0
- snowflake/connector/vendored/urllib3/http2/connection.py +356 -0
- snowflake/connector/vendored/urllib3/http2/probe.py +87 -0
- snowflake/connector/vendored/urllib3/poolmanager.py +651 -0
- snowflake/connector/vendored/urllib3/py.typed +2 -0
- snowflake/connector/vendored/urllib3/response.py +1480 -0
- snowflake/connector/vendored/urllib3/util/__init__.py +42 -0
- snowflake/connector/vendored/urllib3/util/connection.py +137 -0
- snowflake/connector/vendored/urllib3/util/proxy.py +43 -0
- snowflake/connector/vendored/urllib3/util/request.py +263 -0
- snowflake/connector/vendored/urllib3/util/response.py +101 -0
- snowflake/connector/vendored/urllib3/util/retry.py +549 -0
- snowflake/connector/vendored/urllib3/util/ssl_.py +527 -0
- snowflake/connector/vendored/urllib3/util/ssl_match_hostname.py +159 -0
- snowflake/connector/vendored/urllib3/util/ssltransport.py +271 -0
- snowflake/connector/vendored/urllib3/util/timeout.py +275 -0
- snowflake/connector/vendored/urllib3/util/url.py +469 -0
- snowflake/connector/vendored/urllib3/util/util.py +42 -0
- snowflake/connector/vendored/urllib3/util/wait.py +124 -0
- snowflake/connector/version.py +3 -0
- snowflake/connector/wif_util.py +434 -0
- snowflake_connector_python-4.4.0.dist-info/DELVEWHEEL +2 -0
- snowflake_connector_python-4.4.0.dist-info/METADATA +1591 -0
- snowflake_connector_python-4.4.0.dist-info/RECORD +302 -0
- snowflake_connector_python-4.4.0.dist-info/WHEEL +5 -0
- snowflake_connector_python-4.4.0.dist-info/entry_points.txt +4 -0
- snowflake_connector_python-4.4.0.dist-info/licenses/LICENSE.txt +202 -0
- snowflake_connector_python-4.4.0.dist-info/licenses/NOTICE +8 -0
- snowflake_connector_python-4.4.0.dist-info/top_level.txt +1 -0
- snowflake_connector_python.libs/msvcp140-0f885b509a685d2bbfa652fed26b5fb3.dll +0 -0
|
@@ -0,0 +1,917 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import abc
|
|
4
|
+
import gzip
|
|
5
|
+
import json
|
|
6
|
+
import time
|
|
7
|
+
from base64 import b64decode
|
|
8
|
+
from enum import Enum, unique
|
|
9
|
+
from logging import getLogger
|
|
10
|
+
from typing import TYPE_CHECKING, Any, Callable, Iterator, NamedTuple, Sequence
|
|
11
|
+
|
|
12
|
+
from typing_extensions import Self
|
|
13
|
+
|
|
14
|
+
from .arrow_context import ArrowConverterContext
|
|
15
|
+
from .backoff_policies import exponential_backoff
|
|
16
|
+
from .compat import OK, UNAUTHORIZED, urlparse
|
|
17
|
+
from .constants import FIELD_TYPES, IterUnit
|
|
18
|
+
from .errorcode import ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE, ER_NO_PYARROW
|
|
19
|
+
from .errors import Error, InterfaceError, NotSupportedError, ProgrammingError
|
|
20
|
+
from .network import (
|
|
21
|
+
RetryRequest,
|
|
22
|
+
get_http_retryable_error,
|
|
23
|
+
is_retryable_http_code,
|
|
24
|
+
raise_failed_request_error,
|
|
25
|
+
raise_okta_unauthorized_error,
|
|
26
|
+
)
|
|
27
|
+
from .options import installed_pandas
|
|
28
|
+
from .options import pyarrow as pa
|
|
29
|
+
from .secret_detector import SecretDetector
|
|
30
|
+
from .session_manager import HttpConfig, SessionManager, SessionManagerFactory
|
|
31
|
+
from .time_util import TimerContextManager
|
|
32
|
+
|
|
33
|
+
logger = getLogger(__name__)
|
|
34
|
+
|
|
35
|
+
MAX_DOWNLOAD_RETRY = 10
|
|
36
|
+
DOWNLOAD_TIMEOUT = 7 # seconds
|
|
37
|
+
|
|
38
|
+
if TYPE_CHECKING: # pragma: no cover
|
|
39
|
+
from pandas import DataFrame
|
|
40
|
+
from pyarrow import DataType, Table
|
|
41
|
+
|
|
42
|
+
from .connection import SnowflakeConnection
|
|
43
|
+
from .converter import SnowflakeConverterType
|
|
44
|
+
from .cursor import ResultMetadataV2, SnowflakeCursor
|
|
45
|
+
from .vendored.requests import Response
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# emtpy pyarrow type array corresponding to FIELD_TYPES
|
|
49
|
+
FIELD_TYPE_TO_PA_TYPE: list[Callable[[ResultMetadataV2], DataType]] = []
|
|
50
|
+
|
|
51
|
+
# qrmk related constants
|
|
52
|
+
SSE_C_ALGORITHM = "x-amz-server-side-encryption-customer-algorithm"
|
|
53
|
+
SSE_C_KEY = "x-amz-server-side-encryption-customer-key"
|
|
54
|
+
SSE_C_AES = "AES256"
|
|
55
|
+
|
|
56
|
+
_GZIP_MAGIC = b"\x1f\x8b"
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def _ensure_decompressed(response: Response) -> None:
|
|
60
|
+
"""Decompress the response body if HTTP-level gzip decompression was skipped.
|
|
61
|
+
|
|
62
|
+
Cloud storage (S3/GCS/Azure) may serve result-set chunks as raw gzip blobs
|
|
63
|
+
without a ``Content-Encoding: gzip`` header, in which case urllib3's
|
|
64
|
+
transparent decompression never activates. Detect this by inspecting the
|
|
65
|
+
cached *response.content* for the gzip magic number and decompress in-place.
|
|
66
|
+
"""
|
|
67
|
+
content = response.content
|
|
68
|
+
if isinstance(content, bytes) and content[:2] == _GZIP_MAGIC:
|
|
69
|
+
logger.debug(
|
|
70
|
+
"Response body starts with gzip magic number but was not "
|
|
71
|
+
"decompressed by the HTTP stack; decompressing explicitly."
|
|
72
|
+
)
|
|
73
|
+
response._content = gzip.decompress(content)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _create_nanoarrow_iterator(
|
|
77
|
+
data: bytes,
|
|
78
|
+
context: ArrowConverterContext,
|
|
79
|
+
use_dict_result: bool,
|
|
80
|
+
numpy: bool,
|
|
81
|
+
number_to_decimal: bool,
|
|
82
|
+
row_unit: IterUnit,
|
|
83
|
+
check_error_on_every_column: bool = True,
|
|
84
|
+
force_microsecond_precision: bool = False,
|
|
85
|
+
):
|
|
86
|
+
from .nanoarrow_arrow_iterator import PyArrowRowIterator, PyArrowTableIterator
|
|
87
|
+
|
|
88
|
+
logger.debug("Using nanoarrow as the arrow data converter")
|
|
89
|
+
return (
|
|
90
|
+
PyArrowRowIterator(
|
|
91
|
+
None,
|
|
92
|
+
data,
|
|
93
|
+
context,
|
|
94
|
+
use_dict_result,
|
|
95
|
+
numpy,
|
|
96
|
+
number_to_decimal,
|
|
97
|
+
check_error_on_every_column,
|
|
98
|
+
)
|
|
99
|
+
if row_unit == IterUnit.ROW_UNIT
|
|
100
|
+
else PyArrowTableIterator(
|
|
101
|
+
None,
|
|
102
|
+
data,
|
|
103
|
+
context,
|
|
104
|
+
use_dict_result,
|
|
105
|
+
numpy,
|
|
106
|
+
number_to_decimal,
|
|
107
|
+
check_error_on_every_column,
|
|
108
|
+
force_microsecond_precision,
|
|
109
|
+
)
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@unique
|
|
114
|
+
class DownloadMetrics(Enum):
|
|
115
|
+
"""Defines the keywords by which to store metrics for chunks."""
|
|
116
|
+
|
|
117
|
+
download = "download" # Download time in milliseconds
|
|
118
|
+
parse = "parse" # Parsing time to final data types
|
|
119
|
+
load = "load" # Parsing time from initial type to intermediate types
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
class RemoteChunkInfo(NamedTuple):
|
|
123
|
+
"""Small class that holds information about chunks that are given by back-end."""
|
|
124
|
+
|
|
125
|
+
url: str
|
|
126
|
+
uncompressedSize: int
|
|
127
|
+
compressedSize: int
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def create_batches_from_response(
|
|
131
|
+
cursor: SnowflakeCursor,
|
|
132
|
+
_format: str,
|
|
133
|
+
data: dict[str, Any],
|
|
134
|
+
schema: Sequence[ResultMetadataV2],
|
|
135
|
+
) -> list[ResultBatch]:
|
|
136
|
+
column_converters: list[tuple[str, SnowflakeConverterType]] = []
|
|
137
|
+
arrow_context: ArrowConverterContext | None = None
|
|
138
|
+
rowtypes = data["rowtype"]
|
|
139
|
+
total_len: int = data.get("total", 0)
|
|
140
|
+
first_chunk_len = total_len
|
|
141
|
+
rest_of_chunks: list[ResultBatch] = []
|
|
142
|
+
if _format == "json":
|
|
143
|
+
|
|
144
|
+
def col_to_converter(col: dict[str, Any]) -> tuple[str, SnowflakeConverterType]:
|
|
145
|
+
type_name = col["type"].upper()
|
|
146
|
+
python_method = cursor._connection.converter.to_python_method(
|
|
147
|
+
type_name, col
|
|
148
|
+
)
|
|
149
|
+
return type_name, python_method
|
|
150
|
+
|
|
151
|
+
column_converters = [col_to_converter(c) for c in rowtypes]
|
|
152
|
+
else:
|
|
153
|
+
rowset_b64 = data.get("rowsetBase64")
|
|
154
|
+
arrow_context = ArrowConverterContext(cursor._connection._session_parameters)
|
|
155
|
+
if "chunks" in data:
|
|
156
|
+
chunks = data["chunks"]
|
|
157
|
+
logger.debug(f"chunk size={len(chunks)}")
|
|
158
|
+
# prepare the downloader for further fetch
|
|
159
|
+
qrmk = data.get("qrmk")
|
|
160
|
+
chunk_headers: dict[str, Any] = {}
|
|
161
|
+
if "chunkHeaders" in data:
|
|
162
|
+
chunk_headers = {}
|
|
163
|
+
for header_key, header_value in data["chunkHeaders"].items():
|
|
164
|
+
chunk_headers[header_key] = header_value
|
|
165
|
+
if "encryption" not in header_key:
|
|
166
|
+
logger.debug(
|
|
167
|
+
f"added chunk header: key={header_key}, value={header_value}"
|
|
168
|
+
)
|
|
169
|
+
elif qrmk is not None:
|
|
170
|
+
logger.debug(f"qrmk={SecretDetector.mask_secrets(qrmk)}")
|
|
171
|
+
chunk_headers[SSE_C_ALGORITHM] = SSE_C_AES
|
|
172
|
+
chunk_headers[SSE_C_KEY] = qrmk
|
|
173
|
+
|
|
174
|
+
def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
|
|
175
|
+
return RemoteChunkInfo(
|
|
176
|
+
url=c["url"],
|
|
177
|
+
uncompressedSize=c["uncompressedSize"],
|
|
178
|
+
compressedSize=c["compressedSize"],
|
|
179
|
+
)
|
|
180
|
+
|
|
181
|
+
if _format == "json":
|
|
182
|
+
rest_of_chunks = [
|
|
183
|
+
JSONResultBatch(
|
|
184
|
+
c["rowCount"],
|
|
185
|
+
chunk_headers,
|
|
186
|
+
remote_chunk_info(c),
|
|
187
|
+
schema,
|
|
188
|
+
column_converters,
|
|
189
|
+
cursor._use_dict_result,
|
|
190
|
+
json_result_force_utf8_decoding=cursor._connection._json_result_force_utf8_decoding,
|
|
191
|
+
session_manager=cursor._connection._session_manager.clone(),
|
|
192
|
+
)
|
|
193
|
+
for c in chunks
|
|
194
|
+
]
|
|
195
|
+
else:
|
|
196
|
+
rest_of_chunks = [
|
|
197
|
+
ArrowResultBatch(
|
|
198
|
+
c["rowCount"],
|
|
199
|
+
chunk_headers,
|
|
200
|
+
remote_chunk_info(c),
|
|
201
|
+
arrow_context,
|
|
202
|
+
cursor._use_dict_result,
|
|
203
|
+
cursor._connection._numpy,
|
|
204
|
+
schema,
|
|
205
|
+
cursor._connection._arrow_number_to_decimal,
|
|
206
|
+
session_manager=cursor._connection._session_manager.clone(),
|
|
207
|
+
)
|
|
208
|
+
for c in chunks
|
|
209
|
+
]
|
|
210
|
+
for c in rest_of_chunks:
|
|
211
|
+
first_chunk_len -= c.rowcount
|
|
212
|
+
if _format == "json":
|
|
213
|
+
first_chunk = JSONResultBatch.from_data(
|
|
214
|
+
data.get("rowset"),
|
|
215
|
+
first_chunk_len,
|
|
216
|
+
schema,
|
|
217
|
+
column_converters,
|
|
218
|
+
cursor._use_dict_result,
|
|
219
|
+
session_manager=cursor._connection._session_manager.clone(),
|
|
220
|
+
)
|
|
221
|
+
elif rowset_b64 is not None:
|
|
222
|
+
first_chunk = ArrowResultBatch.from_data(
|
|
223
|
+
rowset_b64,
|
|
224
|
+
first_chunk_len,
|
|
225
|
+
arrow_context,
|
|
226
|
+
cursor._use_dict_result,
|
|
227
|
+
cursor._connection._numpy,
|
|
228
|
+
schema,
|
|
229
|
+
cursor._connection._arrow_number_to_decimal,
|
|
230
|
+
session_manager=cursor._connection._session_manager.clone(),
|
|
231
|
+
)
|
|
232
|
+
else:
|
|
233
|
+
logger.error(f"Don't know how to construct ResultBatches from response: {data}")
|
|
234
|
+
first_chunk = ArrowResultBatch.from_data(
|
|
235
|
+
"",
|
|
236
|
+
0,
|
|
237
|
+
arrow_context,
|
|
238
|
+
cursor._use_dict_result,
|
|
239
|
+
cursor._connection._numpy,
|
|
240
|
+
schema,
|
|
241
|
+
cursor._connection._arrow_number_to_decimal,
|
|
242
|
+
session_manager=cursor._connection._session_manager.clone(),
|
|
243
|
+
)
|
|
244
|
+
|
|
245
|
+
return [first_chunk] + rest_of_chunks
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
class ResultBatch(abc.ABC):
|
|
249
|
+
"""Represents what the back-end calls a result chunk.
|
|
250
|
+
|
|
251
|
+
These are parts of a result set of a query. They each know how to retrieve their
|
|
252
|
+
own results and convert them into Python native formats.
|
|
253
|
+
|
|
254
|
+
As you are iterating through a ResultBatch you should check whether the yielded
|
|
255
|
+
value is an ``Exception`` in case there was some error parsing the current row
|
|
256
|
+
we might yield one of these to allow iteration to continue instead of raising the
|
|
257
|
+
``Exception`` when it occurs.
|
|
258
|
+
|
|
259
|
+
These objects are pickleable for easy distribution and replication.
|
|
260
|
+
|
|
261
|
+
Please note that the URLs stored in these do expire. The lifetime is dictated by the
|
|
262
|
+
Snowflake back-end, at the time of writing this this is 6 hours.
|
|
263
|
+
|
|
264
|
+
They can be iterated over multiple times and in different ways. Please follow the
|
|
265
|
+
code in ``cursor.py`` to make sure that you are using this class correctly.
|
|
266
|
+
|
|
267
|
+
"""
|
|
268
|
+
|
|
269
|
+
def __init__(
|
|
270
|
+
self,
|
|
271
|
+
rowcount: int,
|
|
272
|
+
chunk_headers: dict[str, str] | None,
|
|
273
|
+
remote_chunk_info: RemoteChunkInfo | None,
|
|
274
|
+
schema: Sequence[ResultMetadataV2],
|
|
275
|
+
use_dict_result: bool,
|
|
276
|
+
session_manager: SessionManager | None = None,
|
|
277
|
+
) -> None:
|
|
278
|
+
self.rowcount = rowcount
|
|
279
|
+
self._chunk_headers = chunk_headers
|
|
280
|
+
self._remote_chunk_info = remote_chunk_info
|
|
281
|
+
self._schema = schema
|
|
282
|
+
self.schema = (
|
|
283
|
+
[s._to_result_metadata_v1() for s in schema] if schema is not None else None
|
|
284
|
+
)
|
|
285
|
+
self._use_dict_result = use_dict_result
|
|
286
|
+
# Passed to contain the configured Http behavior in case the connection is no longer active for the download
|
|
287
|
+
# Can be overridden with setters if needed.
|
|
288
|
+
self._session_manager = session_manager
|
|
289
|
+
self._metrics: dict[str, int] = {}
|
|
290
|
+
self._data: str | list[tuple[Any, ...]] | None = None
|
|
291
|
+
if self._remote_chunk_info:
|
|
292
|
+
parsed_url = urlparse(self._remote_chunk_info.url)
|
|
293
|
+
path_parts = parsed_url.path.rsplit("/", 1)
|
|
294
|
+
self.id = path_parts[-1]
|
|
295
|
+
else:
|
|
296
|
+
self.id = str(self.rowcount)
|
|
297
|
+
|
|
298
|
+
@property
|
|
299
|
+
def _local(self) -> bool:
|
|
300
|
+
"""Whether this chunk is local."""
|
|
301
|
+
return self._data is not None
|
|
302
|
+
|
|
303
|
+
@property
|
|
304
|
+
def compressed_size(self) -> int | None:
|
|
305
|
+
"""Returns the size of chunk in bytes in compressed form.
|
|
306
|
+
|
|
307
|
+
If it's a local chunk this function returns None.
|
|
308
|
+
"""
|
|
309
|
+
if self._local:
|
|
310
|
+
return None
|
|
311
|
+
return self._remote_chunk_info.compressedSize
|
|
312
|
+
|
|
313
|
+
@property
|
|
314
|
+
def uncompressed_size(self) -> int | None:
|
|
315
|
+
"""Returns the size of chunk in bytes in uncompressed form.
|
|
316
|
+
|
|
317
|
+
If it's a local chunk this function returns None.
|
|
318
|
+
"""
|
|
319
|
+
if self._local:
|
|
320
|
+
return None
|
|
321
|
+
return self._remote_chunk_info.uncompressedSize
|
|
322
|
+
|
|
323
|
+
@property
|
|
324
|
+
def column_names(self) -> list[str]:
|
|
325
|
+
return [col.name for col in self._schema]
|
|
326
|
+
|
|
327
|
+
@property
|
|
328
|
+
def session_manager(self) -> SessionManager | None:
|
|
329
|
+
return self._session_manager
|
|
330
|
+
|
|
331
|
+
@session_manager.setter
|
|
332
|
+
def session_manager(self, session_manager: SessionManager | None) -> None:
|
|
333
|
+
self._session_manager = session_manager
|
|
334
|
+
|
|
335
|
+
@property
|
|
336
|
+
def http_config(self):
|
|
337
|
+
return self._session_manager.config
|
|
338
|
+
|
|
339
|
+
@http_config.setter
|
|
340
|
+
def http_config(self, config: HttpConfig) -> None:
|
|
341
|
+
if self._session_manager:
|
|
342
|
+
self._session_manager.config = config
|
|
343
|
+
else:
|
|
344
|
+
self._session_manager = SessionManagerFactory.get_manager(config=config)
|
|
345
|
+
|
|
346
|
+
def __iter__(
|
|
347
|
+
self,
|
|
348
|
+
) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
|
|
349
|
+
"""Returns an iterator through the data this chunk holds.
|
|
350
|
+
|
|
351
|
+
In case of this chunk being a local one it iterates through the local already
|
|
352
|
+
parsed data and if it's a remote chunk it will download, parse its data and
|
|
353
|
+
return an iterator through it.
|
|
354
|
+
"""
|
|
355
|
+
return self.create_iter()
|
|
356
|
+
|
|
357
|
+
def _download(
|
|
358
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
359
|
+
) -> Response:
|
|
360
|
+
"""Downloads the data that the ``ResultBatch`` is pointing at."""
|
|
361
|
+
sleep_timer = 1
|
|
362
|
+
backoff = (
|
|
363
|
+
connection._backoff_generator
|
|
364
|
+
if connection is not None
|
|
365
|
+
else exponential_backoff()()
|
|
366
|
+
)
|
|
367
|
+
for retry in range(MAX_DOWNLOAD_RETRY):
|
|
368
|
+
try:
|
|
369
|
+
with TimerContextManager() as download_metric:
|
|
370
|
+
logger.debug(f"started downloading result batch id: {self.id}")
|
|
371
|
+
chunk_url = self._remote_chunk_info.url
|
|
372
|
+
request_data = {
|
|
373
|
+
"url": chunk_url,
|
|
374
|
+
"headers": self._chunk_headers,
|
|
375
|
+
"timeout": DOWNLOAD_TIMEOUT,
|
|
376
|
+
}
|
|
377
|
+
# Try to reuse a connection if possible
|
|
378
|
+
|
|
379
|
+
if (
|
|
380
|
+
connection
|
|
381
|
+
and connection.rest
|
|
382
|
+
and connection.rest.session_manager is not None
|
|
383
|
+
):
|
|
384
|
+
# If connection was explicitly passed and not closed yet - we can reuse SessionManager with session pooling
|
|
385
|
+
with connection.rest.use_requests_session(
|
|
386
|
+
request_data["url"]
|
|
387
|
+
) as session:
|
|
388
|
+
logger.debug(
|
|
389
|
+
f"downloading result batch id: {self.id} with existing session {session}"
|
|
390
|
+
)
|
|
391
|
+
response = session.request("get", **request_data)
|
|
392
|
+
elif self._session_manager is not None:
|
|
393
|
+
# If connection is not accessible or was already closed, but cursors are now used to fetch the data - we will only reuse the http setup (through cloned SessionManager without session pooling)
|
|
394
|
+
with self._session_manager.use_session(
|
|
395
|
+
request_data["url"]
|
|
396
|
+
) as session:
|
|
397
|
+
response = session.request("get", **request_data)
|
|
398
|
+
else:
|
|
399
|
+
# If there was no session manager cloned, then we are using a default Session Manager setup, since it is very unlikely to enter this part outside of testing
|
|
400
|
+
logger.debug(
|
|
401
|
+
f"downloading result batch id: {self.id} with new session through local session manager"
|
|
402
|
+
)
|
|
403
|
+
local_session_manager = SessionManagerFactory.get_manager(
|
|
404
|
+
use_pooling=False
|
|
405
|
+
)
|
|
406
|
+
response = local_session_manager.get(**request_data)
|
|
407
|
+
|
|
408
|
+
if response.status_code == OK:
|
|
409
|
+
logger.debug(
|
|
410
|
+
f"successfully downloaded result batch id: {self.id}"
|
|
411
|
+
)
|
|
412
|
+
break
|
|
413
|
+
|
|
414
|
+
# Raise error here to correctly go in to exception clause
|
|
415
|
+
if is_retryable_http_code(response.status_code):
|
|
416
|
+
# retryable server exceptions
|
|
417
|
+
error: Error = get_http_retryable_error(response.status_code)
|
|
418
|
+
raise RetryRequest(error)
|
|
419
|
+
elif response.status_code == UNAUTHORIZED:
|
|
420
|
+
# make a unauthorized error
|
|
421
|
+
raise_okta_unauthorized_error(None, response)
|
|
422
|
+
else:
|
|
423
|
+
raise_failed_request_error(None, chunk_url, "get", response)
|
|
424
|
+
|
|
425
|
+
except (RetryRequest, Exception) as e:
|
|
426
|
+
if retry == MAX_DOWNLOAD_RETRY - 1:
|
|
427
|
+
# Re-throw if we failed on the last retry
|
|
428
|
+
e = e.args[0] if isinstance(e, RetryRequest) else e
|
|
429
|
+
raise e
|
|
430
|
+
sleep_timer = next(backoff)
|
|
431
|
+
logger.exception(
|
|
432
|
+
f"Failed to fetch the large result set batch "
|
|
433
|
+
f"{self.id} for the {retry + 1} th time, "
|
|
434
|
+
f"backing off for {sleep_timer}s for the reason: '{e}'"
|
|
435
|
+
)
|
|
436
|
+
time.sleep(sleep_timer)
|
|
437
|
+
|
|
438
|
+
self._metrics[DownloadMetrics.download.value] = (
|
|
439
|
+
download_metric.get_timing_millis()
|
|
440
|
+
)
|
|
441
|
+
|
|
442
|
+
_ensure_decompressed(response)
|
|
443
|
+
return response
|
|
444
|
+
|
|
445
|
+
@abc.abstractmethod
|
|
446
|
+
def create_iter(
|
|
447
|
+
self, **kwargs
|
|
448
|
+
) -> (
|
|
449
|
+
Iterator[dict | Exception]
|
|
450
|
+
| Iterator[tuple | Exception]
|
|
451
|
+
| Iterator[Table]
|
|
452
|
+
| Iterator[DataFrame]
|
|
453
|
+
):
|
|
454
|
+
"""Downloads the data from from blob storage that this ResultChunk points at.
|
|
455
|
+
|
|
456
|
+
This function is the one that does the actual work for ``self.__iter__``.
|
|
457
|
+
|
|
458
|
+
It is necessary because a ``ResultBatch`` can return multiple types of
|
|
459
|
+
iterators. A good example of this is simply iterating through
|
|
460
|
+
``SnowflakeCursor`` and calling ``fetch_pandas_batches`` on it.
|
|
461
|
+
"""
|
|
462
|
+
raise NotImplementedError()
|
|
463
|
+
|
|
464
|
+
def _check_can_use_pandas(self) -> None:
|
|
465
|
+
if not installed_pandas:
|
|
466
|
+
msg = (
|
|
467
|
+
"Optional dependency: 'pandas' is not installed, please see the following link for install "
|
|
468
|
+
"instructions: https://docs.snowflake.com/en/user-guide/python-connector-pandas.html#installation"
|
|
469
|
+
)
|
|
470
|
+
errno = ER_NO_PYARROW
|
|
471
|
+
|
|
472
|
+
raise Error.errorhandler_make_exception(
|
|
473
|
+
ProgrammingError,
|
|
474
|
+
{
|
|
475
|
+
"msg": msg,
|
|
476
|
+
"errno": errno,
|
|
477
|
+
},
|
|
478
|
+
)
|
|
479
|
+
|
|
480
|
+
@abc.abstractmethod
|
|
481
|
+
def to_pandas(self) -> DataFrame:
|
|
482
|
+
raise NotImplementedError()
|
|
483
|
+
|
|
484
|
+
@abc.abstractmethod
|
|
485
|
+
def to_arrow(self) -> Table:
|
|
486
|
+
raise NotImplementedError()
|
|
487
|
+
|
|
488
|
+
@abc.abstractmethod
|
|
489
|
+
def populate_data(
|
|
490
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
491
|
+
) -> Self:
|
|
492
|
+
"""Downloads the data that the ``ResultBatch`` is pointing at and populates it into self._data.
|
|
493
|
+
Returns the instance itself."""
|
|
494
|
+
raise NotImplementedError()
|
|
495
|
+
|
|
496
|
+
|
|
497
|
+
class JSONResultBatch(ResultBatch):
|
|
498
|
+
def __init__(
|
|
499
|
+
self,
|
|
500
|
+
rowcount: int,
|
|
501
|
+
chunk_headers: dict[str, str] | None,
|
|
502
|
+
remote_chunk_info: RemoteChunkInfo | None,
|
|
503
|
+
schema: Sequence[ResultMetadataV2],
|
|
504
|
+
column_converters: Sequence[tuple[str, SnowflakeConverterType]],
|
|
505
|
+
use_dict_result: bool,
|
|
506
|
+
*,
|
|
507
|
+
json_result_force_utf8_decoding: bool = False,
|
|
508
|
+
session_manager: SessionManager | None = None,
|
|
509
|
+
) -> None:
|
|
510
|
+
super().__init__(
|
|
511
|
+
rowcount,
|
|
512
|
+
chunk_headers,
|
|
513
|
+
remote_chunk_info,
|
|
514
|
+
schema,
|
|
515
|
+
use_dict_result,
|
|
516
|
+
session_manager,
|
|
517
|
+
)
|
|
518
|
+
self._json_result_force_utf8_decoding = json_result_force_utf8_decoding
|
|
519
|
+
self.column_converters = column_converters
|
|
520
|
+
|
|
521
|
+
@classmethod
|
|
522
|
+
def from_data(
|
|
523
|
+
cls,
|
|
524
|
+
data: Sequence[Sequence[Any]],
|
|
525
|
+
data_len: int,
|
|
526
|
+
schema: Sequence[ResultMetadataV2],
|
|
527
|
+
column_converters: Sequence[tuple[str, SnowflakeConverterType]],
|
|
528
|
+
use_dict_result: bool,
|
|
529
|
+
session_manager: SessionManager | None = None,
|
|
530
|
+
):
|
|
531
|
+
"""Initializes a ``JSONResultBatch`` from static, local data."""
|
|
532
|
+
new_chunk = cls(
|
|
533
|
+
len(data),
|
|
534
|
+
None,
|
|
535
|
+
None,
|
|
536
|
+
schema,
|
|
537
|
+
column_converters,
|
|
538
|
+
use_dict_result,
|
|
539
|
+
session_manager=session_manager,
|
|
540
|
+
)
|
|
541
|
+
new_chunk._data = new_chunk._parse(data)
|
|
542
|
+
return new_chunk
|
|
543
|
+
|
|
544
|
+
def _load(self, response: Response) -> list:
|
|
545
|
+
"""This function loads a compressed JSON file into memory.
|
|
546
|
+
|
|
547
|
+
Returns:
|
|
548
|
+
Whatever ``json.loads`` return, but in a list.
|
|
549
|
+
Unfortunately there's no type hint for this.
|
|
550
|
+
For context: https://github.com/python/typing/issues/182
|
|
551
|
+
"""
|
|
552
|
+
# if users specify how to decode the data, we decode the bytes using the specified encoding
|
|
553
|
+
if self._json_result_force_utf8_decoding:
|
|
554
|
+
try:
|
|
555
|
+
read_data = str(response.content, "utf-8", errors="strict")
|
|
556
|
+
except Exception as exc:
|
|
557
|
+
err_msg = f"failed to decode json result content due to error {exc!r}"
|
|
558
|
+
logger.error(err_msg)
|
|
559
|
+
raise Error(msg=err_msg)
|
|
560
|
+
else:
|
|
561
|
+
# note: SNOW-787480 response.apparent_encoding is unreliable, chardet.detect can be wrong which is used by
|
|
562
|
+
# response.text to decode content, check issue: https://github.com/chardet/chardet/issues/148
|
|
563
|
+
read_data = response.text
|
|
564
|
+
return json.loads("".join(["[", read_data, "]"]))
|
|
565
|
+
|
|
566
|
+
def _parse(
|
|
567
|
+
self, downloaded_data
|
|
568
|
+
) -> list[dict | Exception] | list[tuple | Exception]:
|
|
569
|
+
"""Parses downloaded data into its final form."""
|
|
570
|
+
logger.debug(f"parsing for result batch id: {self.id}")
|
|
571
|
+
result_list = []
|
|
572
|
+
if self._use_dict_result:
|
|
573
|
+
for row in downloaded_data:
|
|
574
|
+
row_result = {}
|
|
575
|
+
try:
|
|
576
|
+
for (_t, c), v, col in zip(
|
|
577
|
+
self.column_converters,
|
|
578
|
+
row,
|
|
579
|
+
self._schema,
|
|
580
|
+
):
|
|
581
|
+
row_result[col.name] = v if c is None or v is None else c(v)
|
|
582
|
+
result_list.append(row_result)
|
|
583
|
+
except Exception as error:
|
|
584
|
+
msg = f"Failed to convert: field {col.name}: {_t}::{v}, Error: {error}"
|
|
585
|
+
logger.exception(msg)
|
|
586
|
+
result_list.append(
|
|
587
|
+
Error.errorhandler_make_exception(
|
|
588
|
+
InterfaceError,
|
|
589
|
+
{
|
|
590
|
+
"msg": msg,
|
|
591
|
+
"errno": ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE,
|
|
592
|
+
},
|
|
593
|
+
)
|
|
594
|
+
)
|
|
595
|
+
else:
|
|
596
|
+
for row in downloaded_data:
|
|
597
|
+
row_result = [None] * len(self._schema)
|
|
598
|
+
try:
|
|
599
|
+
idx = 0
|
|
600
|
+
for (_t, c), v, _col in zip(
|
|
601
|
+
self.column_converters,
|
|
602
|
+
row,
|
|
603
|
+
self._schema,
|
|
604
|
+
):
|
|
605
|
+
row_result[idx] = v if c is None or v is None else c(v)
|
|
606
|
+
idx += 1
|
|
607
|
+
result_list.append(tuple(row_result))
|
|
608
|
+
except Exception as error:
|
|
609
|
+
msg = f"Failed to convert: field {_col.name}: {_t}::{v}, Error: {error}"
|
|
610
|
+
logger.exception(msg)
|
|
611
|
+
result_list.append(
|
|
612
|
+
Error.errorhandler_make_exception(
|
|
613
|
+
InterfaceError,
|
|
614
|
+
{
|
|
615
|
+
"msg": msg,
|
|
616
|
+
"errno": ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE,
|
|
617
|
+
},
|
|
618
|
+
)
|
|
619
|
+
)
|
|
620
|
+
return result_list
|
|
621
|
+
|
|
622
|
+
def __repr__(self) -> str:
|
|
623
|
+
return f"JSONResultChunk({self.id})"
|
|
624
|
+
|
|
625
|
+
def _fetch_data(
|
|
626
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
627
|
+
) -> list[dict | Exception] | list[tuple | Exception]:
|
|
628
|
+
response = self._download(connection=connection)
|
|
629
|
+
# Load data to a intermediate form
|
|
630
|
+
logger.debug(f"started loading result batch id: {self.id}")
|
|
631
|
+
with TimerContextManager() as load_metric:
|
|
632
|
+
downloaded_data = self._load(response)
|
|
633
|
+
logger.debug(f"finished loading result batch id: {self.id}")
|
|
634
|
+
self._metrics[DownloadMetrics.load.value] = load_metric.get_timing_millis()
|
|
635
|
+
# Process downloaded data
|
|
636
|
+
with TimerContextManager() as parse_metric:
|
|
637
|
+
parsed_data = self._parse(downloaded_data)
|
|
638
|
+
self._metrics[DownloadMetrics.parse.value] = parse_metric.get_timing_millis()
|
|
639
|
+
return parsed_data
|
|
640
|
+
|
|
641
|
+
def populate_data(
|
|
642
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
643
|
+
) -> Self:
|
|
644
|
+
self._data = self._fetch_data(connection=connection, **kwargs)
|
|
645
|
+
return self
|
|
646
|
+
|
|
647
|
+
def create_iter(
|
|
648
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
649
|
+
) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
|
|
650
|
+
if self._local:
|
|
651
|
+
return iter(self._data)
|
|
652
|
+
return iter(self._fetch_data(connection=connection, **kwargs))
|
|
653
|
+
|
|
654
|
+
def _arrow_fetching_error(self):
|
|
655
|
+
return NotSupportedError(
|
|
656
|
+
f"Trying to use arrow fetching on {type(self)} which "
|
|
657
|
+
f"is not ArrowResultChunk"
|
|
658
|
+
)
|
|
659
|
+
|
|
660
|
+
def to_pandas(self):
|
|
661
|
+
raise self._arrow_fetching_error()
|
|
662
|
+
|
|
663
|
+
def to_arrow(self):
|
|
664
|
+
raise self._arrow_fetching_error()
|
|
665
|
+
|
|
666
|
+
|
|
667
|
+
class ArrowResultBatch(ResultBatch):
|
|
668
|
+
def __init__(
|
|
669
|
+
self,
|
|
670
|
+
rowcount: int,
|
|
671
|
+
chunk_headers: dict[str, str] | None,
|
|
672
|
+
remote_chunk_info: RemoteChunkInfo | None,
|
|
673
|
+
context: ArrowConverterContext,
|
|
674
|
+
use_dict_result: bool,
|
|
675
|
+
numpy: bool,
|
|
676
|
+
schema: Sequence[ResultMetadataV2],
|
|
677
|
+
number_to_decimal: bool,
|
|
678
|
+
session_manager: SessionManager | None = None,
|
|
679
|
+
) -> None:
|
|
680
|
+
super().__init__(
|
|
681
|
+
rowcount,
|
|
682
|
+
chunk_headers,
|
|
683
|
+
remote_chunk_info,
|
|
684
|
+
schema,
|
|
685
|
+
use_dict_result,
|
|
686
|
+
session_manager,
|
|
687
|
+
)
|
|
688
|
+
self._context = context
|
|
689
|
+
self._numpy = numpy
|
|
690
|
+
self._number_to_decimal = number_to_decimal
|
|
691
|
+
|
|
692
|
+
def __repr__(self) -> str:
|
|
693
|
+
return f"ArrowResultChunk({self.id})"
|
|
694
|
+
|
|
695
|
+
def _load(
|
|
696
|
+
self,
|
|
697
|
+
response: Response,
|
|
698
|
+
row_unit: IterUnit,
|
|
699
|
+
force_microsecond_precision: bool = False,
|
|
700
|
+
) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
|
|
701
|
+
"""Creates a ``PyArrowIterator`` from a response.
|
|
702
|
+
|
|
703
|
+
This is used to iterate through results in different ways depending on which
|
|
704
|
+
mode that ``PyArrowIterator`` is in.
|
|
705
|
+
"""
|
|
706
|
+
return _create_nanoarrow_iterator(
|
|
707
|
+
response.content,
|
|
708
|
+
self._context,
|
|
709
|
+
self._use_dict_result,
|
|
710
|
+
self._numpy,
|
|
711
|
+
self._number_to_decimal,
|
|
712
|
+
row_unit,
|
|
713
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
714
|
+
)
|
|
715
|
+
|
|
716
|
+
def _from_data(
|
|
717
|
+
self,
|
|
718
|
+
data: str | bytes,
|
|
719
|
+
iter_unit: IterUnit,
|
|
720
|
+
check_error_on_every_column: bool = True,
|
|
721
|
+
force_microsecond_precision: bool = False,
|
|
722
|
+
) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
|
|
723
|
+
"""Creates a ``PyArrowIterator`` files from a str.
|
|
724
|
+
|
|
725
|
+
This is used to iterate through results in different ways depending on which
|
|
726
|
+
mode that ``PyArrowIterator`` is in.
|
|
727
|
+
"""
|
|
728
|
+
if len(data) == 0:
|
|
729
|
+
return iter([])
|
|
730
|
+
|
|
731
|
+
if isinstance(data, str):
|
|
732
|
+
data = b64decode(data)
|
|
733
|
+
|
|
734
|
+
return _create_nanoarrow_iterator(
|
|
735
|
+
data,
|
|
736
|
+
self._context,
|
|
737
|
+
self._use_dict_result,
|
|
738
|
+
self._numpy,
|
|
739
|
+
self._number_to_decimal,
|
|
740
|
+
iter_unit,
|
|
741
|
+
check_error_on_every_column,
|
|
742
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
743
|
+
)
|
|
744
|
+
|
|
745
|
+
@classmethod
|
|
746
|
+
def from_data(
|
|
747
|
+
cls,
|
|
748
|
+
data: str,
|
|
749
|
+
data_len: int,
|
|
750
|
+
context: ArrowConverterContext,
|
|
751
|
+
use_dict_result: bool,
|
|
752
|
+
numpy: bool,
|
|
753
|
+
schema: Sequence[ResultMetadataV2],
|
|
754
|
+
number_to_decimal: bool,
|
|
755
|
+
session_manager: SessionManager | None = None,
|
|
756
|
+
):
|
|
757
|
+
"""Initializes an ``ArrowResultBatch`` from static, local data."""
|
|
758
|
+
new_chunk = cls(
|
|
759
|
+
data_len,
|
|
760
|
+
None,
|
|
761
|
+
None,
|
|
762
|
+
context,
|
|
763
|
+
use_dict_result,
|
|
764
|
+
numpy,
|
|
765
|
+
schema,
|
|
766
|
+
number_to_decimal,
|
|
767
|
+
session_manager=session_manager,
|
|
768
|
+
)
|
|
769
|
+
new_chunk._data = data
|
|
770
|
+
|
|
771
|
+
return new_chunk
|
|
772
|
+
|
|
773
|
+
def _create_iter(
|
|
774
|
+
self,
|
|
775
|
+
iter_unit: IterUnit,
|
|
776
|
+
connection: SnowflakeConnection | None = None,
|
|
777
|
+
force_microsecond_precision: bool = False,
|
|
778
|
+
) -> Iterator[dict | Exception] | Iterator[tuple | Exception] | Iterator[Table]:
|
|
779
|
+
"""Create an iterator for the ResultBatch. Used by get_arrow_iter."""
|
|
780
|
+
if self._local:
|
|
781
|
+
try:
|
|
782
|
+
return self._from_data(
|
|
783
|
+
self._data,
|
|
784
|
+
iter_unit,
|
|
785
|
+
(
|
|
786
|
+
connection.check_arrow_conversion_error_on_every_column
|
|
787
|
+
if connection
|
|
788
|
+
else None
|
|
789
|
+
),
|
|
790
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
791
|
+
)
|
|
792
|
+
except Exception:
|
|
793
|
+
if connection and getattr(connection, "_debug_arrow_chunk", False):
|
|
794
|
+
logger.debug(f"arrow data can not be parsed: {self._data}")
|
|
795
|
+
raise
|
|
796
|
+
response = self._download(connection=connection)
|
|
797
|
+
logger.debug(f"started loading result batch id: {self.id}")
|
|
798
|
+
with TimerContextManager() as load_metric:
|
|
799
|
+
try:
|
|
800
|
+
loaded_data = self._load(
|
|
801
|
+
response,
|
|
802
|
+
iter_unit,
|
|
803
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
804
|
+
)
|
|
805
|
+
except Exception:
|
|
806
|
+
if connection and getattr(connection, "_debug_arrow_chunk", False):
|
|
807
|
+
logger.debug(f"arrow data can not be parsed: {response}")
|
|
808
|
+
raise
|
|
809
|
+
logger.debug(f"finished loading result batch id: {self.id}")
|
|
810
|
+
self._metrics[DownloadMetrics.load.value] = load_metric.get_timing_millis()
|
|
811
|
+
return loaded_data
|
|
812
|
+
|
|
813
|
+
def _get_arrow_iter(
|
|
814
|
+
self,
|
|
815
|
+
connection: SnowflakeConnection | None = None,
|
|
816
|
+
force_microsecond_precision: bool = False,
|
|
817
|
+
) -> Iterator[Table]:
|
|
818
|
+
"""Returns an iterator for this batch which yields a pyarrow Table"""
|
|
819
|
+
return self._create_iter(
|
|
820
|
+
iter_unit=IterUnit.TABLE_UNIT,
|
|
821
|
+
connection=connection,
|
|
822
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
823
|
+
)
|
|
824
|
+
|
|
825
|
+
def _create_empty_table(self) -> Table:
|
|
826
|
+
"""Returns empty Arrow table based on schema"""
|
|
827
|
+
if installed_pandas:
|
|
828
|
+
# initialize pyarrow type array corresponding to FIELD_TYPES
|
|
829
|
+
FIELD_TYPE_TO_PA_TYPE = [e.pa_type for e in FIELD_TYPES]
|
|
830
|
+
fields = [
|
|
831
|
+
pa.field(s.name, FIELD_TYPE_TO_PA_TYPE[s.type_code](s))
|
|
832
|
+
for s in self._schema
|
|
833
|
+
]
|
|
834
|
+
return pa.schema(fields).empty_table()
|
|
835
|
+
|
|
836
|
+
def to_arrow(
|
|
837
|
+
self,
|
|
838
|
+
connection: SnowflakeConnection | None = None,
|
|
839
|
+
force_microsecond_precision: bool = False,
|
|
840
|
+
) -> Table:
|
|
841
|
+
"""Returns this batch as a pyarrow Table"""
|
|
842
|
+
val = next(
|
|
843
|
+
self._get_arrow_iter(
|
|
844
|
+
connection=connection,
|
|
845
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
846
|
+
),
|
|
847
|
+
None,
|
|
848
|
+
)
|
|
849
|
+
if val is not None:
|
|
850
|
+
return val
|
|
851
|
+
return self._create_empty_table()
|
|
852
|
+
|
|
853
|
+
def to_pandas(
|
|
854
|
+
self,
|
|
855
|
+
connection: SnowflakeConnection | None = None,
|
|
856
|
+
force_microsecond_precision: bool = False,
|
|
857
|
+
**kwargs,
|
|
858
|
+
) -> DataFrame:
|
|
859
|
+
"""Returns this batch as a pandas DataFrame"""
|
|
860
|
+
self._check_can_use_pandas()
|
|
861
|
+
table = self.to_arrow(
|
|
862
|
+
connection=connection,
|
|
863
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
864
|
+
)
|
|
865
|
+
return table.to_pandas(**kwargs)
|
|
866
|
+
|
|
867
|
+
def _get_pandas_iter(
|
|
868
|
+
self,
|
|
869
|
+
connection: SnowflakeConnection | None = None,
|
|
870
|
+
force_microsecond_precision: bool = False,
|
|
871
|
+
**kwargs,
|
|
872
|
+
) -> Iterator[DataFrame]:
|
|
873
|
+
"""An iterator for this batch which yields a pandas DataFrame"""
|
|
874
|
+
iterator_data = []
|
|
875
|
+
dataframe = self.to_pandas(
|
|
876
|
+
connection=connection,
|
|
877
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
878
|
+
**kwargs,
|
|
879
|
+
)
|
|
880
|
+
if not dataframe.empty:
|
|
881
|
+
iterator_data.append(dataframe)
|
|
882
|
+
return iter(iterator_data)
|
|
883
|
+
|
|
884
|
+
def create_iter(
|
|
885
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
886
|
+
) -> (
|
|
887
|
+
Iterator[dict | Exception]
|
|
888
|
+
| Iterator[tuple | Exception]
|
|
889
|
+
| Iterator[Table]
|
|
890
|
+
| Iterator[DataFrame]
|
|
891
|
+
):
|
|
892
|
+
"""The interface used by ResultSet to create an iterator for this ResultBatch."""
|
|
893
|
+
iter_unit: IterUnit = kwargs.pop("iter_unit", IterUnit.ROW_UNIT)
|
|
894
|
+
force_microsecond_precision: bool = kwargs.pop(
|
|
895
|
+
"force_microsecond_precision", False
|
|
896
|
+
)
|
|
897
|
+
if iter_unit == IterUnit.TABLE_UNIT:
|
|
898
|
+
structure = kwargs.pop("structure", "pandas")
|
|
899
|
+
if structure == "pandas":
|
|
900
|
+
return self._get_pandas_iter(
|
|
901
|
+
connection=connection,
|
|
902
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
903
|
+
**kwargs,
|
|
904
|
+
)
|
|
905
|
+
else:
|
|
906
|
+
return self._get_arrow_iter(
|
|
907
|
+
connection=connection,
|
|
908
|
+
force_microsecond_precision=force_microsecond_precision,
|
|
909
|
+
)
|
|
910
|
+
else:
|
|
911
|
+
return self._create_iter(iter_unit=iter_unit, connection=connection)
|
|
912
|
+
|
|
913
|
+
def populate_data(
|
|
914
|
+
self, connection: SnowflakeConnection | None = None, **kwargs
|
|
915
|
+
) -> Self:
|
|
916
|
+
self._data = self._download(connection=connection).content
|
|
917
|
+
return self
|