diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py index d856c19852e7..d9afc9781532 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py @@ -16,10 +16,8 @@ from typing import Any import packaging.version - from google.cloud.bigquery import exceptions - _MIN_PYARROW_VERSION = packaging.version.Version("3.0.0") _MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0") _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0") @@ -27,6 +25,7 @@ _MIN_PANDAS_VERSION_RANGE = packaging.version.Version("1.5.0") _MIN_PYARROW_VERSION_RANGE = packaging.version.Version("10.0.1") +_MIN_PANDAS_GBQ_DELEGATION_VERSION = packaging.version.Version("1.0.0") class PyarrowVersions: @@ -247,3 +246,52 @@ def try_import(self, raise_if_error: bool = False) -> Any: and PYARROW_VERSIONS.try_import() is not None and PYARROW_VERSIONS.installed_version >= _MIN_PYARROW_VERSION_RANGE ) + + +class PandasGBQVersions: + """Version and delegation comparisons for pandas-gbq package.""" + + def __init__(self): + self._installed_version = None + self._delegation_api_version = None + + @property + def installed_version(self) -> packaging.version.Version: + """Return the parsed version of pandas-gbq.""" + if self._installed_version is not None: + return self._installed_version + + try: + import pandas_gbq # type: ignore + + self._installed_version = packaging.version.parse( + getattr(pandas_gbq, "__version__", "0.0.0") + ) + except Exception: + self._installed_version = packaging.version.parse("0.0.0") + return self._installed_version + + @property + def delegation_api_version(self) -> packaging.version.Version: + """Return the delegation API version of pandas-gbq if installed, otherwise 0.0.0.""" + if self._delegation_api_version is not None: + return self._delegation_api_version + + try: + import pandas_gbq # type: ignore + + raw_version = getattr( + pandas_gbq, "_internal_delegation_api_version", "0.0.0" + ) + self._delegation_api_version = packaging.version.parse(str(raw_version)) + except Exception: + self._delegation_api_version = packaging.version.parse("0.0.0") + return self._delegation_api_version + + @property + def is_delegation_supported(self) -> bool: + """True if the installed pandas-gbq version supports query delegation API.""" + return self.delegation_api_version >= _MIN_PANDAS_GBQ_DELEGATION_VERSION + + +PANDAS_GBQ_VERSIONS = PandasGBQVersions() diff --git a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py index 8379c87c18e0..aec553bb7b19 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys +import types from unittest import mock +import packaging.version import pytest try: @@ -31,8 +34,8 @@ except ImportError: pandas = None -from google.cloud.bigquery import _versions_helpers -from google.cloud.bigquery import exceptions +from google import cloud +from google.cloud.bigquery import _versions_helpers, exceptions @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") @@ -59,20 +62,19 @@ def test_try_import_raises_error_w_legacy_pyarrow(): versions.try_import(raise_if_error=True) -@pytest.mark.skipif( - pyarrow is not None, - reason="pyarrow is installed, but this test needs it not to be", -) def test_try_import_raises_error_w_no_pyarrow(): versions = _versions_helpers.PyarrowVersions() - with pytest.raises(exceptions.LegacyPyarrowError): - versions.try_import(raise_if_error=True) + with mock.patch.dict(sys.modules, {"pyarrow": None}): + assert versions.try_import(raise_if_error=False) is None + with pytest.raises(exceptions.LegacyPyarrowError): + versions.try_import(raise_if_error=True) @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") def test_installed_pyarrow_version_returns_cached(): versions = _versions_helpers.PyarrowVersions() versions._installed_version = object() + assert versions.installed_version is versions._installed_version @@ -122,17 +124,27 @@ def test_returns_none_with_legacy_bqstorage(): assert bq_storage is None -@pytest.mark.skipif( - bigquery_storage is not None, - reason="Tests behavior when `google-cloud-bigquery-storage` isn't installed", -) def test_returns_none_with_bqstorage_uninstalled(): - try: - bqstorage_versions = _versions_helpers.BQStorageVersions() - bq_storage = bqstorage_versions.try_import() - except exceptions.LegacyBigQueryStorageError: # pragma: NO COVER - raise ("NotFound error raised when raise_if_error == False.") - assert bq_storage is None + versions = _versions_helpers.BQStorageVersions() + with ( + mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}), + mock.patch.dict(cloud.__dict__), + ): + cloud.__dict__.pop("bigquery_storage", None) + + assert versions.try_import() is None + + +def test_raises_error_with_bqstorage_uninstalled(): + versions = _versions_helpers.BQStorageVersions() + with ( + mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}), + mock.patch.dict(cloud.__dict__), + ): + cloud.__dict__.pop("bigquery_storage", None) + + with pytest.raises(exceptions.BigQueryStorageNotFoundError): + versions.try_import(raise_if_error=True) @pytest.mark.skipif( @@ -153,6 +165,7 @@ def test_raises_error_w_unknown_bqstorage_version(): def test_installed_bqstorage_version_returns_cached(): bqstorage_versions = _versions_helpers.BQStorageVersions() bqstorage_versions._installed_version = object() + assert bqstorage_versions.installed_version is bqstorage_versions._installed_version @@ -220,20 +233,19 @@ def test_try_import_raises_error_w_legacy_pandas(): versions.try_import(raise_if_error=True) -@pytest.mark.skipif( - pandas is not None, - reason="pandas is installed, but this test needs it not to be", -) def test_try_import_raises_error_w_no_pandas(): versions = _versions_helpers.PandasVersions() - with pytest.raises(exceptions.LegacyPandasError): - versions.try_import(raise_if_error=True) + with mock.patch.dict(sys.modules, {"pandas": None}): + assert versions.try_import(raise_if_error=False) is None + with pytest.raises(exceptions.LegacyPandasError): + versions.try_import(raise_if_error=True) @pytest.mark.skipif(pandas is None, reason="pandas is not installed") def test_installed_pandas_version_returns_cached(): versions = _versions_helpers.PandasVersions() versions._installed_version = object() + assert versions.installed_version is versions._installed_version @@ -246,3 +258,146 @@ def test_installed_pandas_version_returns_parsed_version(): assert version.major == 1 assert version.minor == 1 assert version.micro == 0 + + +def test_installed_pandas_gbq_version_returns_cached(): + versions = _versions_helpers.PandasGBQVersions() + versions._installed_version = object() + + assert versions.installed_version is versions._installed_version + + +def test_installed_pandas_gbq_version_returns_parsed_version(): + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.__version__ = "1.2.3" + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.installed_version + + assert version.major == 1 + assert version.minor == 2 + assert version.micro == 3 + + +def test_installed_pandas_gbq_version_falls_back_on_import_error(): + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": None}): + version = versions.installed_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_installed_pandas_gbq_version_falls_back_on_missing_version_attribute(): + mock_pandas_gbq = types.ModuleType("pandas_gbq") + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.installed_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_installed_pandas_gbq_version_falls_back_on_other_error(): + class CorruptPandasGBQ: + @property + def __version__(self): + raise TypeError("Corrupted package") + + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": CorruptPandasGBQ()}): + version = versions.installed_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_pandas_gbq_delegation_api_version_returns_cached(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = object() + + assert versions.delegation_api_version is versions._delegation_api_version + + +def test_pandas_gbq_delegation_api_version_returns_value(): + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq._internal_delegation_api_version = "1.2.3" + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.delegation_api_version + + assert version.major == 1 + assert version.minor == 2 + assert version.micro == 3 + + +def test_pandas_gbq_delegation_api_version_supports_integer_value(): + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq._internal_delegation_api_version = 1 + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.delegation_api_version + + assert version.major == 1 + assert version.minor == 0 + assert version.micro == 0 + + +def test_pandas_gbq_delegation_api_version_falls_back_on_missing_attribute(): + mock_pandas_gbq = types.ModuleType("pandas_gbq") + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.delegation_api_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_pandas_gbq_delegation_api_version_falls_back_on_import_error(): + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": None}): + version = versions.delegation_api_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_pandas_gbq_delegation_api_version_falls_back_on_other_error(): + class CorruptPandasGBQ: + @property + def _internal_delegation_api_version(self): + raise TypeError("Corrupted package") + + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": CorruptPandasGBQ()}): + version = versions.delegation_api_version + + assert version.major == 0 + assert version.minor == 0 + assert version.micro == 0 + + +def test_pandas_gbq_is_delegation_supported_true(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = packaging.version.Version("1.0.0") + + assert versions.is_delegation_supported is True + + +def test_pandas_gbq_is_delegation_supported_with_sub_minor_version(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = packaging.version.Version("1.2.0") + + assert versions.is_delegation_supported is True + + +def test_pandas_gbq_is_delegation_supported_false(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = packaging.version.Version("0.9.9") + + assert versions.is_delegation_supported is False