diff --git a/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py b/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py index 09511b201d..001c2266a5 100644 --- a/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py +++ b/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py @@ -13,7 +13,7 @@ from snowflake.connector import ProgrammingError from snowflake.connector.cursor import SnowflakeCursor -from snowflake.connector.options import pyarrow +from snowflake.snowpark._internal.options import pyarrow from snowflake.snowpark._internal.analyzer.binary_plan_node import ( AsOf, Except, diff --git a/src/snowflake/snowpark/_internal/data_source/datasource_reader.py b/src/snowflake/snowpark/_internal/data_source/datasource_reader.py index de72d40f86..a2a1c83b40 100644 --- a/src/snowflake/snowpark/_internal/data_source/datasource_reader.py +++ b/src/snowflake/snowpark/_internal/data_source/datasource_reader.py @@ -11,7 +11,7 @@ from snowflake.snowpark._internal.data_source.drivers.base_driver import BaseDriver from snowflake.snowpark.exceptions import SnowparkDataframeReaderException from snowflake.snowpark.types import StructType -from snowflake.connector.options import pandas as pd +from snowflake.snowpark._internal.options import pandas as pd import logging logger = logging.getLogger(__name__) diff --git a/src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py b/src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py index 88a65e8825..eafa31def5 100644 --- a/src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py +++ b/src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py @@ -4,7 +4,7 @@ from enum import Enum import datetime from typing import Dict, List, Callable, Any, Optional, TYPE_CHECKING -from snowflake.connector.options import pandas as pd +from snowflake.snowpark._internal.options import pandas as pd from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted from snowflake.snowpark._internal.data_source.datasource_typing import ( diff --git a/src/snowflake/snowpark/_internal/event_table_telemetry.py b/src/snowflake/snowpark/_internal/event_table_telemetry.py index a1b70b7e4d..846d5c7d91 100644 --- a/src/snowflake/snowpark/_internal/event_table_telemetry.py +++ b/src/snowflake/snowpark/_internal/event_table_telemetry.py @@ -8,10 +8,13 @@ from abc import ABC from logging import getLogger from typing import Dict, Optional, Tuple -from snowflake.connector.options import MissingOptionalDependency, ModuleLikeObject import snowflake.snowpark import requests +from snowflake.snowpark._internal.options import ( + MissingOptionalDependency, + ModuleLikeObject, +) from snowflake.snowpark._internal.utils import parse_table_name _logger = getLogger(__name__) diff --git a/src/snowflake/snowpark/_internal/options.py b/src/snowflake/snowpark/_internal/options.py new file mode 100644 index 0000000000..9936dd4363 --- /dev/null +++ b/src/snowflake/snowpark/_internal/options.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# + +from snowflake.connector.version import VERSION as connector_version + +IS_V5_DRIVER: bool = connector_version[0] >= 5 + +if IS_V5_DRIVER: + from snowflake.connector._common.extras import pandas # noqa: F401 + from snowflake.connector._common.extras import ModuleLikeObject # noqa: F401 + from snowflake.connector._common.extras import installed_pandas # noqa: F401 + from snowflake.connector._common.extras import ( + MissingOptionalDependency, + pyarrow, + installed_pyarrow, + ) +else: + from snowflake.connector.options import pandas # noqa: F401 + from snowflake.connector.options import ModuleLikeObject # noqa: F401 + from snowflake.connector.options import installed_pandas # noqa: F401 + from snowflake.connector.options import ( + MissingOptionalDependency, + MissingPandas, + pyarrow, + ) + + # connector.options (v4) never exported installed_pyarrow as its own name. + installed_pyarrow: bool = not isinstance(pyarrow, MissingOptionalDependency) + + +def _missing_pandas() -> MissingOptionalDependency: + # v4 has no __init__ override (no-arg-subclass only); v5 deleted + # MissingPandas in favor of the positional-arg form. + if IS_V5_DRIVER: + return MissingOptionalDependency("pandas") + return MissingPandas() diff --git a/src/snowflake/snowpark/_internal/server_connection.py b/src/snowflake/snowpark/_internal/server_connection.py index 94df8ceb9c..add6656e0c 100644 --- a/src/snowflake/snowpark/_internal/server_connection.py +++ b/src/snowflake/snowpark/_internal/server_connection.py @@ -28,8 +28,7 @@ from snowflake.connector.constants import FIELD_ID_TO_NAME from snowflake.connector.cursor import ResultMetadata, SnowflakeCursor from snowflake.connector.errors import Error, NotSupportedError, ProgrammingError -from snowflake.connector.network import ReauthenticationRequest -from snowflake.connector.options import pandas +from snowflake.snowpark._internal.options import pandas from snowflake.snowpark._internal.analyzer.analyzer_utils import ( quote_name_without_upper_casing, ) @@ -54,6 +53,7 @@ get_plan_telemetry_metrics, ) from snowflake.snowpark._internal.utils import ( + IS_V5_DRIVER, create_rlock, create_thread_local, escape_quotes, @@ -69,6 +69,11 @@ result_set_to_rows, unwrap_stage_location_single_quote, ) + +if IS_V5_DRIVER: + from snowflake.connector.errors import ReauthenticationRequest +else: + from snowflake.connector.network import ReauthenticationRequest from snowflake.snowpark import context from snowflake.snowpark.async_job import AsyncJob, _AsyncResultType from snowflake.snowpark.query_history import QueryListener, QueryRecord diff --git a/src/snowflake/snowpark/_internal/telemetry.py b/src/snowflake/snowpark/_internal/telemetry.py index 80f16d395e..4d012df803 100644 --- a/src/snowflake/snowpark/_internal/telemetry.py +++ b/src/snowflake/snowpark/_internal/telemetry.py @@ -12,11 +12,6 @@ from typing import Any, Dict, List, Optional from snowflake.connector import SnowflakeConnection -from snowflake.connector.telemetry import ( - TelemetryClient as PCTelemetryClient, - TelemetryData as PCTelemetryData, - TelemetryField as PCTelemetryField, -) from snowflake.connector.time_util import get_time_millis from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import ( PlanState, @@ -30,6 +25,7 @@ DescribeQueryTelemetryField, ) from snowflake.snowpark._internal.utils import ( + IS_V5_DRIVER, get_application_name, get_os_name, get_python_version, @@ -39,6 +35,19 @@ generate_random_alphanumeric, ) +if IS_V5_DRIVER: + from snowflake.connector._common.telemetry import ( + TelemetryClient as PCTelemetryClient, + TelemetryData as PCTelemetryData, + TelemetryField as PCTelemetryField, + ) +else: + from snowflake.connector.telemetry import ( + TelemetryClient as PCTelemetryClient, + TelemetryData as PCTelemetryData, + TelemetryField as PCTelemetryField, + ) + try: import psutil diff --git a/src/snowflake/snowpark/_internal/type_utils.py b/src/snowflake/snowpark/_internal/type_utils.py index c7ca2a2a3e..b27e50642a 100644 --- a/src/snowflake/snowpark/_internal/type_utils.py +++ b/src/snowflake/snowpark/_internal/type_utils.py @@ -36,7 +36,7 @@ import snowflake.snowpark.types # type: ignore from snowflake.connector.constants import FIELD_ID_TO_NAME from snowflake.connector.cursor import ResultMetadata -from snowflake.connector.options import installed_pandas, pandas +from snowflake.snowpark._internal.options import installed_pandas, pandas from snowflake.snowpark._internal.utils import quote_name from snowflake.snowpark.row import Row from snowflake.snowpark.types import ( diff --git a/src/snowflake/snowpark/_internal/udf_utils.py b/src/snowflake/snowpark/_internal/udf_utils.py index 7b3b947904..de7d82a0d8 100644 --- a/src/snowflake/snowpark/_internal/udf_utils.py +++ b/src/snowflake/snowpark/_internal/udf_utils.py @@ -29,7 +29,7 @@ from packaging.requirements import Requirement import snowflake.snowpark -from snowflake.connector.options import installed_pandas, pandas +from snowflake.snowpark._internal.options import installed_pandas, pandas from snowflake.snowpark._internal import code_generation, type_utils from snowflake.snowpark._internal.analyzer.datatype_mapper import to_sql, to_sql_no_cast from snowflake.snowpark._internal.telemetry import TelemetryField diff --git a/src/snowflake/snowpark/_internal/utils.py b/src/snowflake/snowpark/_internal/utils.py index 99dc105b37..74c390ad8c 100644 --- a/src/snowflake/snowpark/_internal/utils.py +++ b/src/snowflake/snowpark/_internal/utils.py @@ -54,12 +54,20 @@ from snowflake.connector.constants import FIELD_ID_TO_NAME from snowflake.connector.cursor import ResultMetadata, SnowflakeCursor from snowflake.connector.description import OPERATING_SYSTEM, PLATFORM -from snowflake.connector.options import MissingOptionalDependency, ModuleLikeObject from snowflake.connector.version import VERSION as connector_version from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages from snowflake.snowpark.row import Row from snowflake.snowpark.version import VERSION as snowpark_version +IS_V5_DRIVER: bool = connector_version[0] >= 5 + +from snowflake.snowpark._internal.options import ( # noqa: E402 + MissingOptionalDependency, + ModuleLikeObject, + installed_pandas, +) + + if TYPE_CHECKING: from snowflake.snowpark._internal.analyzer.snowflake_plan import ( SnowflakePlan, @@ -240,24 +248,6 @@ SUPPORTED_TABLE_TYPES = ["temp", "temporary", "transient"] -# TODO: merge fixed pandas importer changes to connector. -def _pandas_importer(): # noqa: E302 - """Helper function to lazily import pandas and return MissingPandas if not installed.""" - from snowflake.connector.options import MissingPandas - - pandas = MissingPandas() - try: - pandas = importlib.import_module("pandas") - # since we enable relative imports without dots this import gives us an issues when ran from test directory - from pandas import DataFrame # NOQA - except ImportError: # pragma: no cover - pass # pragma: no cover - return pandas - - -pandas = _pandas_importer() -installed_pandas = not isinstance(pandas, MissingOptionalDependency) - class TempObjectType(Enum): TABLE = "TABLE" diff --git a/src/snowflake/snowpark/async_job.py b/src/snowflake/snowpark/async_job.py index 7bae2324fc..d4919eed10 100644 --- a/src/snowflake/snowpark/async_job.py +++ b/src/snowflake/snowpark/async_job.py @@ -8,7 +8,7 @@ import snowflake.snowpark from snowflake.connector.errors import DatabaseError -from snowflake.connector.options import pandas +from snowflake.snowpark._internal.options import pandas from snowflake.snowpark._internal.analyzer.analyzer_utils import result_scan_statement from snowflake.snowpark._internal.analyzer.snowflake_plan import Query from snowflake.snowpark._internal.utils import ( diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 42e7b4ef65..dd69e90d65 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -31,7 +31,7 @@ import snowflake.snowpark import snowflake.snowpark.context as context import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto -from snowflake.connector.options import installed_pandas, pandas, pyarrow +from snowflake.snowpark._internal.options import installed_pandas, pandas, pyarrow from snowflake.snowpark._internal.analyzer.binary_plan_node import ( AsOf, diff --git a/src/snowflake/snowpark/mock/_options.py b/src/snowflake/snowpark/mock/_options.py index 7ae10dcda5..5ac0d3d807 100644 --- a/src/snowflake/snowpark/mock/_options.py +++ b/src/snowflake/snowpark/mock/_options.py @@ -4,26 +4,35 @@ import importlib -from snowflake.connector.options import MissingOptionalDependency, MissingPandas +from snowflake.snowpark._internal.options import ( + MissingOptionalDependency, + _missing_pandas, +) +from snowflake.snowpark._internal.utils import IS_V5_DRIVER try: import pandas installed_pandas = True except ImportError: - pandas = MissingPandas() + pandas = _missing_pandas() installed_pandas = False -class MissingNumpy(MissingOptionalDependency): - """The class is specifically for numpy optional dependency.""" +if IS_V5_DRIVER: + from snowflake.connector._common.extras import numpy - _dep_name = "numpy" + installed_numpy = not isinstance(numpy, MissingOptionalDependency) +else: + class MissingNumpy(MissingOptionalDependency): + """The class is specifically for numpy optional dependency.""" -try: - numpy = importlib.import_module("numpy") - installed_numpy = True -except ImportError: - numpy = MissingNumpy() - installed_numpy = False + _dep_name = "numpy" + + try: + numpy = importlib.import_module("numpy") + installed_numpy = True + except ImportError: + numpy = MissingNumpy() + installed_numpy = False diff --git a/src/snowflake/snowpark/relational_grouped_dataframe.py b/src/snowflake/snowpark/relational_grouped_dataframe.py index 189d37d66d..a0fc656b94 100644 --- a/src/snowflake/snowpark/relational_grouped_dataframe.py +++ b/src/snowflake/snowpark/relational_grouped_dataframe.py @@ -8,7 +8,7 @@ from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto import snowflake.snowpark.context as context -from snowflake.connector.options import pandas +from snowflake.snowpark._internal.options import pandas from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted from snowflake.snowpark import functions from snowflake.snowpark._internal.analyzer.expression import ( diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index aea7b7e5de..a461116324 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -42,7 +42,7 @@ import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto import snowflake.snowpark.context as context from snowflake.connector import ProgrammingError, SnowflakeConnection -from snowflake.connector.options import installed_pandas, pandas, pyarrow +from snowflake.snowpark._internal.options import installed_pandas, pandas, pyarrow from snowflake.connector.pandas_tools import write_pandas from snowflake.snowpark import UDFProfiler diff --git a/src/snowflake/snowpark/types.py b/src/snowflake/snowpark/types.py index 3d97816615..a3f8e10241 100644 --- a/src/snowflake/snowpark/types.py +++ b/src/snowflake/snowpark/types.py @@ -16,11 +16,8 @@ import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto # Use correct version from here: -from snowflake.snowpark._internal.utils import ( - installed_pandas, - pandas, - quote_name, -) +from snowflake.snowpark._internal.options import installed_pandas, pandas +from snowflake.snowpark._internal.utils import quote_name # TODO: connector installed_pandas is broken. If pyarrow is not installed, but pandas is this function returns the wrong answer. # The core issue is that in the connector detection of both pandas/arrow are mixed, which is wrong. diff --git a/tests/integ/scala/test_datatype_suite.py b/tests/integ/scala/test_datatype_suite.py index 950e05a060..6b1bbe7cc0 100644 --- a/tests/integ/scala/test_datatype_suite.py +++ b/tests/integ/scala/test_datatype_suite.py @@ -12,7 +12,7 @@ from unittest import mock import snowflake.snowpark.context as context -from snowflake.connector.options import installed_pandas +from snowflake.snowpark._internal.options import installed_pandas from snowflake.snowpark import Row from snowflake.snowpark.dataframe import DataFrame from snowflake.snowpark.exceptions import SnowparkSQLException diff --git a/tests/integ/scala/test_update_delete_merge_suite.py b/tests/integ/scala/test_update_delete_merge_suite.py index ed5e904aff..534c0824b0 100644 --- a/tests/integ/scala/test_update_delete_merge_suite.py +++ b/tests/integ/scala/test_update_delete_merge_suite.py @@ -8,7 +8,7 @@ import pytest -from snowflake.connector.options import installed_pandas +from snowflake.snowpark._internal.options import installed_pandas from snowflake.snowpark import ( DeleteResult, MergeResult, diff --git a/tests/integ/test_cte.py b/tests/integ/test_cte.py index fcb9a6694a..fab97afc67 100644 --- a/tests/integ/test_cte.py +++ b/tests/integ/test_cte.py @@ -11,7 +11,7 @@ import pytest from snowflake.connector.errors import ProgrammingError -from snowflake.connector.options import installed_pandas +from snowflake.snowpark._internal.options import installed_pandas from snowflake.snowpark import Window from snowflake.snowpark._internal.analyzer import analyzer from snowflake.snowpark._internal.analyzer.snowflake_plan import PlanQueryType, Query diff --git a/tests/integ/test_function.py b/tests/integ/test_function.py index 455707e8bf..81cc445293 100644 --- a/tests/integ/test_function.py +++ b/tests/integ/test_function.py @@ -165,7 +165,7 @@ year, ) import functools -from snowflake.connector.options import installed_pandas +from snowflake.snowpark._internal.options import installed_pandas from snowflake.snowpark.functions import udf, vectorized from snowflake.snowpark.udf import UserDefinedFunction from snowflake.snowpark.types import ( diff --git a/tests/unit/test_internal_utils.py b/tests/unit/test_internal_utils.py index f09cfc63d3..6d41a8a40e 100644 --- a/tests/unit/test_internal_utils.py +++ b/tests/unit/test_internal_utils.py @@ -4,11 +4,9 @@ import concurrent.futures import random import pytest -from snowflake.connector.options import MissingPandas -from snowflake.snowpark._internal import utils +from snowflake.snowpark._internal import options, utils from snowflake.snowpark._internal.utils import ( - _pandas_importer, generate_random_alphanumeric, split_snowflake_identifier_with_dot, ) @@ -148,16 +146,6 @@ def test_normalize_path_escapes_backslash_and_quote(raw_path, is_local): ), f"decoded={decoded!r} does not end with {expected_tail!r}" -def test__pandas_importer(): - imported_pandas = _pandas_importer() - try: - import pandas - - assert imported_pandas == pandas - except ImportError: - assert isinstance(imported_pandas, MissingPandas) - - def test_generate_random_alphanumeric(): random.seed(42) random_string1 = generate_random_alphanumeric() @@ -210,3 +198,9 @@ def test_generate_random_alphanumeric(): ) def test_split_snowflake_identifier_with_dot(string, expected_result): assert split_snowflake_identifier_with_dot(string) == expected_result + + +def test_missing_pandas(): + result = options._missing_pandas() + assert isinstance(result, options.MissingOptionalDependency) + assert result._dep_name == "pandas" diff --git a/tests/unit/test_server_connection.py b/tests/unit/test_server_connection.py index a53d42dc16..f2a3761f5e 100644 --- a/tests/unit/test_server_connection.py +++ b/tests/unit/test_server_connection.py @@ -10,9 +10,14 @@ import pytest -from snowflake.connector.network import ReauthenticationRequest from snowflake.snowpark import Session from snowflake.snowpark._internal.analyzer.snowflake_plan import Query, SnowflakePlan +from snowflake.snowpark._internal.utils import IS_V5_DRIVER + +if IS_V5_DRIVER: + from snowflake.connector.errors import ReauthenticationRequest +else: + from snowflake.connector.network import ReauthenticationRequest from snowflake.snowpark.exceptions import ( SnowparkFetchDataException, SnowparkQueryCancelledException,