diff --git a/src/snowflake/connector/options.py b/src/snowflake/connector/options.py index 6aea0ee34..be9f73cc9 100644 --- a/src/snowflake/connector/options.py +++ b/src/snowflake/connector/options.py @@ -7,7 +7,7 @@ import importlib import os import warnings -from importlib.metadata import distributions +from importlib.metadata import PackageNotFoundError, distribution from logging import getLogger from types import ModuleType from typing import Union @@ -85,13 +85,13 @@ def _import_or_missing_pandas_option() -> ( os.environ["ARROW_DEFAULT_MEMORY_POOL"] = "system" # Check whether we have the currently supported pyarrow installed - installed_packages = { - package.metadata["Name"]: package for package in distributions() - } - if {"pyarrow", "snowflake-connector-python"} <= installed_packages.keys(): - dependencies = installed_packages[ - "snowflake-connector-python" - ].metadata.get_all("Requires-Dist", []) + try: + pyarrow_dist = distribution("pyarrow") + snowflake_connector_dist = distribution("snowflake-connector-python") + + dependencies = snowflake_connector_dist.metadata.get_all( + "Requires-Dist", [] + ) pandas_pyarrow_extra = None for dependency in dependencies: dep = Requirement(dependency) @@ -103,16 +103,15 @@ def _import_or_missing_pandas_option() -> ( pandas_pyarrow_extra = dep break - installed_pyarrow_version = installed_packages["pyarrow"].version + installed_pyarrow_version = pyarrow_dist.version if not pandas_pyarrow_extra.specifier.contains(installed_pyarrow_version): warn_incompatible_dep( "pyarrow", installed_pyarrow_version, pandas_pyarrow_extra ) - else: + except PackageNotFoundError as e: logger.info( - "Cannot determine if compatible pyarrow is installed because of missing package(s) from " - "{}".format(list(installed_packages.keys())) + f"Cannot determine if compatible pyarrow is installed because of missing package(s): {e}" ) return pandas, pyarrow, True except ImportError: diff --git a/test/integ/pandas/test_unit_options.py b/test/integ/pandas/test_unit_options.py index 473212c9f..e992b2cb2 100644 --- a/test/integ/pandas/test_unit_options.py +++ b/test/integ/pandas/test_unit_options.py @@ -18,7 +18,7 @@ MissingPandas = None _import_or_missing_pandas_option = None -from importlib.metadata import distributions +from importlib.metadata import PackageNotFoundError, distribution @pytest.mark.skipif( @@ -30,18 +30,15 @@ def test_pandas_option_reporting(caplog): This issue was brought to attention in: https://github.com/snowflakedb/snowflake-connector-python/issues/412 """ - modified_distributions = list( - d - for d in distributions() - if d.metadata["Name"] - not in ( - "pyarrow", - "snowflake-connecctor-python", - ) - ) + + def modified_distribution(name, *args, **kwargs): + if name in ["pyarrow", "snowflake-connector-python"]: + raise PackageNotFoundError("TestErrorMessage") + return distribution(name, *args, **kwargs) + with mock.patch( - "snowflake.connector.options.distributions", - return_value=modified_distributions, + "snowflake.connector.options.distribution", + wraps=modified_distribution, ): caplog.set_level(logging.DEBUG, "snowflake.connector") pandas, pyarrow, installed_pandas = _import_or_missing_pandas_option() @@ -49,6 +46,7 @@ def test_pandas_option_reporting(caplog): assert not isinstance(pandas, MissingPandas) assert not isinstance(pyarrow, MissingPandas) assert ( - "Cannot determine if compatible pyarrow is installed because of missing package(s) " - "from " - ) in caplog.text + "Cannot determine if compatible pyarrow is installed because of missing package(s)" + in caplog.text + ) + assert "TestErrorMessage" in caplog.text