diff --git a/packages/google-cloud-bigquery/.coveragerc b/packages/google-cloud-bigquery/.coveragerc index e78e7a931e09..f8b9e9e79f03 100644 --- a/packages/google-cloud-bigquery/.coveragerc +++ b/packages/google-cloud-bigquery/.coveragerc @@ -12,3 +12,5 @@ exclude_lines = pragma: (no cover|NO COVER) # Ignore debug-only repr def __repr__ + except ImportError: + except ImportError as .*: 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..12ccf62512ef 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") @@ -247,3 +245,45 @@ 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 + + return packaging.version.parse(getattr(pandas_gbq, "__version__", "0.0.0")) + except Exception: + return packaging.version.parse("0.0.0") + + @property + def delegation_api_version(self) -> int: + """Return the delegation API version of pandas-gbq if installed, otherwise 0.""" + if self._delegation_api_version is not None: + return self._delegation_api_version + + try: + import pandas_gbq # type: ignore + + return int(getattr(pandas_gbq, "_internal_delegation_api_version", 0)) + except Exception: + return 0 + + @property + def is_delegation_supported(self) -> bool: + """True if the installed pandas-gbq version supports query delegation API (version >= 1).""" + return self.delegation_api_version >= 1 + + +PANDAS_GBQ_VERSIONS = PandasGBQVersions() diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py index 30bc9d27a8b6..5d951a4f3be3 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py @@ -23,10 +23,10 @@ from __future__ import print_function -import re import ast import copy import functools +import re import sys import time import warnings @@ -39,14 +39,12 @@ except ImportError: raise ImportError("This module can only be loaded in IPython.") -from google.api_core import client_info -from google.api_core import client_options -from google.api_core.exceptions import NotFound import google.auth # type: ignore -from google.cloud import bigquery import google.cloud.bigquery.dataset -from google.cloud.bigquery import _versions_helpers -from google.cloud.bigquery import exceptions +from google.api_core import client_info, client_options +from google.api_core.exceptions import NotFound +from google.cloud import bigquery +from google.cloud.bigquery import _versions_helpers, exceptions from google.cloud.bigquery.dbapi import _helpers from google.cloud.bigquery.magics import line_arg_parser as lap @@ -231,7 +229,7 @@ def progress_bar_type(self, value): # their code. if bigquery_magics is not None: context = bigquery_magics.context -else: +else: # pragma: NO COVER context = Context() @@ -498,8 +496,9 @@ def _cell_magic(line, query): raise rebranded_error from exc except lap.exceptions.ParseError as exc: rebranded_error = ValueError( - "Unrecognized input, are option values correct? " - "Error details: {}".format(exc.args[0]) + "Unrecognized input, are option values correct? Error details: {}".format( + exc.args[0] + ) ) raise rebranded_error from exc diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 870cdcc5d2ab..70496332d6bd 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -21,9 +21,8 @@ import functools import operator import typing -from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union, Sequence - import warnings +from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, Union try: import pandas # type: ignore @@ -56,30 +55,33 @@ _read_wkt = wkt.loads import google.api_core.exceptions -from google.api_core.page_iterator import HTTPIterator - import google.cloud._helpers # type: ignore -from google.cloud.bigquery import _helpers -from google.cloud.bigquery import _pandas_helpers -from google.cloud.bigquery import _versions_helpers +from google.api_core.page_iterator import HTTPIterator +from google.cloud.bigquery import ( + _helpers, + _pandas_helpers, + _string_references, + _versions_helpers, + external_config, +) from google.cloud.bigquery import exceptions as bq_exceptions +from google.cloud.bigquery import schema as _schema from google.cloud.bigquery._tqdm_helpers import get_progress_bar from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration from google.cloud.bigquery.enums import DefaultPandasDTypes from google.cloud.bigquery.external_config import ExternalConfig -from google.cloud.bigquery import schema as _schema -from google.cloud.bigquery.schema import _build_schema_resource -from google.cloud.bigquery.schema import _parse_schema_resource -from google.cloud.bigquery.schema import _to_schema_fields -from google.cloud.bigquery import external_config -from google.cloud.bigquery import _string_references +from google.cloud.bigquery.schema import ( + _build_schema_resource, + _parse_schema_resource, + _to_schema_fields, +) if typing.TYPE_CHECKING: # pragma: NO COVER # Unconditionally import optional dependencies again to tell pytype that # they are not None, avoiding false "no attribute" errors. + import geopandas # type: ignore import pandas import pyarrow - import geopandas # type: ignore from google.cloud import bigquery_storage # type: ignore from google.cloud.bigquery.dataset import DatasetReference @@ -540,9 +542,9 @@ def biglake_configuration(self, value): api_repr = value if value is not None: api_repr = value.to_api_repr() - self._properties[ - self._PROPERTY_TO_API_FIELD["biglake_configuration"] - ] = api_repr + self._properties[self._PROPERTY_TO_API_FIELD["biglake_configuration"]] = ( + api_repr + ) @property def require_partition_filter(self): @@ -556,9 +558,9 @@ def require_partition_filter(self): @require_partition_filter.setter def require_partition_filter(self, value): - self._properties[ - self._PROPERTY_TO_API_FIELD["require_partition_filter"] - ] = value + self._properties[self._PROPERTY_TO_API_FIELD["require_partition_filter"]] = ( + value + ) @property def schema(self): @@ -656,9 +658,9 @@ def encryption_configuration(self, value): api_repr = value if value is not None: api_repr = value.to_api_repr() - self._properties[ - self._PROPERTY_TO_API_FIELD["encryption_configuration"] - ] = api_repr + self._properties[self._PROPERTY_TO_API_FIELD["encryption_configuration"]] = ( + api_repr + ) @property def created(self): @@ -797,7 +799,7 @@ def time_partitioning(self, value): api_repr = value.to_api_repr() elif value is not None: raise ValueError( - "value must be google.cloud.bigquery.table.TimePartitioning " "or None" + "value must be google.cloud.bigquery.table.TimePartitioning or None" ) self._properties[self._PROPERTY_TO_API_FIELD["time_partitioning"]] = api_repr @@ -933,9 +935,9 @@ def expires(self, value): if not isinstance(value, datetime.datetime) and value is not None: raise ValueError("Pass a datetime, or None") value_ms = google.cloud._helpers._millis_from_datetime(value) - self._properties[ - self._PROPERTY_TO_API_FIELD["expires"] - ] = _helpers._str_or_none(value_ms) + self._properties[self._PROPERTY_TO_API_FIELD["expires"]] = ( + _helpers._str_or_none(value_ms) + ) @property def friendly_name(self): @@ -1131,9 +1133,9 @@ def external_data_configuration(self, value): api_repr = value if value is not None: api_repr = value.to_api_repr() - self._properties[ - self._PROPERTY_TO_API_FIELD["external_data_configuration"] - ] = api_repr + self._properties[self._PROPERTY_TO_API_FIELD["external_data_configuration"]] = ( + api_repr + ) @property def snapshot_definition(self) -> Optional["SnapshotDefinition"]: @@ -2709,6 +2711,14 @@ def to_dataframe( is not supported dtype. """ + if not _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported: + warnings.warn( + "Retrieving dataframes via the core client is deprecated. " + "Please install 'pandas-gbq' for the new high-performance backend.", + PendingDeprecationWarning, + stacklevel=2, + ) + _pandas_helpers.verify_pandas_imports() if geography_as_object and shapely is None: @@ -2801,6 +2811,44 @@ def to_dataframe( create_bqstorage_client = False bqstorage_client = None + if _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported: + import pandas_gbq # type: ignore + + if ( + self.client + and hasattr(self.client, "_connection") + and hasattr(self.client._connection, "_client_info") + ): + client_info = self.client._connection._client_info + if client_info: + ua = client_info.user_agent or "" + if "pandas-gbq" not in ua: + pandas_gbq_version = getattr(pandas_gbq, "__version__", "0.0.0") + client_info.user_agent = ( + f"{ua} pandas-gbq/{pandas_gbq_version}".strip() + ) + + return pandas_gbq.pandas.from_row_iterator( + self, + bqstorage_client=bqstorage_client, + dtypes=dtypes, + progress_bar_type=progress_bar_type, + create_bqstorage_client=create_bqstorage_client, + geography_as_object=geography_as_object, + bool_dtype=bool_dtype, + int_dtype=int_dtype, + float_dtype=float_dtype, + string_dtype=string_dtype, + date_dtype=date_dtype, + datetime_dtype=datetime_dtype, + time_dtype=time_dtype, + timestamp_dtype=timestamp_dtype, + range_date_dtype=range_date_dtype, + range_datetime_dtype=range_datetime_dtype, + range_timestamp_dtype=range_timestamp_dtype, + timeout=timeout, + ) + record_batch = self.to_arrow( progress_bar_type=progress_bar_type, bqstorage_client=bqstorage_client, @@ -2990,8 +3038,7 @@ def to_geodataframe( ) if not geography_columns: raise TypeError( - "There must be at least one GEOGRAPHY column" - " to create a GeoDataFrame" + "There must be at least one GEOGRAPHY column to create a GeoDataFrame" ) if geography_column: diff --git a/packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py index 34da6370e039..707aa4dc1d45 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py @@ -18,14 +18,13 @@ import decimal import functools import gc +import importlib.metadata as metadata import operator import queue import time +import warnings from typing import Union from unittest import mock -import warnings - -import importlib.metadata as metadata try: import pandas @@ -45,13 +44,13 @@ geopandas = None import pytest - from google import api_core - -from google.cloud.bigquery import exceptions -from google.cloud.bigquery import _pyarrow_helpers -from google.cloud.bigquery import _versions_helpers -from google.cloud.bigquery import schema +from google.cloud.bigquery import ( + _pyarrow_helpers, + _versions_helpers, + exceptions, + schema, +) from google.cloud.bigquery._pandas_helpers import determine_requested_streams pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import() @@ -1831,8 +1830,7 @@ def test__download_table_bqstorage( expected_call_count, expected_maxsize, ): - from google.cloud.bigquery import dataset - from google.cloud.bigquery import table + from google.cloud.bigquery import dataset, table queue_used = None # A reference to the queue used by code under test. @@ -1885,10 +1883,9 @@ def test__download_table_bqstorage_shuts_down_workers( the child threads are also stopped. """ pytest.importorskip("google.cloud.bigquery_storage_v1") - from google.cloud.bigquery import dataset - from google.cloud.bigquery import table import google.cloud.bigquery_storage_v1.reader import google.cloud.bigquery_storage_v1.types + from google.cloud.bigquery import dataset, table monkeypatch.setattr( _versions_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None @@ -2211,10 +2208,10 @@ def test_determine_requested_streams_invalid_max_stream_count(): bigquery_storage is None, reason="Requires google-cloud-bigquery-storage" ) def test__download_table_bqstorage_w_timeout_error(module_under_test): - from google.cloud.bigquery import dataset - from google.cloud.bigquery import table from unittest import mock + from google.cloud.bigquery import dataset, table + mock_bqstorage_client = mock.create_autospec( bigquery_storage.BigQueryReadClient, instance=True ) @@ -2248,10 +2245,10 @@ def slow_download_stream( bigquery_storage is None, reason="Requires google-cloud-bigquery-storage" ) def test__download_table_bqstorage_w_timeout_success(module_under_test): - from google.cloud.bigquery import dataset - from google.cloud.bigquery import table from unittest import mock + from google.cloud.bigquery import dataset, table + mock_bqstorage_client = mock.create_autospec( bigquery_storage.BigQueryReadClient, instance=True ) @@ -2409,3 +2406,38 @@ def test_download_arrow_bqstorage_passes_timeout_to_create_read_session( assert retry_policy is not None # Check if deadline is set correctly in the retry policy assert retry_policy._deadline == timeout + + +@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") +def test_dataframe_to_bq_schema_w_unused_schema_field(module_under_test): + with mock.patch.object(module_under_test, "pandas_gbq", None): + with pytest.raises( + ValueError, match="bq_schema contains fields not present in dataframe" + ): + module_under_test.dataframe_to_bq_schema( + pandas.DataFrame(), (schema.SchemaField("not_in_df", "STRING"),) + ) + + +@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") +@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`") +def test_get_schema_by_pyarrow_bignumeric(module_under_test): + series = pandas.Series([decimal.Decimal("1.12345678901")]) + result = module_under_test._get_schema_by_pyarrow("col", series) + assert result is not None + assert result.field_type == "BIGNUMERIC" + + +@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") +@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`") +def test_get_types_mapper_range_timestamp_mismatch(module_under_test): + if not hasattr(pandas, "ArrowDtype"): + return + range_ts = pandas.ArrowDtype( + pyarrow.struct( + [("start", pyarrow.timestamp("us")), ("end", pyarrow.timestamp("us"))] + ) + ) + mapper = module_under_test.default_types_mapper(range_timestamp_dtype=range_ts) + unmatched_struct = pyarrow.struct([("other", pyarrow.int64())]) + assert mapper(unmatched_struct) is None diff --git a/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py index c12a526de5d3..123b5bfdedc0 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py @@ -35,8 +35,7 @@ def test_bq_to_arrow_scalars(module_under_test): # but we'd like this to map as closely to the BQ Storage API as # possible, which uses the string() dtype, as JSON support in Arrow # predates JSON support in BigQuery by several years. - module_under_test.bq_to_arrow_scalars("JSON")() - == pyarrow.string() + module_under_test.bq_to_arrow_scalars("JSON")() == pyarrow.string() ) assert module_under_test.bq_to_arrow_scalars("UNKNOWN_TYPE") is None @@ -44,3 +43,16 @@ def test_bq_to_arrow_scalars(module_under_test): def test_arrow_scalar_ids_to_bq(module_under_test): assert module_under_test.arrow_scalar_ids_to_bq(pyarrow.bool_().id) == "BOOL" assert module_under_test.arrow_scalar_ids_to_bq("UNKNOWN_TYPE") is None + + +def test_pyarrow_helpers_when_pyarrow_none(module_under_test): + import importlib + import sys + from unittest import mock + + with mock.patch.dict(sys.modules, {"pyarrow": None}): + importlib.reload(module_under_test) + assert module_under_test.pyarrow is None + assert module_under_test.arrow_scalar_ids_to_bq(1) is None + + importlib.reload(module_under_test) 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..06ce47104cb5 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__versions_helpers.py @@ -31,8 +31,7 @@ except ImportError: pandas = None -from google.cloud.bigquery import _versions_helpers -from google.cloud.bigquery import exceptions +from google.cloud.bigquery import _versions_helpers, exceptions @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") @@ -59,14 +58,14 @@ 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(): + import sys + 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") @@ -122,17 +121,29 @@ 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 + import sys + + from google import cloud + + versions = _versions_helpers.BQStorageVersions() + with mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}): + with mock.patch.dict(cloud.__dict__): + cloud.__dict__.pop("bigquery_storage", None) + assert versions.try_import() is None + + +def test_raises_error_with_bqstorage_uninstalled(): + import sys + + from google import cloud + + versions = _versions_helpers.BQStorageVersions() + with mock.patch.dict(sys.modules, {"google.cloud.bigquery_storage": None}): + with 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( @@ -220,14 +231,14 @@ 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(): + import sys + 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") @@ -246,3 +257,108 @@ 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(): + import sys + + 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(): + import sys + + 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_other_error(): + import sys + + # Simulate a corrupted package raising an error on import/property access + 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(): + import sys + + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq._internal_delegation_api_version = 42 + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + version = versions.delegation_api_version + + assert version == 42 + + +def test_pandas_gbq_delegation_api_version_falls_back_on_import_error(): + import sys + + versions = _versions_helpers.PandasGBQVersions() + with mock.patch.dict(sys.modules, {"pandas_gbq": None}): + version = versions.delegation_api_version + + assert version == 0 + + +def test_pandas_gbq_delegation_api_version_falls_back_on_other_error(): + import sys + + 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 == 0 + + +def test_pandas_gbq_is_delegation_supported_true(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = 1 + assert versions.is_delegation_supported is True + + +def test_pandas_gbq_is_delegation_supported_false(): + versions = _versions_helpers.PandasGBQVersions() + versions._delegation_api_version = 0 + assert versions.is_delegation_supported is False diff --git a/packages/google-cloud-bigquery/tests/unit/test_magics.py b/packages/google-cloud-bigquery/tests/unit/test_magics.py index 03a3a2dbbdba..8b4858c1916f 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_magics.py +++ b/packages/google-cloud-bigquery/tests/unit/test_magics.py @@ -147,6 +147,8 @@ def test_context_with_default_credentials(): """When Application Default Credentials are set, the context credentials will be created the first time it is called """ + magics.context._credentials = None + magics.context._project = None assert magics.context._credentials is None assert magics.context._project is None @@ -164,6 +166,16 @@ def test_context_with_default_credentials(): assert default_mock.call_count == 2 +def test_context_fallback_when_bigquery_magics_none(): + ctx = magics.Context() + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + with mock.patch("google.auth.default", return_value=(credentials_mock, "proj-123")): + assert ctx.credentials is credentials_mock + assert ctx.project == "proj-123" + + @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") def test_context_with_default_connection(monkeypatch): @@ -674,9 +686,11 @@ def test_bigquery_magic_with_bqstorage_from_argument( google.cloud.bigquery.job.QueryJob, instance=True ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), warnings.catch_warnings(record=True) as warned: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + warnings.catch_warnings(record=True) as warned, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api", sql) @@ -842,11 +856,12 @@ def test_bigquery_magic_w_max_results_query_job_results_fails(monkeypatch): ) query_job_mock.result.side_effect = [[], OSError] - with pytest.raises( - OSError - ), client_query_patch as client_query_mock, ( - default_patch - ), close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + client_query_patch as client_query_mock, + default_patch, + close_transports_patch as close_transports, + ): client_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", "--max_results=5", sql) @@ -1965,9 +1980,10 @@ def test_bigquery_magic_nonexisting_query_variable(monkeypatch): ip.user_ns.pop("custom_query", None) # Make sure the variable does NOT exist. cell_body = "$custom_query" # Referring to a non-existing variable name. - with pytest.raises( - NameError, match=r".*custom_query does not exist.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r".*custom_query does not exist.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -1988,9 +2004,10 @@ def test_bigquery_magic_empty_query_variable_name(monkeypatch): ) cell_body = "$" # Not referring to any variable (name omitted). - with pytest.raises( - NameError, match=r"(?i).*missing query variable name.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r"(?i).*missing query variable name.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2016,9 +2033,10 @@ def test_bigquery_magic_query_variable_non_string(ipython_ns_cleanup, monkeypatc ip.user_ns["custom_query"] = object() cell_body = "$custom_query" # Referring to a non-string variable. - with pytest.raises( - TypeError, match=r".*must be a string or a bytes-like.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(TypeError, match=r".*must be a string or a bytes-like.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2183,9 +2201,11 @@ def test_bigquery_magic_create_dataset_fails(monkeypatch): autospec=True, ) - with pytest.raises( - OSError - ), create_dataset_if_necessary_patch, close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + create_dataset_if_necessary_patch, + close_transports_patch as close_transports, + ): ip.run_cell_magic( "bigquery", "--destination_table dataset_id.table_id", diff --git a/packages/google-cloud-bigquery/tests/unit/test_table.py b/packages/google-cloud-bigquery/tests/unit/test_table.py index 5701143a62d4..22d59c653357 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_table.py +++ b/packages/google-cloud-bigquery/tests/unit/test_table.py @@ -19,21 +19,16 @@ import time import types import unittest -from unittest import mock import warnings - -import pytest +from unittest import mock import google.api_core.exceptions -from test_utils.imports import maybe_fail_import - -from google.cloud.bigquery import _versions_helpers -from google.cloud.bigquery import exceptions -from google.cloud.bigquery import external_config -from google.cloud.bigquery import schema +import pytest +from google.cloud.bigquery import _versions_helpers, exceptions, external_config, schema +from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.enums import DefaultPandasDTypes from google.cloud.bigquery.table import TableReference -from google.cloud.bigquery.dataset import DatasetReference +from test_utils.imports import maybe_fail_import def _mock_client(): @@ -381,9 +376,7 @@ def test_from_api_repr(self): def test___repr__(self): dataset = DatasetReference("project1", "dataset1") table1 = self._make_one(dataset, "table1") - expected = ( - "TableReference(DatasetReference('project1', 'dataset1'), " "'table1')" - ) + expected = "TableReference(DatasetReference('project1', 'dataset1'), 'table1')" self.assertEqual(repr(table1), expected) def test___str__(self): @@ -414,6 +407,7 @@ def _make_one(self, *args, **kw): def _setUpConstants(self): import datetime + from google.cloud._helpers import UTC self.WHEN_TS = 1437767599.006 @@ -618,10 +612,10 @@ def test_ctor_string(self): self.assertEqual(table.table_id, "some_tbl") def test_ctor_tablelistitem(self): - from google.cloud.bigquery.table import Table, TableListItem - import datetime - from google.cloud._helpers import _millis, UTC + + from google.cloud._helpers import UTC, _millis + from google.cloud.bigquery.table import Table, TableListItem self.WHEN_TS = 1437767599.125 self.EXP_TIME = datetime.datetime(2015, 8, 1, 23, 59, 59, tzinfo=UTC) @@ -818,8 +812,8 @@ def test_schema_setter_valid_mapping_representation(self): def test_props_set_by_server(self): import datetime - from google.cloud._helpers import UTC - from google.cloud._helpers import _millis + + from google.cloud._helpers import UTC, _millis CREATED = datetime.datetime(2015, 7, 29, 12, 13, 22, tzinfo=UTC) MODIFIED = datetime.datetime(2015, 7, 29, 14, 47, 15, tzinfo=UTC) @@ -1162,6 +1156,7 @@ def test_expires_setter_bad_value(self): def test_expires_setter(self): import datetime + from google.cloud._helpers import UTC WHEN = datetime.datetime(2015, 7, 28, 16, 39, tzinfo=UTC) @@ -1374,8 +1369,8 @@ def test_from_api_repr_bare(self): def test_from_api_repr_w_properties(self): import datetime - from google.cloud._helpers import UTC - from google.cloud._helpers import _millis + + from google.cloud._helpers import UTC, _millis RESOURCE = self._make_resource() RESOURCE["view"] = {"query": "select fullname, age from person_ages"} @@ -1389,8 +1384,8 @@ def test_from_api_repr_w_properties(self): def test_from_api_repr_w_partial_streamingbuffer(self): import datetime - from google.cloud._helpers import UTC - from google.cloud._helpers import _millis + + from google.cloud._helpers import UTC, _millis RESOURCE = self._make_resource() self.OLDEST_TIME = datetime.datetime(2015, 8, 1, 23, 59, 59, tzinfo=UTC) @@ -1554,8 +1549,7 @@ def test__build_resource_w_custom_field_not_in__properties(self): table._build_resource(["bad"]) def test_range_partitioning(self): - from google.cloud.bigquery.table import RangePartitioning - from google.cloud.bigquery.table import PartitionRange + from google.cloud.bigquery.table import PartitionRange, RangePartitioning table = self._make_one("proj.dset.tbl") assert table.range_partitioning is None @@ -1588,8 +1582,7 @@ def test_require_partitioning_filter(self): assert table.require_partition_filter is None def test_time_partitioning_getter(self): - from google.cloud.bigquery.table import TimePartitioning - from google.cloud.bigquery.table import TimePartitioningType + from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType dataset = DatasetReference(self.PROJECT, self.DS_ID) table_ref = dataset.table(self.TABLE_NAME) @@ -1646,8 +1639,7 @@ def test_time_partitioning_getter_w_empty(self): self.assertIs(warning.category, PendingDeprecationWarning) def test_time_partitioning_setter(self): - from google.cloud.bigquery.table import TimePartitioning - from google.cloud.bigquery.table import TimePartitioningType + from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType dataset = DatasetReference(self.PROJECT, self.DS_ID) table_ref = dataset.table(self.TABLE_NAME) @@ -1837,9 +1829,7 @@ def test___repr__(self): dataset = DatasetReference("project1", "dataset1") table1 = self._make_one(TableReference(dataset, "table1")) expected = ( - "Table(TableReference(" - "DatasetReference('project1', 'dataset1'), " - "'table1'))" + "Table(TableReference(DatasetReference('project1', 'dataset1'), 'table1'))" ) self.assertEqual(repr(table1), expected) @@ -1903,7 +1893,7 @@ def _call_fut(self, mapping, schema): return _row_from_mapping(mapping, schema) def test__row_from_mapping_wo_schema(self): - from google.cloud.bigquery.table import Table, _TABLE_HAS_NO_SCHEMA + from google.cloud.bigquery.table import _TABLE_HAS_NO_SCHEMA, Table MAPPING = {"full_name": "Phred Phlyntstone", "age": 32} dataset = DatasetReference(self.PROJECT, self.DS_ID) @@ -2581,8 +2571,7 @@ def _make_one_from_data(self, schema=(), rows=()): return self._make_one(_mock_client(), api_request, path, schema) def test_constructor(self): - from google.cloud.bigquery.table import _item_to_row - from google.cloud.bigquery.table import _rows_page_start + from google.cloud.bigquery.table import _item_to_row, _rows_page_start client = _mock_client() path = "/some/path" @@ -2880,7 +2869,8 @@ def test__should_use_bqstorage_returns_true_if_no_cached_results(self): def test__should_use_bqstorage_returns_false_if_page_token_set(self): iterator = self._make_one( - page_token="abc", first_page_response=None # not cached + page_token="abc", + first_page_response=None, # not cached ) result = iterator._should_use_bqstorage( bqstorage_client=None, create_bqstorage_client=True @@ -2889,7 +2879,8 @@ def test__should_use_bqstorage_returns_false_if_page_token_set(self): def test__should_use_bqstorage_returns_false_if_max_results_set(self): iterator = self._make_one( - max_results=10, first_page_response=None # not cached + max_results=10, + first_page_response=None, # not cached ) result = iterator._should_use_bqstorage( bqstorage_client=None, create_bqstorage_client=True @@ -3048,14 +3039,13 @@ def test_to_arrow_iterable_w_bqstorage(self): pyarrow = pytest.importorskip("pyarrow") pytest.importorskip("google.cloud.bigquery_storage") from google.cloud import bigquery_storage + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut from google.cloud.bigquery_storage_v1 import reader from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, ) - from google.cloud.bigquery import schema - from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport @@ -3422,9 +3412,9 @@ def test_to_arrow_w_bqstorage(self): pytest.importorskip("numpy") pyarrow = pytest.importorskip("pyarrow") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage from google.cloud.bigquery_storage_v1 import reader from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, @@ -3506,9 +3496,9 @@ def test_to_arrow_w_bqstorage_creates_client(self): pytest.importorskip("numpy") pytest.importorskip("pyarrow") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, ) @@ -3574,9 +3564,9 @@ def test_to_arrow_w_bqstorage_no_streams(self): pytest.importorskip("numpy") pyarrow = pytest.importorskip("pyarrow") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) session = bigquery_storage.types.ReadSession() @@ -3749,9 +3739,9 @@ def test_to_dataframe_iterable_w_bqstorage(self): pandas = pytest.importorskip("pandas") pyarrow = pytest.importorskip("pyarrow") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage from google.cloud.bigquery_storage_v1 import reader from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, @@ -3823,9 +3813,9 @@ def test_to_dataframe_iterable_w_bqstorage_max_results_warning(self): pytest.importorskip("numpy") pandas = pytest.importorskip("pandas") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) @@ -4151,6 +4141,8 @@ def test_to_dataframe_tqdm_error(self): # dependency and are unrelated to the code under test. if "Pyparsing" in warning.category.__name__: continue + if issubclass(warning.category, PendingDeprecationWarning): + continue self.assertIn( warning.category, [UserWarning, DeprecationWarning, tqdm.TqdmExperimentalWarning], @@ -4176,6 +4168,7 @@ def test_to_dataframe_w_empty_results(self): def test_to_dataframe_w_various_types_nullable(self): pandas = pytest.importorskip("pandas") import datetime + from google.cloud.bigquery.schema import SchemaField schema = [ @@ -4855,9 +4848,9 @@ def test_to_dataframe_w_bqstorage_creates_client(self): pytest.importorskip("numpy") pytest.importorskip("pandas") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, ) @@ -4889,9 +4882,9 @@ def test_to_dataframe_w_bqstorage_no_streams(self): pytest.importorskip("numpy") pytest.importorskip("pandas") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) session = bigquery_storage.types.ReadSession() @@ -4919,8 +4912,8 @@ def test_to_dataframe_w_bqstorage_logs_session(self): pytest.importorskip("google.cloud.bigquery_storage") pytest.importorskip("pandas") pytest.importorskip("pyarrow") - from google.cloud.bigquery.table import Table from google.cloud import bigquery_storage + from google.cloud.bigquery.table import Table bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) session = bigquery_storage.types.ReadSession() @@ -4999,9 +4992,9 @@ def test_to_dataframe_w_bqstorage_nonempty(self): pytest.importorskip("google.cloud.bigquery_storage") pytest.importorskip("pandas") pyarrow = pytest.importorskip("pyarrow") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage from google.cloud.bigquery_storage_v1 import reader from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( grpc as big_query_read_grpc_transport, @@ -5604,8 +5597,7 @@ def test_to_geodataframe_no_geog(self): with self.assertRaisesRegex( TypeError, re.escape( - "There must be at least one GEOGRAPHY column" - " to create a GeoDataFrame" + "There must be at least one GEOGRAPHY column to create a GeoDataFrame" ), ): row_iterator.to_geodataframe(create_bqstorage_client=False) @@ -5723,6 +5715,305 @@ def test_rowiterator_to_geodataframe_delegation(self, to_dataframe): self.assertEqual([v.__class__.__name__ for v in df.g], ["Point"]) + def test_to_dataframe_delegated_when_supported_no_range_types(self): + import sys + + db_dtypes = pytest.importorskip("db_dtypes") + pandas = pytest.importorskip("pandas") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + + mock_pandas_gbq.pandas.from_row_iterator.assert_called_once_with( + row_iterator, + bqstorage_client=None, + dtypes={}, + progress_bar_type="tqdm", + create_bqstorage_client=True, + geography_as_object=False, + bool_dtype=pandas.BooleanDtype(), + int_dtype=pandas.Int64Dtype(), + float_dtype=None, + string_dtype=None, + date_dtype=DefaultPandasDTypes.DATE_DTYPE, + datetime_dtype=None, + time_dtype=db_dtypes.TimeDtype(), + timestamp_dtype=None, + range_date_dtype=None, + range_datetime_dtype=None, + range_timestamp_dtype=None, + timeout=5.0, + ) + self.assertEqual(df, mock.sentinel.dataframe) + + def test_to_dataframe_delegated_when_supported_with_range_types(self): + import sys + + db_dtypes = pytest.importorskip("db_dtypes") + pandas = pytest.importorskip("pandas") + if not hasattr(pandas, "ArrowDtype"): + pytest.skip("pandas.ArrowDtype is not available in this environment.") + pyarrow = pytest.importorskip("pyarrow") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + True, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + + expected_range_date = pandas.ArrowDtype( + pyarrow.struct( + [("start", pyarrow.date32()), ("end", pyarrow.date32())] + ) + ) + expected_range_datetime = pandas.ArrowDtype( + pyarrow.struct( + [ + ("start", pyarrow.timestamp("us")), + ("end", pyarrow.timestamp("us")), + ] + ) + ) + expected_range_timestamp = pandas.ArrowDtype( + pyarrow.struct( + [ + ("start", pyarrow.timestamp("us", tz="UTC")), + ("end", pyarrow.timestamp("us", tz="UTC")), + ] + ) + ) + + mock_pandas_gbq.pandas.from_row_iterator.assert_called_once_with( + row_iterator, + bqstorage_client=None, + dtypes={}, + progress_bar_type="tqdm", + create_bqstorage_client=True, + geography_as_object=False, + bool_dtype=pandas.BooleanDtype(), + int_dtype=pandas.Int64Dtype(), + float_dtype=None, + string_dtype=None, + date_dtype=DefaultPandasDTypes.DATE_DTYPE, + datetime_dtype=None, + time_dtype=db_dtypes.TimeDtype(), + timestamp_dtype=None, + range_date_dtype=expected_range_date, + range_datetime_dtype=expected_range_datetime, + range_timestamp_dtype=expected_range_timestamp, + timeout=5.0, + ) + self.assertEqual(df, mock.sentinel.dataframe) + + def test_to_dataframe_not_delegated_when_unsupported(self): + import sys + + pandas = pytest.importorskip("pandas") + pytest.importorskip("pyarrow") + mock_pandas_gbq = mock.Mock() + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + + with warnings.catch_warnings(record=True) as warned: + warnings.simplefilter("always") + df = row_iterator.to_dataframe(create_bqstorage_client=False) + + mock_pandas_gbq.pandas.from_row_iterator.assert_not_called() + self.assertIsInstance(df, pandas.DataFrame) + self.assertEqual(df.name.tolist(), ["foo"]) + + deprecation_warnings = [ + w + for w in warned + if issubclass(w.category, PendingDeprecationWarning) + and "pandas-gbq" in str(w.message) + ] + self.assertEqual(len(deprecation_warnings), 1) + + def test_to_dataframe_delegated_updates_user_agent(self): + import sys + + pytest.importorskip("db_dtypes") + pytest.importorskip("pandas") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + mock_client_info = mock.Mock() + mock_client_info.user_agent = "gl-python/3.10.0" + + mock_client = _mock_client() + mock_client._connection = mock.Mock(_client_info=mock_client_info) + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + row_iterator.client = mock_client + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + self.assertEqual(df, mock.sentinel.dataframe) + self.assertEqual( + mock_client_info.user_agent, + "gl-python/3.10.0 pandas-gbq/1.0.0", + ) + + def test_to_dataframe_delegated_does_not_duplicate_user_agent(self): + import sys + + pytest.importorskip("db_dtypes") + pytest.importorskip("pandas") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + mock_client_info = mock.Mock() + mock_client_info.user_agent = "gl-python/3.10.0 pandas-gbq/1.0.0" + + mock_client = _mock_client() + mock_client._connection = mock.Mock(_client_info=mock_client_info) + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + row_iterator.client = mock_client + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + self.assertEqual(df, mock.sentinel.dataframe) + self.assertEqual( + mock_client_info.user_agent, + "gl-python/3.10.0 pandas-gbq/1.0.0", + ) + + def test_to_dataframe_delegated_when_client_info_is_none(self): + import sys + + pytest.importorskip("db_dtypes") + pytest.importorskip("pandas") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + mock_client = _mock_client() + mock_client._connection = mock.Mock(_client_info=None) + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + row_iterator.client = mock_client + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + self.assertEqual(df, mock.sentinel.dataframe) + + def test_to_dataframe_delegated_when_user_agent_is_none(self): + import sys + + pytest.importorskip("db_dtypes") + pytest.importorskip("pandas") + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.pandas.from_row_iterator.return_value = mock.sentinel.dataframe + mock_pandas_gbq.__version__ = "1.0.0" + + mock_client_info = mock.Mock() + mock_client_info.user_agent = None + + mock_client = _mock_client() + mock_client._connection = mock.Mock(_client_info=mock_client_info) + + with mock.patch( + "google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported", + new_callable=mock.PropertyMock, + return_value=True, + ): + with mock.patch( + "google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW", + False, + ): + with mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}): + row_iterator = self._make_one_from_data( + (("name", "STRING"),), (("foo",),) + ) + row_iterator.client = mock_client + df = row_iterator.to_dataframe( + progress_bar_type="tqdm", timeout=5.0 + ) + self.assertEqual(df, mock.sentinel.dataframe) + self.assertEqual( + mock_client_info.user_agent, + "pandas-gbq/1.0.0", + ) + class TestPartitionRange(unittest.TestCase): def _get_target_class(self): @@ -6329,10 +6620,10 @@ def test_constructor_defaults(self): def test_constructor_explicit(self): from google.cloud.bigquery.table import ( - PrimaryKey, + ColumnReference, ForeignKey, + PrimaryKey, TableReference, - ColumnReference, ) primary_key = PrimaryKey(columns=["my_pk_id"]) @@ -6364,10 +6655,10 @@ def test_constructor_explicit_with_none(self): def test__eq__other_type(self): from google.cloud.bigquery.table import ( - PrimaryKey, + ColumnReference, ForeignKey, + PrimaryKey, TableReference, - ColumnReference, ) table_constraint = self._make_one( @@ -6602,8 +6893,8 @@ def test_table_constraint_eq_parametrized( ColumnReference, ForeignKey, PrimaryKey, - TableReference, TableConstraints, + TableReference, ) # Helper function to create a PrimaryKey object or None @@ -6850,9 +7141,9 @@ def test_table_reference_to_bqstorage_v1_stable(table_path): def test_to_arrow_iterable_w_bqstorage_max_stream_count(preserve_order): pytest.importorskip("pandas") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) session = bigquery_storage.types.ReadSession() @@ -6887,9 +7178,9 @@ def test_to_arrow_iterable_w_bqstorage_max_stream_count(preserve_order): def test_to_dataframe_iterable_w_bqstorage_max_stream_count(preserve_order): pytest.importorskip("pandas") pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud import bigquery_storage from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - from google.cloud import bigquery_storage bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) session = bigquery_storage.types.ReadSession()