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/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py index fa03156e2e26..20fab496506b 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py @@ -34,89 +34,99 @@ __version__ = bigquery_version.__version__ -from google.cloud.bigquery.client import Client -from google.cloud.bigquery.dataset import AccessEntry -from google.cloud.bigquery.dataset import Dataset -from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery import enums -from google.cloud.bigquery.enums import AutoRowIDs -from google.cloud.bigquery.enums import DecimalTargetType -from google.cloud.bigquery.enums import KeyResultStatementKind -from google.cloud.bigquery.enums import SqlTypeNames -from google.cloud.bigquery.enums import StandardSqlTypeNames -from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError -from google.cloud.bigquery.exceptions import LegacyPandasError -from google.cloud.bigquery.exceptions import LegacyPyarrowError -from google.cloud.bigquery.external_config import ExternalConfig -from google.cloud.bigquery.external_config import BigtableOptions -from google.cloud.bigquery.external_config import BigtableColumnFamily -from google.cloud.bigquery.external_config import BigtableColumn -from google.cloud.bigquery.external_config import CSVOptions -from google.cloud.bigquery.external_config import GoogleSheetsOptions -from google.cloud.bigquery.external_config import ExternalSourceFormat -from google.cloud.bigquery.external_config import HivePartitioningOptions -from google.cloud.bigquery.format_options import AvroOptions -from google.cloud.bigquery.format_options import ParquetOptions +from google.cloud.bigquery.client import Client +from google.cloud.bigquery.dataset import AccessEntry, Dataset, DatasetReference +from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration +from google.cloud.bigquery.enums import ( + AutoRowIDs, + DecimalTargetType, + KeyResultStatementKind, + SqlTypeNames, + StandardSqlTypeNames, +) +from google.cloud.bigquery.exceptions import ( + LegacyBigQueryStorageError, + LegacyPandasError, + LegacyPyarrowError, +) +from google.cloud.bigquery.external_config import ( + BigtableColumn, + BigtableColumnFamily, + BigtableOptions, + CSVOptions, + ExternalConfig, + ExternalSourceFormat, + GoogleSheetsOptions, + HivePartitioningOptions, +) +from google.cloud.bigquery.format_options import AvroOptions, ParquetOptions +from google.cloud.bigquery.job import ( + Compression, + CopyJob, + CopyJobConfig, + CreateDisposition, + DestinationFormat, + DmlStats, + Encoding, + ExtractJob, + ExtractJobConfig, + LoadJob, + LoadJobConfig, + OperationType, + QueryJob, + QueryJobConfig, + QueryPriority, + SchemaUpdateOption, + ScriptOptions, + SourceFormat, + TransactionInfo, + UnknownJob, + WriteDisposition, +) from google.cloud.bigquery.job.base import SessionInfo -from google.cloud.bigquery.job import Compression -from google.cloud.bigquery.job import CopyJob -from google.cloud.bigquery.job import CopyJobConfig -from google.cloud.bigquery.job import CreateDisposition -from google.cloud.bigquery.job import DestinationFormat -from google.cloud.bigquery.job import DmlStats -from google.cloud.bigquery.job import Encoding -from google.cloud.bigquery.job import ExtractJob -from google.cloud.bigquery.job import ExtractJobConfig -from google.cloud.bigquery.job import LoadJob -from google.cloud.bigquery.job import LoadJobConfig -from google.cloud.bigquery.job import OperationType -from google.cloud.bigquery.job import QueryJob -from google.cloud.bigquery.job import QueryJobConfig -from google.cloud.bigquery.job import QueryPriority -from google.cloud.bigquery.job import SchemaUpdateOption -from google.cloud.bigquery.job import ScriptOptions -from google.cloud.bigquery.job import SourceFormat -from google.cloud.bigquery.job import UnknownJob -from google.cloud.bigquery.job import TransactionInfo -from google.cloud.bigquery.job import WriteDisposition -from google.cloud.bigquery.model import Model -from google.cloud.bigquery.model import ModelReference -from google.cloud.bigquery.query import ArrayQueryParameter -from google.cloud.bigquery.query import ArrayQueryParameterType -from google.cloud.bigquery.query import ConnectionProperty -from google.cloud.bigquery.query import ScalarQueryParameter -from google.cloud.bigquery.query import ScalarQueryParameterType -from google.cloud.bigquery.query import RangeQueryParameter -from google.cloud.bigquery.query import RangeQueryParameterType -from google.cloud.bigquery.query import SqlParameterScalarTypes -from google.cloud.bigquery.query import StructQueryParameter -from google.cloud.bigquery.query import StructQueryParameterType -from google.cloud.bigquery.query import UDFResource +from google.cloud.bigquery.model import Model, ModelReference +from google.cloud.bigquery.query import ( + ArrayQueryParameter, + ArrayQueryParameterType, + ConnectionProperty, + RangeQueryParameter, + RangeQueryParameterType, + ScalarQueryParameter, + ScalarQueryParameterType, + SqlParameterScalarTypes, + StructQueryParameter, + StructQueryParameterType, + UDFResource, +) from google.cloud.bigquery.retry import DEFAULT_RETRY -from google.cloud.bigquery.routine import DeterminismLevel -from google.cloud.bigquery.routine import Routine -from google.cloud.bigquery.routine import RoutineArgument -from google.cloud.bigquery.routine import RoutineReference -from google.cloud.bigquery.routine import RoutineType -from google.cloud.bigquery.routine import RemoteFunctionOptions -from google.cloud.bigquery.routine import ExternalRuntimeOptions -from google.cloud.bigquery.schema import PolicyTagList -from google.cloud.bigquery.schema import SchemaField -from google.cloud.bigquery.schema import FieldElementType -from google.cloud.bigquery.standard_sql import StandardSqlDataType -from google.cloud.bigquery.standard_sql import StandardSqlField -from google.cloud.bigquery.standard_sql import StandardSqlStructType -from google.cloud.bigquery.standard_sql import StandardSqlTableType -from google.cloud.bigquery.table import PartitionRange -from google.cloud.bigquery.table import RangePartitioning -from google.cloud.bigquery.table import Row -from google.cloud.bigquery.table import SnapshotDefinition -from google.cloud.bigquery.table import CloneDefinition -from google.cloud.bigquery.table import Table -from google.cloud.bigquery.table import TableReference -from google.cloud.bigquery.table import TimePartitioningType -from google.cloud.bigquery.table import TimePartitioning -from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration +from google.cloud.bigquery.routine import ( + DeterminismLevel, + ExternalRuntimeOptions, + RemoteFunctionOptions, + Routine, + RoutineArgument, + RoutineReference, + RoutineType, +) +from google.cloud.bigquery.schema import FieldElementType, PolicyTagList, SchemaField +from google.cloud.bigquery.standard_sql import ( + StandardSqlDataType, + StandardSqlField, + StandardSqlStructType, + StandardSqlTableType, +) +from google.cloud.bigquery.table import ( + CloneDefinition, + PartitionRange, + RangePartitioning, + Row, + SnapshotDefinition, + Table, + TableReference, + TimePartitioning, + TimePartitioningType, +) try: import bigquery_magics # type: ignore diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_helpers.py index a35fe1677a8f..b3760ec52df5 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_helpers.py @@ -19,23 +19,24 @@ import decimal import json import math -import re import os +import re import textwrap import warnings from typing import Any, Optional, Tuple, Type, Union from dateutil import relativedelta -from google.cloud._helpers import UTC # type: ignore -from google.cloud._helpers import _date_from_iso8601_date -from google.cloud._helpers import _datetime_from_microseconds -from google.cloud._helpers import _RFC3339_MICROS -from google.cloud._helpers import _RFC3339_NO_FRACTION -from google.cloud._helpers import _to_bytes -from google.cloud.bigquery import enums - -from google.auth import credentials as ga_credentials # type: ignore from google.api_core import client_options as client_options_lib +from google.auth import credentials as ga_credentials # type: ignore +from google.cloud._helpers import ( + _RFC3339_MICROS, + _RFC3339_NO_FRACTION, + UTC, # type: ignore + _date_from_iso8601_date, + _datetime_from_microseconds, + _to_bytes, +) +from google.cloud.bigquery import enums TimeoutType = Union[float, None] diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index 30f89759ee2b..3cd46503c40d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -41,20 +41,17 @@ import dataclasses import datetime import functools -import uuid import textwrap -from typing import Any, Callable, Dict, Optional, TYPE_CHECKING, Union +import uuid import warnings +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union import google.api_core.exceptions as core_exceptions -from google.api_core import retry as retries - -from google.cloud.bigquery import enums -from google.cloud.bigquery import job import google.cloud.bigquery.job.query import google.cloud.bigquery.query -from google.cloud.bigquery import table import google.cloud.bigquery.retry +from google.api_core import retry as retries +from google.cloud.bigquery import enums, job, table from google.cloud.bigquery.retry import POLLING_DEFAULT_VALUE # Avoid circular imports diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_pandas_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_pandas_helpers.py index 7bd9f99b625f..812319f5b32c 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_pandas_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_pandas_helpers.py @@ -20,22 +20,18 @@ """ import concurrent.futures -from datetime import datetime import functools -from itertools import islice import logging import queue import threading import time import warnings -from typing import Any, Union, Optional, Callable, Generator, List - +from datetime import datetime +from itertools import islice +from typing import Any, Callable, Generator, List, Optional, Union -from google.cloud.bigquery import _pyarrow_helpers -from google.cloud.bigquery import _versions_helpers +from google.cloud.bigquery import _pyarrow_helpers, _versions_helpers, schema from google.cloud.bigquery import retry as bq_retry -from google.cloud.bigquery import schema - try: import pandas # type: ignore diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_string_references.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_string_references.py index 16a145a05a96..29147de4c408 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_string_references.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_string_references.py @@ -21,7 +21,6 @@ import re from typing import TypedDict, Union - ParsedDatasetReference = TypedDict( "ParsedDatasetReference", { 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/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index ce8768b68b30..5064ee85a87e 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -54,8 +54,8 @@ from google.api_core.iam import Policy from google.cloud import exceptions # pytype: disable=import-error from google.cloud.client import ( - ClientWithProject, # type: ignore # pytype: disable=import-error -) + ClientWithProject, +) # type: ignore # pytype: disable=import-error from google.resumable_media.requests import ( MultipartUpload, # type: ignore ResumableUpload, diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dataset.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dataset.py index fa3342de1628..a4e176430de8 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dataset.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dataset.py @@ -18,19 +18,15 @@ import copy import json - import typing -from typing import Optional, List, Dict, Any, Union +from typing import Any, Dict, List, Optional, Union import google.cloud._helpers # type: ignore - -from google.cloud.bigquery import _helpers +from google.cloud.bigquery import _helpers, _string_references, external_config +from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration from google.cloud.bigquery.model import ModelReference from google.cloud.bigquery.routine import Routine, RoutineReference from google.cloud.bigquery.table import Table, TableReference -from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration -from google.cloud.bigquery import external_config -from google.cloud.bigquery import _string_references def _get_table_reference(self, table_id: str) -> TableReference: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/__init__.py index d1a723949b10..cdb0b4b24f11 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/__init__.py @@ -21,32 +21,34 @@ https://www.python.org/dev/peps/pep-0249/ """ -from google.cloud.bigquery.dbapi.connection import connect -from google.cloud.bigquery.dbapi.connection import Connection +from google.cloud.bigquery.dbapi.connection import Connection, connect from google.cloud.bigquery.dbapi.cursor import Cursor -from google.cloud.bigquery.dbapi.exceptions import Warning -from google.cloud.bigquery.dbapi.exceptions import Error -from google.cloud.bigquery.dbapi.exceptions import InterfaceError -from google.cloud.bigquery.dbapi.exceptions import DatabaseError -from google.cloud.bigquery.dbapi.exceptions import DataError -from google.cloud.bigquery.dbapi.exceptions import OperationalError -from google.cloud.bigquery.dbapi.exceptions import IntegrityError -from google.cloud.bigquery.dbapi.exceptions import InternalError -from google.cloud.bigquery.dbapi.exceptions import ProgrammingError -from google.cloud.bigquery.dbapi.exceptions import NotSupportedError -from google.cloud.bigquery.dbapi.types import Binary -from google.cloud.bigquery.dbapi.types import Date -from google.cloud.bigquery.dbapi.types import DateFromTicks -from google.cloud.bigquery.dbapi.types import Time -from google.cloud.bigquery.dbapi.types import TimeFromTicks -from google.cloud.bigquery.dbapi.types import Timestamp -from google.cloud.bigquery.dbapi.types import TimestampFromTicks -from google.cloud.bigquery.dbapi.types import BINARY -from google.cloud.bigquery.dbapi.types import DATETIME -from google.cloud.bigquery.dbapi.types import NUMBER -from google.cloud.bigquery.dbapi.types import ROWID -from google.cloud.bigquery.dbapi.types import STRING - +from google.cloud.bigquery.dbapi.exceptions import ( + DatabaseError, + DataError, + Error, + IntegrityError, + InterfaceError, + InternalError, + NotSupportedError, + OperationalError, + ProgrammingError, + Warning, +) +from google.cloud.bigquery.dbapi.types import ( + BINARY, + DATETIME, + NUMBER, + ROWID, + STRING, + Binary, + Date, + DateFromTicks, + Time, + TimeFromTicks, + Timestamp, + TimestampFromTicks, +) apilevel = "2.0" diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/_helpers.py index a4ab05ce82a3..e3c955bc2268 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/_helpers.py @@ -13,19 +13,18 @@ # limitations under the License. -from collections import abc as collections_abc import datetime import decimal import functools import numbers import re import typing +from collections import abc as collections_abc from google.cloud import bigquery -from google.cloud.bigquery import table, query +from google.cloud.bigquery import query, table from google.cloud.bigquery.dbapi import exceptions - _NUMERIC_SERVER_MIN = decimal.Decimal("-9.9999999999999999999999999999999999999E+28") _NUMERIC_SERVER_MAX = decimal.Decimal("9.9999999999999999999999999999999999999E+28") diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py index b0d7ef895141..8c031c5319d0 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py @@ -17,8 +17,7 @@ import weakref from google.cloud import bigquery -from google.cloud.bigquery.dbapi import cursor -from google.cloud.bigquery.dbapi import _helpers +from google.cloud.bigquery.dbapi import _helpers, cursor @_helpers.raise_on_closed("Operating on a closed connection.") diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/cursor.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/cursor.py index 4beec6929a05..9293083d9316 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/cursor.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/cursor.py @@ -17,8 +17,8 @@ from __future__ import annotations import collections -from collections import abc as collections_abc import re +from collections import abc as collections_abc from typing import Optional try: @@ -29,11 +29,9 @@ # Having BQ Storage available implies that pyarrow >=1.0.0 is available, too. _ARROW_COMPRESSION_SUPPORT = True -from google.cloud.bigquery import job -from google.cloud.bigquery.dbapi import _helpers -from google.cloud.bigquery.dbapi import exceptions import google.cloud.exceptions # type: ignore - +from google.cloud.bigquery import job +from google.cloud.bigquery.dbapi import _helpers, exceptions # Per PEP 249: A 7-item sequence containing information describing one result # column. The first two items (name and type_code) are mandatory, the other diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/types.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/types.py index 717593ae1336..89ca07ae70ab 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/types.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/types.py @@ -22,7 +22,6 @@ import datetime - Date = datetime.date Time = datetime.time Timestamp = datetime.datetime diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/external_config.py b/packages/google-cloud-bigquery/google/cloud/bigquery/external_config.py index 7e76f93b57fe..bee2d6d27723 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/external_config.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/external_config.py @@ -25,14 +25,15 @@ import typing from typing import Any, Dict, FrozenSet, Iterable, Optional, Union -from google.cloud.bigquery._helpers import _to_bytes -from google.cloud.bigquery._helpers import _bytes_to_json -from google.cloud.bigquery._helpers import _int_or_none -from google.cloud.bigquery._helpers import _str_or_none -from google.cloud.bigquery import _helpers +from google.cloud.bigquery import _helpers, schema +from google.cloud.bigquery._helpers import ( + _bytes_to_json, + _int_or_none, + _str_or_none, + _to_bytes, +) from google.cloud.bigquery.enums import SourceColumnMatch from google.cloud.bigquery.format_options import AvroOptions, ParquetOptions -from google.cloud.bigquery import schema from google.cloud.bigquery.schema import SchemaField diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/__init__.py index 4cda6596516c..5eacdd7e337b 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/__init__.py @@ -14,41 +14,42 @@ """Define API Jobs.""" -from google.cloud.bigquery.job.base import _AsyncJob -from google.cloud.bigquery.job.base import _error_result_to_exception -from google.cloud.bigquery.job.base import _DONE_STATE -from google.cloud.bigquery.job.base import _JobConfig -from google.cloud.bigquery.job.base import _JobReference -from google.cloud.bigquery.job.base import ReservationUsage -from google.cloud.bigquery.job.base import ScriptStatistics -from google.cloud.bigquery.job.base import ScriptStackFrame -from google.cloud.bigquery.job.base import TransactionInfo -from google.cloud.bigquery.job.base import UnknownJob -from google.cloud.bigquery.job.copy_ import CopyJob -from google.cloud.bigquery.job.copy_ import CopyJobConfig -from google.cloud.bigquery.job.copy_ import OperationType -from google.cloud.bigquery.job.extract import ExtractJob -from google.cloud.bigquery.job.extract import ExtractJobConfig -from google.cloud.bigquery.job.load import LoadJob -from google.cloud.bigquery.job.load import LoadJobConfig -from google.cloud.bigquery.job.query import _contains_order_by -from google.cloud.bigquery.job.query import DmlStats -from google.cloud.bigquery.job.query import QueryJob -from google.cloud.bigquery.job.query import QueryJobConfig -from google.cloud.bigquery.job.query import QueryPlanEntry -from google.cloud.bigquery.job.query import QueryPlanEntryStep -from google.cloud.bigquery.job.query import ScriptOptions -from google.cloud.bigquery.job.query import TimelineEntry -from google.cloud.bigquery.job.query import IncrementalResultStats -from google.cloud.bigquery.enums import Compression -from google.cloud.bigquery.enums import CreateDisposition -from google.cloud.bigquery.enums import DestinationFormat -from google.cloud.bigquery.enums import Encoding -from google.cloud.bigquery.enums import QueryPriority -from google.cloud.bigquery.enums import SchemaUpdateOption -from google.cloud.bigquery.enums import SourceFormat -from google.cloud.bigquery.enums import WriteDisposition - +from google.cloud.bigquery.enums import ( + Compression, + CreateDisposition, + DestinationFormat, + Encoding, + QueryPriority, + SchemaUpdateOption, + SourceFormat, + WriteDisposition, +) +from google.cloud.bigquery.job.base import ( + _DONE_STATE, + ReservationUsage, + ScriptStackFrame, + ScriptStatistics, + TransactionInfo, + UnknownJob, + _AsyncJob, + _error_result_to_exception, + _JobConfig, + _JobReference, +) +from google.cloud.bigquery.job.copy_ import CopyJob, CopyJobConfig, OperationType +from google.cloud.bigquery.job.extract import ExtractJob, ExtractJobConfig +from google.cloud.bigquery.job.load import LoadJob, LoadJobConfig +from google.cloud.bigquery.job.query import ( + DmlStats, + IncrementalResultStats, + QueryJob, + QueryJobConfig, + QueryPlanEntry, + QueryPlanEntryStep, + ScriptOptions, + TimelineEntry, + _contains_order_by, +) # Include classes previously in job.py for backwards compatibility. __all__ = [ diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/base.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/base.py index 7576fc9aa64e..0feea493cc92 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/base.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/base.py @@ -14,24 +14,19 @@ """Base classes and helpers for job classes.""" -from collections import namedtuple import copy import http import threading import typing +from collections import namedtuple from typing import ClassVar, Dict, Optional, Sequence -from google.api_core import retry as retries -from google.api_core import exceptions import google.api_core.future.polling - +from google.api_core import exceptions +from google.api_core import retry as retries from google.cloud.bigquery import _helpers from google.cloud.bigquery._helpers import _int_or_none -from google.cloud.bigquery.retry import ( - DEFAULT_GET_JOB_TIMEOUT, - DEFAULT_RETRY, -) - +from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT, DEFAULT_RETRY _DONE_STATE = "DONE" _STOPPED_REASON = "stopped" diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/copy_.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/copy_.py index 5c52aeed6470..79ed801e971e 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/copy_.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/copy_.py @@ -17,14 +17,11 @@ import typing from typing import Optional -from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration from google.cloud.bigquery import _helpers +from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration +from google.cloud.bigquery.job.base import _AsyncJob, _JobConfig, _JobReference from google.cloud.bigquery.table import TableReference -from google.cloud.bigquery.job.base import _AsyncJob -from google.cloud.bigquery.job.base import _JobConfig -from google.cloud.bigquery.job.base import _JobReference - class OperationType: """Different operation types supported in table copy job. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/extract.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/extract.py index 64ec39b7603c..08568b821444 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/extract.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/extract.py @@ -17,13 +17,9 @@ import typing from google.cloud.bigquery import _helpers +from google.cloud.bigquery.job.base import _AsyncJob, _JobConfig, _JobReference from google.cloud.bigquery.model import ModelReference -from google.cloud.bigquery.table import Table -from google.cloud.bigquery.table import TableListItem -from google.cloud.bigquery.table import TableReference -from google.cloud.bigquery.job.base import _AsyncJob -from google.cloud.bigquery.job.base import _JobConfig -from google.cloud.bigquery.job.base import _JobReference +from google.cloud.bigquery.table import Table, TableListItem, TableReference class ExtractJobConfig(_JobConfig): diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/load.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/load.py index 9c74f7124393..f3a20387e4c5 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/load.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/load.py @@ -15,22 +15,21 @@ """Classes for load jobs.""" import typing -from typing import FrozenSet, List, Iterable, Optional, Union +from typing import FrozenSet, Iterable, List, Optional, Union +from google.cloud.bigquery import _helpers from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration from google.cloud.bigquery.enums import SourceColumnMatch from google.cloud.bigquery.external_config import HivePartitioningOptions from google.cloud.bigquery.format_options import ParquetOptions -from google.cloud.bigquery import _helpers -from google.cloud.bigquery.schema import SchemaField -from google.cloud.bigquery.schema import _to_schema_fields -from google.cloud.bigquery.table import RangePartitioning -from google.cloud.bigquery.table import TableReference -from google.cloud.bigquery.table import TimePartitioning -from google.cloud.bigquery.job.base import _AsyncJob -from google.cloud.bigquery.job.base import _JobConfig -from google.cloud.bigquery.job.base import _JobReference +from google.cloud.bigquery.job.base import _AsyncJob, _JobConfig, _JobReference from google.cloud.bigquery.query import ConnectionProperty +from google.cloud.bigquery.schema import SchemaField, _to_schema_fields +from google.cloud.bigquery.table import ( + RangePartitioning, + TableReference, + TimePartitioning, +) class ColumnNameCharacterMap: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py index e85e43f6a164..2307f25ecd64 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py @@ -21,42 +21,39 @@ import typing from typing import Any, Dict, Iterable, List, Optional, Union +import requests from google.api_core import exceptions from google.api_core import retry as retries -import requests - -from google.cloud.bigquery.dataset import Dataset -from google.cloud.bigquery.dataset import DatasetListItem -from google.cloud.bigquery.dataset import DatasetReference +from google.cloud.bigquery import _helpers +from google.cloud.bigquery._tqdm_helpers import wait_for_query +from google.cloud.bigquery.dataset import Dataset, DatasetListItem, DatasetReference from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration -from google.cloud.bigquery.enums import KeyResultStatementKind, DefaultPandasDTypes +from google.cloud.bigquery.enums import DefaultPandasDTypes, KeyResultStatementKind from google.cloud.bigquery.external_config import ExternalConfig -from google.cloud.bigquery import _helpers +from google.cloud.bigquery.job.base import _AsyncJob, _JobConfig, _JobReference from google.cloud.bigquery.query import ( - _query_param_from_api_repr, ArrayQueryParameter, ConnectionProperty, ScalarQueryParameter, StructQueryParameter, UDFResource, + _query_param_from_api_repr, ) from google.cloud.bigquery.retry import ( - DEFAULT_RETRY, DEFAULT_JOB_RETRY, + DEFAULT_RETRY, POLLING_DEFAULT_VALUE, ) from google.cloud.bigquery.routine import RoutineReference from google.cloud.bigquery.schema import SchemaField -from google.cloud.bigquery.table import _EmptyRowIterator -from google.cloud.bigquery.table import RangePartitioning -from google.cloud.bigquery.table import _table_arg_to_table_ref -from google.cloud.bigquery.table import TableReference, PropertyGraphReference -from google.cloud.bigquery.table import TimePartitioning -from google.cloud.bigquery._tqdm_helpers import wait_for_query - -from google.cloud.bigquery.job.base import _AsyncJob -from google.cloud.bigquery.job.base import _JobConfig -from google.cloud.bigquery.job.base import _JobReference +from google.cloud.bigquery.table import ( + PropertyGraphReference, + RangePartitioning, + TableReference, + TimePartitioning, + _EmptyRowIterator, + _table_arg_to_table_ref, +) try: import pandas # type: ignore @@ -66,8 +63,8 @@ if typing.TYPE_CHECKING: # pragma: NO COVER # Assumption: type checks are only used by library developers and CI environments # that have all optional dependencies installed, thus no conditional imports. - import pandas # type: ignore import geopandas # type: ignore + import pandas # type: ignore import pyarrow # type: ignore from google.cloud import bigquery_storage from google.cloud.bigquery.client import Client diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/__init__.py index d228a35bb134..704bb2a580b8 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/__init__.py @@ -14,7 +14,6 @@ from google.cloud.bigquery.magics.magics import context - # For backwards compatibility we need to make the context available in the path # google.cloud.bigquery.magics.context __all__ = ("context",) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/__init__.py index 9471446c516b..bcaab2c7c579 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/__init__.py @@ -12,17 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.bigquery.magics.line_arg_parser.exceptions import ParseError from google.cloud.bigquery.magics.line_arg_parser.exceptions import ( DuplicateQueryParamsError, + ParseError, QueryParamsParseError, ) -from google.cloud.bigquery.magics.line_arg_parser.lexer import Lexer -from google.cloud.bigquery.magics.line_arg_parser.lexer import TokenType +from google.cloud.bigquery.magics.line_arg_parser.lexer import Lexer, TokenType from google.cloud.bigquery.magics.line_arg_parser.parser import Parser from google.cloud.bigquery.magics.line_arg_parser.visitors import QueryParamsExtractor - __all__ = ( "DuplicateQueryParamsError", "Lexer", diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/lexer.py b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/lexer.py index 71b287d01e83..180dd8b9528f 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/lexer.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/lexer.py @@ -12,13 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections import namedtuple -from collections import OrderedDict +import enum import itertools import re - -import enum - +from collections import OrderedDict, namedtuple Token = namedtuple("Token", ("type_", "lexeme", "pos")) StateTransition = namedtuple("StateTransition", ("new_state", "total_offset")) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/parser.py b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/parser.py index b9da20cd7862..18c0c0180b7f 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/parser.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/magics/line_arg_parser/parser.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.bigquery.magics.line_arg_parser import DuplicateQueryParamsError -from google.cloud.bigquery.magics.line_arg_parser import ParseError -from google.cloud.bigquery.magics.line_arg_parser import QueryParamsParseError -from google.cloud.bigquery.magics.line_arg_parser import TokenType +from google.cloud.bigquery.magics.line_arg_parser import ( + DuplicateQueryParamsError, + ParseError, + QueryParamsParseError, + TokenType, +) class ParseNode(object): 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/model.py b/packages/google-cloud-bigquery/google/cloud/bigquery/model.py index 16581be5a492..f58de58582ab 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/model.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/model.py @@ -24,8 +24,7 @@ from typing import Any, Dict, Optional, Sequence, Union import google.cloud._helpers # type: ignore -from google.cloud.bigquery import _helpers -from google.cloud.bigquery import standard_sql +from google.cloud.bigquery import _helpers, standard_sql from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/opentelemetry_tracing.py b/packages/google-cloud-bigquery/google/cloud/bigquery/opentelemetry_tracing.py index b5f6bf9912ed..3cfc3c7a2fa8 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/opentelemetry_tracing.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/opentelemetry_tracing.py @@ -14,12 +14,15 @@ import logging from contextlib import contextmanager + from google.api_core.exceptions import GoogleAPICallError # type: ignore logger = logging.getLogger(__name__) try: from opentelemetry import trace # type: ignore - from opentelemetry.instrumentation.utils import http_status_to_status_code # type: ignore + from opentelemetry.instrumentation.utils import ( + http_status_to_status_code, + ) # type: ignore from opentelemetry.trace.status import Status # type: ignore HAS_OPENTELEMETRY = True diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/query.py index 170ed2976849..616a568e005a 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/query.py @@ -14,18 +14,19 @@ """BigQuery query processing.""" -from collections import OrderedDict import copy import datetime import decimal -from typing import Any, cast, Optional, Dict, Union +from collections import OrderedDict +from typing import Any, Dict, Optional, Union, cast -from google.cloud.bigquery.table import _parse_schema_resource from google.cloud.bigquery import _helpers -from google.cloud.bigquery._helpers import _rows_from_json -from google.cloud.bigquery._helpers import _SCALAR_VALUE_TO_JSON_PARAM -from google.cloud.bigquery._helpers import _SUPPORTED_RANGE_ELEMENTS - +from google.cloud.bigquery._helpers import ( + _SCALAR_VALUE_TO_JSON_PARAM, + _SUPPORTED_RANGE_ELEMENTS, + _rows_from_json, +) +from google.cloud.bigquery.table import _parse_schema_resource _SCALAR_VALUE_TYPE = Optional[ Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date] diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/retry.py b/packages/google-cloud-bigquery/google/cloud/bigquery/retry.py index 4e78e7d28dcb..773779b6a4f9 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/retry.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/retry.py @@ -14,11 +14,10 @@ import logging -from google.api_core import exceptions -from google.api_core import retry import google.api_core.future.polling -from google.auth import exceptions as auth_exceptions # type: ignore import requests.exceptions +from google.api_core import exceptions, retry +from google.auth import exceptions as auth_exceptions # type: ignore _LOGGER = logging.getLogger(__name__) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/routine/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery/routine/__init__.py index 0251039579ca..b09f4b5ce1fd 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/routine/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/routine/__init__.py @@ -16,13 +16,14 @@ from google.cloud.bigquery.enums import DeterminismLevel -from google.cloud.bigquery.routine.routine import Routine -from google.cloud.bigquery.routine.routine import RoutineArgument -from google.cloud.bigquery.routine.routine import RoutineReference -from google.cloud.bigquery.routine.routine import RoutineType -from google.cloud.bigquery.routine.routine import RemoteFunctionOptions -from google.cloud.bigquery.routine.routine import ExternalRuntimeOptions - +from google.cloud.bigquery.routine.routine import ( + ExternalRuntimeOptions, + RemoteFunctionOptions, + Routine, + RoutineArgument, + RoutineReference, + RoutineType, +) __all__ = ( "DeterminismLevel", diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/routine/routine.py b/packages/google-cloud-bigquery/google/cloud/bigquery/routine/routine.py index c5aa8750ea70..35786aadc752 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/routine/routine.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/routine/routine.py @@ -20,8 +20,7 @@ import google.cloud._helpers # type: ignore from google.cloud.bigquery import _helpers -from google.cloud.bigquery.standard_sql import StandardSqlDataType -from google.cloud.bigquery.standard_sql import StandardSqlTableType +from google.cloud.bigquery.standard_sql import StandardSqlDataType, StandardSqlTableType class RoutineType: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/schema.py b/packages/google-cloud-bigquery/google/cloud/bigquery/schema.py index 1809df21f0b9..87b29c0d9896 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/schema.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/schema.py @@ -15,16 +15,14 @@ """Schemas for BigQuery tables / queries.""" from __future__ import annotations + import enum import typing -from typing import Any, cast, Dict, Iterable, Optional, Union, Sequence +from typing import Any, Dict, Iterable, Optional, Sequence, Union, cast -from google.cloud.bigquery import _helpers -from google.cloud.bigquery import standard_sql -from google.cloud.bigquery import enums +from google.cloud.bigquery import _helpers, enums, standard_sql from google.cloud.bigquery.enums import StandardSqlTypeNames - _STRUCT_TYPES = ("RECORD", "STRUCT") # SQL types reference: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index 870cdcc5d2ab..c3c97b04c395 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 @@ -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 @@ -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/google/cloud/bigquery_v2/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/__init__.py index 55486a39a2de..bb8ac07a7d88 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/__init__.py @@ -17,20 +17,23 @@ import warnings from .types.encryption_config import EncryptionConfiguration -from .types.model import DeleteModelRequest -from .types.model import GetModelRequest -from .types.model import ListModelsRequest -from .types.model import ListModelsResponse -from .types.model import Model -from .types.model import PatchModelRequest +from .types.model import ( + DeleteModelRequest, + GetModelRequest, + ListModelsRequest, + ListModelsResponse, + Model, + PatchModelRequest, +) from .types.model_reference import ModelReference -from .types.standard_sql import StandardSqlDataType -from .types.standard_sql import StandardSqlField -from .types.standard_sql import StandardSqlStructType -from .types.standard_sql import StandardSqlTableType +from .types.standard_sql import ( + StandardSqlDataType, + StandardSqlField, + StandardSqlStructType, + StandardSqlTableType, +) from .types.table_reference import TableReference - _LEGACY_MSG = ( "Legacy proto-based types from bigquery_v2 are not maintained anymore, " "use types defined in google.cloud.bigquery instead." diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/__init__.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/__init__.py index c36b30969b08..c038bcd74899 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/__init__.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/__init__.py @@ -13,9 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from .encryption_config import ( - EncryptionConfiguration, -) +from .encryption_config import EncryptionConfiguration from .model import ( DeleteModelRequest, GetModelRequest, @@ -24,18 +22,14 @@ Model, PatchModelRequest, ) -from .model_reference import ( - ModelReference, -) +from .model_reference import ModelReference from .standard_sql import ( StandardSqlDataType, StandardSqlField, StandardSqlStructType, StandardSqlTableType, ) -from .table_reference import ( - TableReference, -) +from .table_reference import TableReference __all__ = ( "EncryptionConfiguration", diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/encryption_config.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/encryption_config.py index 9f57acb7c1f5..5e80328c9700 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/encryption_config.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/encryption_config.py @@ -14,10 +14,8 @@ # limitations under the License. # import proto # type: ignore - from google.protobuf import wrappers_pb2 # type: ignore - __protobuf__ = proto.module( package="google.cloud.bigquery.v2", manifest={ diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model.py index f32e15eb12b6..3ad9a2e7ebf7 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model.py @@ -14,14 +14,16 @@ # limitations under the License. # import proto # type: ignore - -from google.cloud.bigquery_v2.types import encryption_config +from google.cloud.bigquery_v2.types import ( + encryption_config, + standard_sql, + table_reference, +) from google.cloud.bigquery_v2.types import model_reference as gcb_model_reference -from google.cloud.bigquery_v2.types import standard_sql -from google.cloud.bigquery_v2.types import table_reference -from google.protobuf import timestamp_pb2 # type: ignore -from google.protobuf import wrappers_pb2 # type: ignore - +from google.protobuf import ( + timestamp_pb2, # type: ignore + wrappers_pb2, # type: ignore +) __protobuf__ = proto.module( package="google.cloud.bigquery.v2", diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model_reference.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model_reference.py index cde139ebe5cd..d462938a00e1 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model_reference.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/model_reference.py @@ -15,7 +15,6 @@ # import proto # type: ignore - __protobuf__ = proto.module( package="google.cloud.bigquery.v2", manifest={ diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/standard_sql.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/standard_sql.py index 3be5304fc153..426139d85436 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/standard_sql.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/standard_sql.py @@ -15,7 +15,6 @@ # import proto # type: ignore - __protobuf__ = proto.module( package="google.cloud.bigquery.v2", manifest={ diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/table_reference.py b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/table_reference.py index c02eb206f364..65759ca2c24d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/table_reference.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery_v2/types/table_reference.py @@ -15,7 +15,6 @@ # import proto # type: ignore - __protobuf__ = proto.module( package="google.cloud.bigquery.v2", manifest={ diff --git a/packages/google-cloud-bigquery/noxfile.py b/packages/google-cloud-bigquery/noxfile.py index d4accfea91a4..6614c580f371 100644 --- a/packages/google-cloud-bigquery/noxfile.py +++ b/packages/google-cloud-bigquery/noxfile.py @@ -597,7 +597,7 @@ def format(session: nox.sessions.Session) -> None: to format code to uniform standard. """ session.install(BLACK_VERSION, ISORT_VERSION) - python_files = [path for path in os.listdir(".") if path.endswith(".py")] + python_files = ["google", "tests", "setup.py", "noxfile.py"] # Use the --fss option to sort imports using strict alphabetical order. # See https://pycqa.github.io/isort/docs/configuration/options.html#force-sort-within-sections diff --git a/packages/google-cloud-bigquery/tests/system/conftest.py b/packages/google-cloud-bigquery/tests/system/conftest.py index 7ab4ded9bed3..b7633a88487e 100644 --- a/packages/google-cloud-bigquery/tests/system/conftest.py +++ b/packages/google-cloud-bigquery/tests/system/conftest.py @@ -19,11 +19,10 @@ import pytest import test_utils.prefixer - from google.cloud import bigquery from google.cloud.bigquery import enums -from . import helpers +from . import helpers prefixer = test_utils.prefixer.Prefixer("python-bigquery", "tests/system") diff --git a/packages/google-cloud-bigquery/tests/system/helpers.py b/packages/google-cloud-bigquery/tests/system/helpers.py index 6a8e142c2a50..5d66116bb96c 100644 --- a/packages/google-cloud-bigquery/tests/system/helpers.py +++ b/packages/google-cloud-bigquery/tests/system/helpers.py @@ -19,7 +19,6 @@ import google.api_core.exceptions import test_utils.retry - from google.cloud._helpers import UTC _naive = datetime.datetime(2016, 12, 5, 12, 41, 9) diff --git a/packages/google-cloud-bigquery/tests/system/test_job_retry.py b/packages/google-cloud-bigquery/tests/system/test_job_retry.py index 520545493290..9c9bed0a568e 100644 --- a/packages/google-cloud-bigquery/tests/system/test_job_retry.py +++ b/packages/google-cloud-bigquery/tests/system/test_job_retry.py @@ -35,7 +35,7 @@ def test_query_retry_539(bigquery_client, dataset_id, job_retry_on_query): See: https://github.com/googleapis/python-bigquery/issues/539 """ from google.api_core import exceptions - from google.api_core.retry import if_exception_type, Retry + from google.api_core.retry import Retry, if_exception_type table_name = f"{dataset_id}.t539" diff --git a/packages/google-cloud-bigquery/tests/unit/conftest.py b/packages/google-cloud-bigquery/tests/unit/conftest.py index 5070a199bc58..2bba7e9caff8 100644 --- a/packages/google-cloud-bigquery/tests/unit/conftest.py +++ b/packages/google-cloud-bigquery/tests/unit/conftest.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import mock import threading +from unittest import mock import pytest diff --git a/packages/google-cloud-bigquery/tests/unit/helpers.py b/packages/google-cloud-bigquery/tests/unit/helpers.py index c5414138e8f5..47f26e1ea99c 100644 --- a/packages/google-cloud-bigquery/tests/unit/helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/helpers.py @@ -14,10 +14,9 @@ from unittest import mock -import pytest - import google.cloud.bigquery.client import google.cloud.bigquery.dataset +import pytest def make_connection(*responses): diff --git a/packages/google-cloud-bigquery/tests/unit/job/helpers.py b/packages/google-cloud-bigquery/tests/unit/job/helpers.py index 24ba2fa997c0..7d931a352749 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/job/helpers.py @@ -16,7 +16,8 @@ from google.api_core import exceptions -from ..helpers import make_connection, make_client as __make_client +from ..helpers import make_client as __make_client +from ..helpers import make_connection def _make_client(project="test-project", connection=None): @@ -103,6 +104,7 @@ def _make_one(self, *args, **kw): def _setUpConstants(self): import datetime + from google.cloud._helpers import UTC self.WHEN_TS = 1437767599.006 diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_async_job_retry.py b/packages/google-cloud-bigquery/tests/unit/job/test_async_job_retry.py index 4d9cccf932c3..a4c6a7d385fc 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_async_job_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_async_job_retry.py @@ -15,11 +15,10 @@ from unittest import mock import google.api_core.retry +import google.cloud.bigquery.job from google.api_core import exceptions from . import helpers -import google.cloud.bigquery.job - PROJECT = "test-project" JOB_ID = "test-job-id" diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_base.py b/packages/google-cloud-bigquery/tests/unit/job/test_base.py index 4209048205a5..3376d16c4176 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_base.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_base.py @@ -17,16 +17,12 @@ import unittest from unittest import mock -from google.api_core.future import polling import pytest - +from google.api_core.future import polling from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT from ..helpers import make_connection - -from .helpers import _make_client -from .helpers import _make_retriable_exception -from .helpers import _make_job_resource +from .helpers import _make_client, _make_job_resource, _make_retriable_exception class Test__error_result_to_exception(unittest.TestCase): @@ -327,6 +323,7 @@ def test_user_email(self): @staticmethod def _datetime_and_millis(): import datetime + from google.cloud._helpers import _millis now = datetime.datetime.now(datetime.timezone.utc).replace( @@ -645,8 +642,8 @@ def test__begin_explicit(self): self.assertEqual(job._properties, expected) def test_exists_defaults_miss(self): - from google.cloud.exceptions import NotFound from google.cloud.bigquery.retry import DEFAULT_RETRY + from google.cloud.exceptions import NotFound job = self._set_properties_job() job._properties["jobReference"]["location"] = self.LOCATION @@ -719,7 +716,7 @@ def test_exists_w_timeout(self): ) def test_reload_defaults(self): - from google.cloud.bigquery.retry import DEFAULT_RETRY, DEFAULT_GET_JOB_TIMEOUT + from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT, DEFAULT_RETRY resource = { "jobReference": { diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_copy.py b/packages/google-cloud-bigquery/tests/unit/job/test_copy.py index 8e2845316d13..3100cfa79d9e 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_copy.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_copy.py @@ -12,14 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime from unittest import mock from ..helpers import make_connection - -from .helpers import _Base -from .helpers import _make_client - -import datetime +from .helpers import _Base, _make_client class TestCopyJobConfig(_Base): @@ -43,9 +40,11 @@ def test_ctor_defaults(self): assert config.operation_type == OperationType.OPERATION_TYPE_UNSPECIFIED def test_ctor_w_properties(self): - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import OperationType - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CreateDisposition, + OperationType, + WriteDisposition, + ) create_disposition = CreateDisposition.CREATE_NEVER write_disposition = WriteDisposition.WRITE_TRUNCATE @@ -382,10 +381,11 @@ def test_begin_w_bound_client(self): self._verifyResourceProperties(job, RESOURCE) def test_begin_w_alternate_client(self): - from google.cloud.bigquery.job import CopyJobConfig - - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CopyJobConfig, + CreateDisposition, + WriteDisposition, + ) PATH = "/projects/%s/jobs" % (self.PROJECT,) RESOURCE = self._make_resource(ended=True) diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_extract.py b/packages/google-cloud-bigquery/tests/unit/job/test_extract.py index ebf9f09e6991..ae51d2995cc9 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_extract.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_extract.py @@ -16,9 +16,7 @@ from unittest import mock from ..helpers import make_connection - -from .helpers import _Base -from .helpers import _make_client +from .helpers import _Base, _make_client class TestExtractJobConfig(_Base): @@ -306,9 +304,11 @@ def test_begin_w_bound_client(self): def test_begin_w_alternate_client(self): from google.cloud.bigquery.dataset import DatasetReference - from google.cloud.bigquery.job import Compression - from google.cloud.bigquery.job import DestinationFormat - from google.cloud.bigquery.job import ExtractJobConfig + from google.cloud.bigquery.job import ( + Compression, + DestinationFormat, + ExtractJobConfig, + ) PATH = "/projects/%s/jobs" % (self.PROJECT,) RESOURCE = self._make_resource(ended=True) diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_load.py b/packages/google-cloud-bigquery/tests/unit/job/test_load.py index b551d52dd327..098c394c14e1 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_load.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_load.py @@ -15,12 +15,11 @@ import copy from unittest import mock -from ..helpers import make_connection - -from .helpers import _Base -from .helpers import _make_client from google.cloud.bigquery.enums import SourceColumnMatch +from ..helpers import make_connection +from .helpers import _Base, _make_client + class TestLoadJob(_Base): JOB_TYPE = "load" @@ -251,8 +250,8 @@ def test_ctor(self): self.assertIsNone(job.timestamp_format) def test_ctor_w_config(self): - from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.job import LoadJobConfig + from google.cloud.bigquery.schema import SchemaField client = _make_client(project=self.PROJECT) full_name = SchemaField("full_name", "STRING", mode="REQUIRED") @@ -342,8 +341,8 @@ def test_schema_setter(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, 8, 11, 12, 13, 22, tzinfo=UTC) STARTED = datetime.datetime(2015, 8, 11, 13, 47, 15, tzinfo=UTC) @@ -600,10 +599,12 @@ def test_begin_w_autodetect(self): self._verifyResourceProperties(job, resource) def test_begin_w_alternate_client(self): - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import LoadJobConfig - from google.cloud.bigquery.job import SchemaUpdateOption - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CreateDisposition, + LoadJobConfig, + SchemaUpdateOption, + WriteDisposition, + ) from google.cloud.bigquery.schema import SchemaField PATH = "/projects/%s/jobs" % (self.PROJECT,) diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_load_config.py b/packages/google-cloud-bigquery/tests/unit/job/test_load_config.py index 2e046bfbfb97..36786c7055d0 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_load_config.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_load_config.py @@ -716,8 +716,7 @@ def test_range_partitioning_w_value(self): object_under_test.range_partitioning.range_.interval == 10 def test_range_partitioning_setter(self): - from google.cloud.bigquery.table import PartitionRange - from google.cloud.bigquery.table import RangePartitioning + from google.cloud.bigquery.table import PartitionRange, RangePartitioning object_under_test = self._get_target_class()() object_under_test.range_partitioning = RangePartitioning( @@ -743,8 +742,7 @@ def test_time_partitioning_miss(self): self.assertIsNone(config.time_partitioning) def test_time_partitioning_hit(self): - from google.cloud.bigquery.table import TimePartitioning - from google.cloud.bigquery.table import TimePartitioningType + from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType field = "creation_date" year_ms = 86400 * 1000 * 365 @@ -769,8 +767,7 @@ def test_time_partitioning_hit(self): assert "TimePartitioning.require_partition_filter" in str(warning) 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 field = "creation_date" year_ms = 86400 * 1000 * 365 @@ -1102,11 +1099,10 @@ def test_from_api_repr(self): SourceFormat, WriteDisposition, ) + from google.cloud.bigquery.job.load import ColumnNameCharacterMap from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType - from google.cloud.bigquery.job.load import ColumnNameCharacterMap - config = LoadJobConfig.from_api_repr(self.RESOURCE) self.assertTrue(config.allow_jagged_rows) @@ -1137,16 +1133,16 @@ def test_from_api_repr(self): self.assertEqual(config.timestamp_target_precision, [6, 12]) def test_to_api_repr(self): + from google.cloud.bigquery.format_options import ParquetOptions from google.cloud.bigquery.job import ( CreateDisposition, LoadJobConfig, SourceFormat, WriteDisposition, ) + from google.cloud.bigquery.job.load import ColumnNameCharacterMap from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType - from google.cloud.bigquery.format_options import ParquetOptions - from google.cloud.bigquery.job.load import ColumnNameCharacterMap config = LoadJobConfig() config.allow_jagged_rows = True diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query.py b/packages/google-cloud-bigquery/tests/unit/job/test_query.py index 3720e2a06c13..69820a6c692f 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query.py @@ -20,18 +20,15 @@ import types from unittest import mock -import requests - -from google.cloud.bigquery.client import _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS import google.cloud.bigquery._job_helpers import google.cloud.bigquery.query +import requests +from google.cloud.bigquery.client import _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT from google.cloud.bigquery.table import _EmptyRowIterator from ..helpers import make_connection - -from .helpers import _Base -from .helpers import _make_client +from .helpers import _Base, _make_client class TestQueryJob(_Base): @@ -340,9 +337,11 @@ def test_from_api_repr_with_transaction_info(self): self._verifyResourceProperties(job, RESOURCE) def test_from_api_repr_w_properties(self): - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import SchemaUpdateOption - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CreateDisposition, + SchemaUpdateOption, + WriteDisposition, + ) client = _make_client(project=self.PROJECT) RESOURCE = self._make_resource() @@ -372,8 +371,7 @@ def test_cancelled(self): def test_query_plan(self): from google.cloud._helpers import _RFC3339_MICROS - from google.cloud.bigquery.job import QueryPlanEntry - from google.cloud.bigquery.job import QueryPlanEntryStep + from google.cloud.bigquery.job import QueryPlanEntry, QueryPlanEntryStep plan_entries = [ { @@ -760,9 +758,11 @@ def test_timeline(self): self.assertEqual(job.timeline[0].slot_millis, 101) def test_undeclared_query_parameters(self): - from google.cloud.bigquery.query import ArrayQueryParameter - from google.cloud.bigquery.query import ScalarQueryParameter - from google.cloud.bigquery.query import StructQueryParameter + from google.cloud.bigquery.query import ( + ArrayQueryParameter, + ScalarQueryParameter, + StructQueryParameter, + ) undeclared = [ { @@ -1856,11 +1856,13 @@ def test_begin_w_bound_client(self): def test_begin_w_alternate_client(self): from google.cloud.bigquery.dataset import DatasetReference - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import QueryJobConfig - from google.cloud.bigquery.job import QueryPriority - from google.cloud.bigquery.job import SchemaUpdateOption - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CreateDisposition, + QueryJobConfig, + QueryPriority, + SchemaUpdateOption, + WriteDisposition, + ) PATH = "/projects/%s/jobs" % (self.PROJECT,) TABLE = "TABLE" @@ -2084,10 +2086,12 @@ def test_begin_w_positional_query_parameter(self): self._verifyResourceProperties(job, RESOURCE) def test_begin_w_table_defs(self): + from google.cloud.bigquery.external_config import ( + BigtableColumn, + BigtableColumnFamily, + ExternalConfig, + ) from google.cloud.bigquery.job import QueryJobConfig - from google.cloud.bigquery.external_config import ExternalConfig - from google.cloud.bigquery.external_config import BigtableColumn - from google.cloud.bigquery.external_config import BigtableColumnFamily PATH = "/projects/%s/jobs" % (self.PROJECT,) RESOURCE = self._make_resource() diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query_config.py b/packages/google-cloud-bigquery/tests/unit/job/test_query_config.py index a63a14b73f46..67ce4f546f21 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query_config.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query_config.py @@ -103,8 +103,7 @@ def test_range_partitioning_w_value(self): object_under_test.range_partitioning.range_.interval == 10 def test_range_partitioning_setter(self): - from google.cloud.bigquery.table import PartitionRange - from google.cloud.bigquery.table import RangePartitioning + from google.cloud.bigquery.table import PartitionRange, RangePartitioning object_under_test = self._get_target_class()() object_under_test.range_partitioning = RangePartitioning( @@ -295,8 +294,7 @@ def test_to_api_repr_with_script_options_none(self): self.assertIsNone(config.script_options) def test_to_api_repr_with_script_options(self): - from google.cloud.bigquery import KeyResultStatementKind - from google.cloud.bigquery import ScriptOptions + from google.cloud.bigquery import KeyResultStatementKind, ScriptOptions config = self._make_one() config.script_options = ScriptOptions( @@ -317,8 +315,7 @@ def test_to_api_repr_with_script_options(self): ) def test_from_api_repr_with_script_options(self): - from google.cloud.bigquery import KeyResultStatementKind - from google.cloud.bigquery import ScriptOptions + from google.cloud.bigquery import KeyResultStatementKind, ScriptOptions resource = { "query": { diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query_job_retry.py b/packages/google-cloud-bigquery/tests/unit/job/test_query_job_retry.py index c8355b68832c..5a3105ff9305 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query_job_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query_job_retry.py @@ -11,14 +11,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import concurrent.futures from unittest import mock -import concurrent.futures import freezegun -from google.api_core import exceptions import google.api_core.retry import pytest - +from google.api_core import exceptions from google.cloud.bigquery.client import _MIN_GET_QUERY_RESULTS_TIMEOUT from google.cloud.bigquery.job import QueryJob from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT @@ -27,7 +26,6 @@ from ..helpers import make_connection from .helpers import _make_client - PROJECT = "test-project" JOB_ID = "test-job-id" QUERY = "select count(*) from persons" diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query_pandas.py b/packages/google-cloud-bigquery/tests/unit/job/test_query_pandas.py index 82477745aabe..dad747d034d3 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query_pandas.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query_pandas.py @@ -18,16 +18,15 @@ from unittest import mock import pytest +from google.cloud.bigquery.enums import DefaultPandasDTypes from ..helpers import make_connection -from .helpers import _make_client -from .helpers import _make_job_resource -from google.cloud.bigquery.enums import DefaultPandasDTypes +from .helpers import _make_client, _make_job_resource try: - from google.cloud import bigquery_storage import google.cloud.bigquery_storage_v1.reader import google.cloud.bigquery_storage_v1.services.big_query_read.client + from google.cloud import bigquery_storage except (ImportError, AttributeError): bigquery_storage = None @@ -312,9 +311,9 @@ def test_to_arrow_max_results_no_progress_bar(): @mock.patch("google.cloud.bigquery._tqdm_helpers.tqdm") def test_to_arrow_w_tqdm_w_query_plan(tqdm_mock): from google.cloud.bigquery import table + from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL from google.cloud.bigquery.job import QueryJob as target_class from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL begun_resource = _make_job_resource(job_type="query") rows = [ @@ -369,9 +368,9 @@ def test_to_arrow_w_tqdm_w_query_plan(tqdm_mock): @mock.patch("google.cloud.bigquery._tqdm_helpers.tqdm") def test_to_arrow_w_tqdm_w_pending_status(tqdm_mock): from google.cloud.bigquery import table + from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL from google.cloud.bigquery.job import QueryJob as target_class from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL begun_resource = _make_job_resource(job_type="query") rows = [ @@ -783,9 +782,9 @@ def test_to_dataframe_with_progress_bar(tqdm_mock): @mock.patch("google.cloud.bigquery._tqdm_helpers.tqdm") def test_to_dataframe_w_tqdm_pending(tqdm_mock): from google.cloud.bigquery import table + from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL from google.cloud.bigquery.job import QueryJob as target_class from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL begun_resource = _make_job_resource(job_type="query") schema = [ @@ -838,9 +837,9 @@ def test_to_dataframe_w_tqdm_pending(tqdm_mock): @mock.patch("google.cloud.bigquery._tqdm_helpers.tqdm") def test_to_dataframe_w_tqdm(tqdm_mock): from google.cloud.bigquery import table + from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL from google.cloud.bigquery.job import QueryJob as target_class from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL begun_resource = _make_job_resource(job_type="query") schema = [ @@ -898,9 +897,9 @@ def test_to_dataframe_w_tqdm(tqdm_mock): @mock.patch("google.cloud.bigquery._tqdm_helpers.tqdm") def test_to_dataframe_w_tqdm_max_results(tqdm_mock): from google.cloud.bigquery import table + from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL from google.cloud.bigquery.job import QueryJob as target_class from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery._tqdm_helpers import _PROGRESS_BAR_UPDATE_INTERVAL begun_resource = _make_job_resource(job_type="query") schema = [ diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query_stats.py b/packages/google-cloud-bigquery/tests/unit/job/test_query_stats.py index c7c7a31e05be..7de53c55f213 100644 --- a/packages/google-cloud-bigquery/tests/unit/job/test_query_stats.py +++ b/packages/google-cloud-bigquery/tests/unit/job/test_query_stats.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .helpers import _Base import datetime +from .helpers import _Base + class TestBiEngineStats: @staticmethod diff --git a/packages/google-cloud-bigquery/tests/unit/line_arg_parser/test_parser.py b/packages/google-cloud-bigquery/tests/unit/line_arg_parser/test_parser.py index b170d536aeee..2d6c62fcf5db 100644 --- a/packages/google-cloud-bigquery/tests/unit/line_arg_parser/test_parser.py +++ b/packages/google-cloud-bigquery/tests/unit/line_arg_parser/test_parser.py @@ -36,8 +36,7 @@ def test_consume_expected_eol(parser_class): def test_consume_unexpected_eol(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. @@ -49,8 +48,7 @@ def test_consume_unexpected_eol(parser_class): def test_input_line_unexpected_input(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. @@ -66,8 +64,7 @@ def test_input_line_unexpected_input(parser_class): def test_destination_var_unexpected_input(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. @@ -82,8 +79,7 @@ def test_destination_var_unexpected_input(parser_class): def test_option_value_unexpected_input(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. @@ -133,8 +129,7 @@ def test_dict_items_trailing_comma(parser_class): def test_dict_item_unknown_input(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. @@ -148,8 +143,7 @@ def test_dict_item_unknown_input(parser_class): def test_pyvalue_list_containing_dict(parser_class): from google.cloud.bigquery.magics.line_arg_parser import TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token - from google.cloud.bigquery.magics.line_arg_parser.parser import PyDict - from google.cloud.bigquery.magics.line_arg_parser.parser import PyList + from google.cloud.bigquery.magics.line_arg_parser.parser import PyDict, PyList # A simple iterable of Tokens is sufficient. fake_lexer = [ @@ -180,8 +174,7 @@ def test_pyvalue_list_containing_dict(parser_class): def test_pyvalue_invalid_token(parser_class): - from google.cloud.bigquery.magics.line_arg_parser import ParseError - from google.cloud.bigquery.magics.line_arg_parser import TokenType + from google.cloud.bigquery.magics.line_arg_parser import ParseError, TokenType from google.cloud.bigquery.magics.line_arg_parser.lexer import Token # A simple iterable of Tokens is sufficient. diff --git a/packages/google-cloud-bigquery/tests/unit/model/test_model.py b/packages/google-cloud-bigquery/tests/unit/model/test_model.py index 279a954c78ea..3ac7f4725855 100644 --- a/packages/google-cloud-bigquery/tests/unit/model/test_model.py +++ b/packages/google-cloud-bigquery/tests/unit/model/test_model.py @@ -16,11 +16,9 @@ import datetime -import pytest - - import google.cloud._helpers import google.cloud.bigquery.model +import pytest KMS_KEY_NAME = "projects/1/locations/us/keyRings/1/cryptoKeys/1" diff --git a/packages/google-cloud-bigquery/tests/unit/routine/test_routine.py b/packages/google-cloud-bigquery/tests/unit/routine/test_routine.py index 965c6b2eb882..4b4b059f5557 100644 --- a/packages/google-cloud-bigquery/tests/unit/routine/test_routine.py +++ b/packages/google-cloud-bigquery/tests/unit/routine/test_routine.py @@ -15,9 +15,8 @@ import datetime -import pytest - import google.cloud._helpers +import pytest from google.cloud import bigquery @@ -54,8 +53,7 @@ def test_ctor_w_string(target_class): def test_ctor_w_properties(target_class): - from google.cloud.bigquery.routine import RoutineArgument - from google.cloud.bigquery.routine import RoutineReference + from google.cloud.bigquery.routine import RoutineArgument, RoutineReference routine_id = "my-proj.my_dset.my_routine" arguments = [ @@ -140,8 +138,7 @@ def test_ctor_invalid_external_runtime_options(target_class): def test_from_api_repr(target_class): - from google.cloud.bigquery.routine import RoutineArgument - from google.cloud.bigquery.routine import RoutineReference + from google.cloud.bigquery.routine import RoutineArgument, RoutineReference creation_time = datetime.datetime( 2010, 5, 19, 16, 0, 0, tzinfo=google.cloud._helpers.UTC @@ -232,9 +229,11 @@ def test_from_api_repr(target_class): def test_from_api_repr_tvf_function(target_class): - from google.cloud.bigquery.routine import RoutineArgument - from google.cloud.bigquery.routine import RoutineReference - from google.cloud.bigquery.routine import RoutineType + from google.cloud.bigquery.routine import ( + RoutineArgument, + RoutineReference, + RoutineType, + ) StandardSqlDataType = bigquery.standard_sql.StandardSqlDataType StandardSqlField = bigquery.standard_sql.StandardSqlField diff --git a/packages/google-cloud-bigquery/tests/unit/test__helpers.py b/packages/google-cloud-bigquery/tests/unit/test__helpers.py index 4e53236e39e7..140e6b5a3022 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__helpers.py @@ -16,13 +16,13 @@ import decimal import json import os -import warnings -import pytest -import packaging import unittest +import warnings from unittest import mock import google.api_core +import packaging +import pytest from google.cloud.bigquery._helpers import _isinstance_or_raise @@ -50,8 +50,8 @@ def test_with_dict_empty(self): self.assertEqual("googleapis.com", _get_client_universe(options)) def test_with_client_options(self): - from google.cloud.bigquery._helpers import _get_client_universe from google.api_core import client_options + from google.cloud.bigquery._helpers import _get_client_universe options = client_options.from_dict({"universe_domain": "foo.com"}) self.assertEqual("foo.com", _get_client_universe(options)) @@ -71,8 +71,8 @@ def test_with_environ_and_dict(self): @mock.patch.dict(os.environ, {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": "foo.com"}) def test_with_environ_and_empty_options(self): - from google.cloud.bigquery._helpers import _get_client_universe from google.api_core import client_options + from google.cloud.bigquery._helpers import _get_client_universe options = client_options.from_dict({}) self.assertEqual("foo.com", _get_client_universe(options)) @@ -93,6 +93,7 @@ def test_with_none(self): def test_with_no_universe_creds(self): from google.cloud.bigquery._helpers import _validate_universe + from .helpers import make_creds creds = make_creds(None) @@ -101,6 +102,7 @@ def test_with_no_universe_creds(self): def test_with_matched_universe_creds(self): from google.cloud.bigquery._helpers import _validate_universe + from .helpers import make_creds creds = make_creds("googleapis.com") @@ -109,6 +111,7 @@ def test_with_matched_universe_creds(self): def test_with_mismatched_universe_creds(self): from google.cloud.bigquery._helpers import _validate_universe + from .helpers import make_creds creds = make_creds("foo.com") @@ -1146,9 +1149,10 @@ def fake_isinstance(instance, target_class): def test_decimal_as_float_api_repr(): """Make sure decimals get converted to float.""" - import google.cloud.bigquery.query from decimal import Decimal + import google.cloud.bigquery.query + param = google.cloud.bigquery.query.ScalarQueryParameter( "x", "FLOAT64", Decimal(42) ) diff --git a/packages/google-cloud-bigquery/tests/unit/test__http.py b/packages/google-cloud-bigquery/tests/unit/test__http.py index fd7ecdc428a2..dc898065ca54 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__http.py +++ b/packages/google-cloud-bigquery/tests/unit/test__http.py @@ -38,8 +38,7 @@ def _make_one(self, *args, **kw): return self._get_target_class()(*args, **kw) def test_build_api_url_no_extra_query_params(self): - from urllib.parse import parse_qsl - from urllib.parse import urlsplit + from urllib.parse import parse_qsl, urlsplit conn = self._make_one(object()) uri = conn.build_api_url("/foo") @@ -52,8 +51,7 @@ def test_build_api_url_no_extra_query_params(self): self.assertEqual(parms, {}) def test_build_api_url_w_custom_endpoint(self): - from urllib.parse import parse_qsl - from urllib.parse import urlsplit + from urllib.parse import parse_qsl, urlsplit custom_endpoint = "https://foo-bigquery.googleapis.com" conn = self._make_one(object(), api_endpoint=custom_endpoint) @@ -67,8 +65,7 @@ def test_build_api_url_w_custom_endpoint(self): self.assertEqual(parms, {}) def test_build_api_url_w_extra_query_params(self): - from urllib.parse import parse_qsl - from urllib.parse import urlsplit + from urllib.parse import parse_qsl, urlsplit conn = self._make_one(object()) uri = conn.build_api_url("/foo", {"bar": "baz"}) diff --git a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py index 19390c7ec2bf..1378eb7edfa9 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__job_helpers.py @@ -16,12 +16,9 @@ from unittest import mock import google.api_core.exceptions -from google.api_core import retry as retries import pytest - -from google.cloud.bigquery import _job_helpers -from google.cloud.bigquery import enums -from google.cloud.bigquery import retry +from google.api_core import retry as retries +from google.cloud.bigquery import _job_helpers, enums, retry from google.cloud.bigquery.client import Client from google.cloud.bigquery.job import copy_ as job_copy from google.cloud.bigquery.job import extract as job_extract diff --git a/packages/google-cloud-bigquery/tests/unit/test__job_helpers_retry.py b/packages/google-cloud-bigquery/tests/unit/test__job_helpers_retry.py index 3ea4b1aae6c7..55759320a5f7 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__job_helpers_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/test__job_helpers_retry.py @@ -14,9 +14,8 @@ import freezegun import google.api_core.exceptions -from google.api_core import retry as retries import pytest - +from google.api_core import retry as retries from google.cloud.bigquery import _job_helpers from . import helpers 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..a4e9c0c9dc78 100644 --- a/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py @@ -44,3 +44,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_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 9df85b7372e5..b93e8dc14919 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -6841,10 +6841,10 @@ def test_insert_rows_w_wrong_arg(self): client.insert_rows_json(table, ROW) def test_insert_rows_json_w_ssl_error(self): + import requests.exceptions from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.table import Table - import requests.exceptions PROJECT = "PROJECT" DS_ID = "DS_ID" diff --git a/packages/google-cloud-bigquery/tests/unit/test_client_bigframes.py b/packages/google-cloud-bigquery/tests/unit/test_client_bigframes.py index 0260da5e4bc8..a4254e8c0df1 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client_bigframes.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client_bigframes.py @@ -19,15 +19,13 @@ import datetime from unittest import mock -import pytest - import google.auth.credentials +import google.cloud.bigquery.client +import pytest from google.api_core import exceptions from google.cloud import bigquery -import google.cloud.bigquery.client from google.cloud.bigquery import _job_helpers - PROJECT = "test-project" LOCATION = "test-location" diff --git a/packages/google-cloud-bigquery/tests/unit/test_client_resumable_media_upload.py b/packages/google-cloud-bigquery/tests/unit/test_client_resumable_media_upload.py index 642c18d15c11..da9852fdf04e 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client_resumable_media_upload.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client_resumable_media_upload.py @@ -12,19 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import mock import email import http.client import io import json +from unittest import mock import pytest - from google.cloud.bigquery.table import TableReference from .helpers import make_connection - PROJECT = "test-project" TABLE_REF = TableReference.from_string(f"{PROJECT}.test_dataset.test_table") EXPECTED_CONFIGURATION = { @@ -113,13 +111,13 @@ def _mock_transport(status_code, headers, content=b""): def _initiate_resumable_upload_helper(num_retries=None, mtls=False): + from google.cloud.bigquery.client import ( + _DEFAULT_CHUNKSIZE, + _GENERIC_CONTENT_TYPE, + _get_upload_headers, + ) + from google.cloud.bigquery.job import LoadJob, LoadJobConfig, SourceFormat from google.resumable_media.requests import ResumableUpload - from google.cloud.bigquery.client import _DEFAULT_CHUNKSIZE - from google.cloud.bigquery.client import _GENERIC_CONTENT_TYPE - from google.cloud.bigquery.client import _get_upload_headers - from google.cloud.bigquery.job import LoadJob - from google.cloud.bigquery.job import LoadJobConfig - from google.cloud.bigquery.job import SourceFormat # Create mocks to be checked for doing transport. resumable_url = "http://test.invalid?upload_id=hey-you" @@ -198,9 +196,7 @@ def _do_multipart_upload_success_helper( get_boundary, num_retries=None, project=None, mtls=False ): from google.cloud.bigquery.client import _get_upload_headers - from google.cloud.bigquery.job import LoadJob - from google.cloud.bigquery.job import LoadJobConfig - from google.cloud.bigquery.job import SourceFormat + from google.cloud.bigquery.job import LoadJob, LoadJobConfig, SourceFormat fake_transport = _mock_transport(http.client.OK, {}) client = _make_client(_http=fake_transport) diff --git a/packages/google-cloud-bigquery/tests/unit/test_client_retry.py b/packages/google-cloud-bigquery/tests/unit/test_client_retry.py index f0e7ac88f302..e860b3eec925 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client_retry.py @@ -18,8 +18,8 @@ import google.api_core.exceptions from google.cloud.bigquery import job as bqjob from google.cloud.bigquery.retry import DEFAULT_RETRY -from .helpers import make_connection +from .helpers import make_connection PROJECT = "test-project" diff --git a/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py b/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py index b144471ca503..6c767a9407ad 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py +++ b/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py @@ -14,11 +14,12 @@ from unittest import mock -from google.cloud.bigquery.dataset import Dataset, DatasetReference -from .helpers import make_connection, dataset_polymorphic, make_client import google.cloud.bigquery.dataset -from google.cloud.bigquery.retry import DEFAULT_TIMEOUT import pytest +from google.cloud.bigquery.dataset import Dataset, DatasetReference +from google.cloud.bigquery.retry import DEFAULT_TIMEOUT + +from .helpers import dataset_polymorphic, make_client, make_connection @dataset_polymorphic diff --git a/packages/google-cloud-bigquery/tests/unit/test_dataset.py b/packages/google-cloud-bigquery/tests/unit/test_dataset.py index 98466544b5ae..0f7048715934 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_dataset.py +++ b/packages/google-cloud-bigquery/tests/unit/test_dataset.py @@ -15,7 +15,6 @@ import unittest from unittest import mock -from google.cloud.bigquery.routine.routine import Routine, RoutineReference import pytest from google.cloud.bigquery.dataset import ( AccessEntry, @@ -25,6 +24,7 @@ Table, TableReference, ) +from google.cloud.bigquery.routine.routine import Routine, RoutineReference class TestAccessEntry(unittest.TestCase): @@ -894,6 +894,7 @@ def _make_one(self, *args, **kw): def _setUpConstants(self): import datetime + from google.cloud._helpers import UTC self.WHEN_TS = 1437767599.006 @@ -1032,7 +1033,7 @@ def test_ctor_string_wo_project_id(self): self._make_one("some_dset") def test_ctor_explicit(self): - from google.cloud.bigquery.dataset import DatasetReference, AccessEntry + from google.cloud.bigquery.dataset import AccessEntry, DatasetReference phred = AccessEntry("OWNER", "userByEmail", "phred@example.com") bharney = AccessEntry("OWNER", "userByEmail", "bharney@example.com") diff --git a/packages/google-cloud-bigquery/tests/unit/test_dbapi__helpers.py b/packages/google-cloud-bigquery/tests/unit/test_dbapi__helpers.py index 9907df97b275..fcd3b8dcb5a1 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_dbapi__helpers.py +++ b/packages/google-cloud-bigquery/tests/unit/test_dbapi__helpers.py @@ -19,12 +19,10 @@ import re import unittest -import pytest - import google.cloud._helpers +import pytest from google.cloud.bigquery import query, table -from google.cloud.bigquery.dbapi import _helpers -from google.cloud.bigquery.dbapi import exceptions +from google.cloud.bigquery.dbapi import _helpers, exceptions class TestQueryParameters(unittest.TestCase): @@ -595,8 +593,8 @@ def _expected_error_match(expect): ], ) def test_complex_query_parameter_type_errors(type_, value, expect): - from google.cloud.bigquery.dbapi._helpers import complex_query_parameter from google.cloud.bigquery.dbapi import exceptions + from google.cloud.bigquery.dbapi._helpers import complex_query_parameter with pytest.raises( exceptions.ProgrammingError, diff --git a/packages/google-cloud-bigquery/tests/unit/test_dbapi_connection.py b/packages/google-cloud-bigquery/tests/unit/test_dbapi_connection.py index 8047462243dd..92458971e255 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_dbapi_connection.py +++ b/packages/google-cloud-bigquery/tests/unit/test_dbapi_connection.py @@ -13,10 +13,11 @@ # limitations under the License. import gc -import pytest import unittest from unittest import mock +import pytest + class TestConnection(unittest.TestCase): @staticmethod @@ -77,8 +78,7 @@ def test_ctor_w_bqstorage_client(self): @mock.patch("google.cloud.bigquery.Client", autospec=True) def test_connect_wo_client(self, mock_client): - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import Connection + from google.cloud.bigquery.dbapi import Connection, connect connection = connect() self.assertIsInstance(connection, Connection) @@ -87,8 +87,7 @@ def test_connect_wo_client(self, mock_client): def test_connect_w_client(self): pytest.importorskip("google.cloud.bigquery_storage") - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import Connection + from google.cloud.bigquery.dbapi import Connection, connect mock_client = self._mock_client() mock_bqstorage_client = self._mock_bqstorage_client() @@ -103,8 +102,7 @@ def test_connect_w_client(self): def test_connect_w_both_clients(self): pytest.importorskip("google.cloud.bigquery_storage") - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import Connection + from google.cloud.bigquery.dbapi import Connection, connect mock_client = self._mock_client() mock_bqstorage_client = self._mock_bqstorage_client() @@ -124,8 +122,7 @@ def test_connect_w_both_clients(self): def test_connect_prefer_bqstorage_client_false(self): pytest.importorskip("google.cloud.bigquery_storage") - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import Connection + from google.cloud.bigquery.dbapi import Connection, connect mock_client = self._mock_client() mock_bqstorage_client = self._mock_bqstorage_client() diff --git a/packages/google-cloud-bigquery/tests/unit/test_dbapi_cursor.py b/packages/google-cloud-bigquery/tests/unit/test_dbapi_cursor.py index c5cad8c91085..db4a7e56be15 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_dbapi_cursor.py +++ b/packages/google-cloud-bigquery/tests/unit/test_dbapi_cursor.py @@ -17,10 +17,8 @@ import unittest from unittest import mock -import pytest - import google.cloud.bigquery.table as bq_table - +import pytest from google.api_core import exceptions from tests.unit.helpers import _to_pyarrow @@ -180,8 +178,7 @@ def _mock_results(self, total_rows=0, schema=None, num_dml_affected_rows=None): return mock_results def test_ctor(self): - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import Cursor + from google.cloud.bigquery.dbapi import Cursor, connect connection = connect(self._mock_client()) cursor = self._make_one(connection) @@ -519,8 +516,8 @@ def test_execute_w_default_config(self): self.assertIsNone(used_config) def test_execute_custom_job_config_wo_default_config(self): - from google.cloud.bigquery.dbapi import connect from google.cloud.bigquery import job + from google.cloud.bigquery.dbapi import connect config = job.QueryJobConfig(use_legacy_sql=True) client = self._mock_client(rows=[], num_dml_affected_rows=0) @@ -533,8 +530,8 @@ def test_execute_custom_job_config_wo_default_config(self): self.assertEqual(kwargs["job_config"], config) def test_execute_custom_job_config_w_default_config(self): - from google.cloud.bigquery.dbapi import connect from google.cloud.bigquery import job + from google.cloud.bigquery.dbapi import connect client = self._mock_client(rows=[], num_dml_affected_rows=0) connection = connect(client) @@ -563,8 +560,8 @@ def test_execute_w_dml(self): self.assertEqual(rows, []) def test_execute_w_query(self): - from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery import dbapi + from google.cloud.bigquery.schema import SchemaField connection = dbapi.connect( self._mock_client( @@ -607,9 +604,9 @@ def test_execute_w_query(self): self.assertIsNone(row) def test_execute_w_query_dry_run(self): + from google.cloud.bigquery import dbapi from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.schema import SchemaField - from google.cloud.bigquery import dbapi connection = dbapi.connect( self._mock_client( @@ -637,10 +634,8 @@ def test_execute_w_query_dry_run(self): def test_execute_raises_if_result_raises(self): import google.cloud.exceptions - from google.cloud.bigquery import client - from google.cloud.bigquery.dbapi import connect - from google.cloud.bigquery.dbapi import exceptions + from google.cloud.bigquery.dbapi import connect, exceptions client = mock.create_autospec(client.Client) client.query_and_wait.side_effect = google.cloud.exceptions.GoogleCloudError("") @@ -702,7 +697,7 @@ def test_query_job_wo_execute(self): self.assertIsNone(cursor.query_job) def test_query_job_w_execute(self): - from google.cloud.bigquery import dbapi, QueryJob + from google.cloud.bigquery import QueryJob, dbapi connection = dbapi.connect(self._mock_client()) cursor = connection.cursor() @@ -722,7 +717,7 @@ def test_query_job_w_execute_no_job(self): self.assertIsNone(cursor.query_job) def test_query_job_w_executemany(self): - from google.cloud.bigquery import dbapi, QueryJob + from google.cloud.bigquery import QueryJob, dbapi connection = dbapi.connect(self._mock_client()) cursor = connection.cursor() @@ -932,8 +927,8 @@ def test__extract_types(inp, expect): ], ) def test__extract_types_fail(match, inp): - from google.cloud.bigquery.dbapi.cursor import _extract_types as et from google.cloud.bigquery.dbapi import exceptions + from google.cloud.bigquery.dbapi.cursor import _extract_types as et with pytest.raises(exceptions.ProgrammingError, match=match): et(inp) diff --git a/packages/google-cloud-bigquery/tests/unit/test_dbapi_types.py b/packages/google-cloud-bigquery/tests/unit/test_dbapi_types.py index 7319aa0161c3..8b2b6d88729f 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_dbapi_types.py +++ b/packages/google-cloud-bigquery/tests/unit/test_dbapi_types.py @@ -15,9 +15,8 @@ import datetime import unittest -import pytest - import google.cloud._helpers +import pytest from google.cloud.bigquery.dbapi import types diff --git a/packages/google-cloud-bigquery/tests/unit/test_delete_dataset.py b/packages/google-cloud-bigquery/tests/unit/test_delete_dataset.py index b48beb1473d5..c6063a71ea9a 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_delete_dataset.py +++ b/packages/google-cloud-bigquery/tests/unit/test_delete_dataset.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .helpers import make_connection, make_client, dataset_polymorphic import google.api_core.exceptions -from google.cloud.bigquery.retry import DEFAULT_TIMEOUT import pytest +from google.cloud.bigquery.retry import DEFAULT_TIMEOUT + +from .helpers import dataset_polymorphic, make_client, make_connection @dataset_polymorphic diff --git a/packages/google-cloud-bigquery/tests/unit/test_external_config.py b/packages/google-cloud-bigquery/tests/unit/test_external_config.py index ea827a560eec..dfa7bbccaa44 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_external_config.py +++ b/packages/google-cloud-bigquery/tests/unit/test_external_config.py @@ -14,14 +14,12 @@ import base64 import copy -from typing import Any, Dict, Optional import unittest - -from google.cloud.bigquery import external_config -from google.cloud.bigquery import schema -from google.cloud.bigquery.enums import SourceColumnMatch +from typing import Any, Dict, Optional import pytest +from google.cloud.bigquery import external_config, schema +from google.cloud.bigquery.enums import SourceColumnMatch class TestExternalConfig(unittest.TestCase): diff --git a/packages/google-cloud-bigquery/tests/unit/test_job_retry.py b/packages/google-cloud-bigquery/tests/unit/test_job_retry.py index fa55e8f6a864..5320ae095ac8 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_job_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/test_job_retry.py @@ -16,19 +16,16 @@ import re from unittest import mock -import pytest - +import freezegun import google.api_core.exceptions import google.api_core.retry -import freezegun +import google.cloud.bigquery.retry +import pytest import requests.exceptions - from google.cloud.bigquery import _job_helpers -import google.cloud.bigquery.retry from .helpers import make_client, make_connection - _RETRY_NOT_FOUND = { "job_retry": google.api_core.retry.Retry( predicate=google.api_core.retry.if_exception_type( @@ -635,9 +632,9 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): ], ) def test_retry_load_job_result(result_retry_param, PROJECT, DS_ID): + import google.cloud.bigquery.retry from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.job.load import LoadJob - import google.cloud.bigquery.retry client = make_client() conn = client._connection = make_connection( diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_datasets.py b/packages/google-cloud-bigquery/tests/unit/test_list_datasets.py index 4ef99fd865a1..ee12ab901471 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_datasets.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_datasets.py @@ -15,8 +15,8 @@ from unittest import mock import pytest - from google.cloud.bigquery.retry import DEFAULT_TIMEOUT + from .helpers import make_connection diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_jobs.py b/packages/google-cloud-bigquery/tests/unit/test_list_jobs.py index edb85af0aa9d..8ae7c78ffc99 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_jobs.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_jobs.py @@ -16,8 +16,8 @@ from unittest import mock import pytest - from google.cloud.bigquery.retry import DEFAULT_TIMEOUT + from .helpers import make_connection @@ -25,12 +25,14 @@ "extra,query", [({}, {}), (dict(page_size=42), dict(maxResults=42))] ) def test_list_jobs_defaults(client, PROJECT, DS_ID, extra, query): - from google.cloud.bigquery.job import CopyJob - from google.cloud.bigquery.job import CreateDisposition - from google.cloud.bigquery.job import ExtractJob - from google.cloud.bigquery.job import LoadJob - from google.cloud.bigquery.job import QueryJob - from google.cloud.bigquery.job import WriteDisposition + from google.cloud.bigquery.job import ( + CopyJob, + CreateDisposition, + ExtractJob, + LoadJob, + QueryJob, + WriteDisposition, + ) SOURCE_TABLE = "source_table" DESTINATION_TABLE = "destination_table" diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_models.py b/packages/google-cloud-bigquery/tests/unit/test_list_models.py index 04932d3572b3..0e42604f947d 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_models.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_models.py @@ -13,9 +13,9 @@ # limitations under the License. import pytest - from google.cloud.bigquery.retry import DEFAULT_TIMEOUT -from .helpers import make_connection, dataset_polymorphic + +from .helpers import dataset_polymorphic, make_connection def test_list_models_empty_w_timeout(client, PROJECT, DS_ID): diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_projects.py b/packages/google-cloud-bigquery/tests/unit/test_list_projects.py index 5260e5246474..539f13a08eb0 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_projects.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_projects.py @@ -15,8 +15,8 @@ from unittest import mock import pytest - from google.cloud.bigquery.retry import DEFAULT_TIMEOUT + from .helpers import make_connection diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_routines.py b/packages/google-cloud-bigquery/tests/unit/test_list_routines.py index 80e62d6bdf01..997c9f8125d2 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_routines.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_routines.py @@ -13,9 +13,9 @@ # limitations under the License. import pytest - from google.cloud.bigquery.retry import DEFAULT_TIMEOUT -from .helpers import make_connection, dataset_polymorphic + +from .helpers import dataset_polymorphic, make_connection def test_list_routines_empty_w_timeout(client): diff --git a/packages/google-cloud-bigquery/tests/unit/test_list_tables.py b/packages/google-cloud-bigquery/tests/unit/test_list_tables.py index 8360f6605630..1105fe841936 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_list_tables.py +++ b/packages/google-cloud-bigquery/tests/unit/test_list_tables.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest - import google.cloud.bigquery.dataset +import pytest from google.cloud.bigquery.retry import DEFAULT_TIMEOUT -from .helpers import make_connection, dataset_polymorphic + +from .helpers import dataset_polymorphic, make_connection @dataset_polymorphic diff --git a/packages/google-cloud-bigquery/tests/unit/test_magics.py b/packages/google-cloud-bigquery/tests/unit/test_magics.py index 03a3a2dbbdba..db14f090fd02 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_magics.py +++ b/packages/google-cloud-bigquery/tests/unit/test_magics.py @@ -37,7 +37,7 @@ bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") IPython = pytest.importorskip("IPython") -interactiveshell = pytest.importorskip("IPython.terminal.interactiveshell") +interactiveshell = pytest.importorskip("IPython.core.interactiveshell") tools = pytest.importorskip("IPython.testing.tools") io = pytest.importorskip("IPython.utils.io") pandas = pytest.importorskip("pandas") @@ -58,8 +58,7 @@ def use_local_magics_context(monkeypatch): @pytest.fixture(scope="session") def ipython(): config = tools.default_config() - config.TerminalInteractiveShell.simple_prompt = True - shell = interactiveshell.TerminalInteractiveShell.instance(config=config) + shell = interactiveshell.InteractiveShell.instance(config=config) return shell @@ -147,6 +146,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 +165,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 +685,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 +855,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 +1979,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 +2003,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 +2032,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 +2200,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_opentelemetry_tracing.py b/packages/google-cloud-bigquery/tests/unit/test_opentelemetry_tracing.py index 57132a1b99b5..43b196a2ff84 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_opentelemetry_tracing.py +++ b/packages/google-cloud-bigquery/tests/unit/test_opentelemetry_tracing.py @@ -35,7 +35,6 @@ raise ImportError(msg) from exc import pytest - from google.cloud.bigquery import opentelemetry_tracing TEST_SPAN_NAME = "bar" @@ -197,8 +196,7 @@ def test_optional_job_attributes(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_no_data_leakage(setup): import google.auth.credentials - from google.cloud.bigquery import client - from google.cloud.bigquery import job + from google.cloud.bigquery import client, job mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) test_client = client.Client( @@ -245,8 +243,8 @@ def test_default_no_data_leakage(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_span_creation_error(setup): import google.auth.credentials - from google.cloud.bigquery import client from google.api_core.exceptions import GoogleAPICallError, InvalidArgument + from google.cloud.bigquery import client mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) test_client = client.Client( diff --git a/packages/google-cloud-bigquery/tests/unit/test_query.py b/packages/google-cloud-bigquery/tests/unit/test_query.py index adb43bcd9b4f..289c19ac5c12 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query.py @@ -169,8 +169,10 @@ def test_from_api_repr(self): self.assertEqual(field._type, "STRING") def test_to_api_repr(self): - from google.cloud.bigquery.query import ScalarQueryParameterType - from google.cloud.bigquery.query import StructQueryParameterType + from google.cloud.bigquery.query import ( + ScalarQueryParameterType, + StructQueryParameterType, + ) array_item_type = StructQueryParameterType( ScalarQueryParameterType("INTEGER", name="weight", description="in kg"), @@ -223,8 +225,10 @@ def test_raises_error_without_any_fields(self): self._make_one() def test_from_api_repr(self): - from google.cloud.bigquery.query import ArrayQueryParameterType - from google.cloud.bigquery.query import ScalarQueryParameterType + from google.cloud.bigquery.query import ( + ArrayQueryParameterType, + ScalarQueryParameterType, + ) api_resource = { "type": "STRUCT", @@ -637,7 +641,7 @@ def test_to_api_repr_w_timestamp_datetime(self): self.assertEqual(param.to_api_repr(), EXPECTED) def test_to_api_repr_w_timestamp_micros(self): - from google.cloud._helpers import _microseconds_from_datetime, UTC + from google.cloud._helpers import UTC, _microseconds_from_datetime now = datetime.datetime.now(UTC) seconds = _microseconds_from_datetime(now) / 1.0e6 @@ -650,7 +654,7 @@ def test_to_api_repr_w_timestamp_micros(self): self.assertEqual(param.to_api_repr(), EXPECTED) def test_to_api_repr_w_datetime_datetime(self): - from google.cloud._helpers import _datetime_to_rfc3339, UTC + from google.cloud._helpers import UTC, _datetime_to_rfc3339 now = datetime.datetime.now(UTC) EXPECTED = { @@ -664,7 +668,7 @@ def test_to_api_repr_w_datetime_datetime(self): self.assertEqual(param.to_api_repr(), EXPECTED) def test_to_api_repr_w_datetime_string(self): - from google.cloud._helpers import _datetime_to_rfc3339, UTC + from google.cloud._helpers import UTC, _datetime_to_rfc3339 now = datetime.datetime.now(UTC) now_str = _datetime_to_rfc3339(now) @@ -1415,8 +1419,10 @@ def test_to_api_repr_w_record_type(self): self.assertEqual(param.to_api_repr(), EXPECTED) def test_to_api_repr_w_empty_array_of_records_type(self): - from google.cloud.bigquery.query import ScalarQueryParameterType - from google.cloud.bigquery.query import StructQueryParameterType + from google.cloud.bigquery.query import ( + ScalarQueryParameterType, + StructQueryParameterType, + ) EXPECTED = { "parameterType": { @@ -1498,8 +1504,10 @@ def test___repr__array_type_scalar_type_instance(self): self.assertEqual(repr(int_items), expected) def test___repr__array_type_struct_type_instance(self): - from google.cloud.bigquery.query import ScalarQueryParameterType - from google.cloud.bigquery.query import StructQueryParameterType + from google.cloud.bigquery.query import ( + ScalarQueryParameterType, + StructQueryParameterType, + ) struct_items = self._make_one( "struct_items", diff --git a/packages/google-cloud-bigquery/tests/unit/test_retry.py b/packages/google-cloud-bigquery/tests/unit/test_retry.py index a249d1909909..ffdd95ca5e53 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_retry.py +++ b/packages/google-cloud-bigquery/tests/unit/test_retry.py @@ -111,8 +111,8 @@ def test_w_unstructured_bad_gateway(self): def test_DEFAULT_JOB_RETRY_predicate(): - from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY from google.api_core.exceptions import ClientError + from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY assert not DEFAULT_JOB_RETRY._predicate(TypeError()) assert not DEFAULT_JOB_RETRY._predicate(ClientError("fail")) @@ -137,8 +137,8 @@ def test_DEFAULT_JOB_RETRY_deadline(): def test_DEFAULT_JOB_RETRY_job_rate_limit_exceeded_retry_predicate(): """Tests the retry predicate specifically for jobRateLimitExceeded.""" - from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY from google.api_core.exceptions import ClientError + from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY # Non-ClientError exceptions should never trigger a retry assert not DEFAULT_JOB_RETRY._predicate(TypeError()) diff --git a/packages/google-cloud-bigquery/tests/unit/test_schema.py b/packages/google-cloud-bigquery/tests/unit/test_schema.py index f61b22035ed5..c3d8c045de47 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_schema.py +++ b/packages/google-cloud-bigquery/tests/unit/test_schema.py @@ -17,12 +17,10 @@ from unittest import mock import pytest - from google.cloud import bigquery -from google.cloud.bigquery import enums -from google.cloud.bigquery.standard_sql import StandardSqlStructType -from google.cloud.bigquery import schema +from google.cloud.bigquery import enums, schema from google.cloud.bigquery.schema import PolicyTagList +from google.cloud.bigquery.standard_sql import StandardSqlStructType class TestSchemaField(unittest.TestCase): diff --git a/packages/google-cloud-bigquery/tests/unit/test_standard_sql_types.py b/packages/google-cloud-bigquery/tests/unit/test_standard_sql_types.py index 3ed912b5a4a8..45bd5d80c149 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_standard_sql_types.py +++ b/packages/google-cloud-bigquery/tests/unit/test_standard_sql_types.py @@ -15,7 +15,6 @@ from unittest import mock import pytest - from google.cloud import bigquery as bq @@ -77,8 +76,10 @@ def test_to_api_repr_struct_type_field_types_missing(self): assert result == {"typeKind": "STRUCT"} def test_to_api_repr_struct_type_w_field_types(self): - from google.cloud.bigquery.standard_sql import StandardSqlField - from google.cloud.bigquery.standard_sql import StandardSqlStructType + from google.cloud.bigquery.standard_sql import ( + StandardSqlField, + StandardSqlStructType, + ) StandardSqlDataType = self._get_target_class() TypeNames = bq.StandardSqlTypeNames @@ -202,8 +203,10 @@ def test_from_api_repr_array_type_missing_element_type(self): assert result == expected def test_from_api_repr_struct_type_nested(self): - from google.cloud.bigquery.standard_sql import StandardSqlField - from google.cloud.bigquery.standard_sql import StandardSqlStructType + from google.cloud.bigquery.standard_sql import ( + StandardSqlField, + StandardSqlStructType, + ) klass = self._get_target_class() TypeNames = bq.StandardSqlTypeNames @@ -269,8 +272,10 @@ def test_from_api_repr_struct_type_missing_struct_info(self): assert result == expected def test_from_api_repr_struct_type_incomplete_field_info(self): - from google.cloud.bigquery.standard_sql import StandardSqlField - from google.cloud.bigquery.standard_sql import StandardSqlStructType + from google.cloud.bigquery.standard_sql import ( + StandardSqlField, + StandardSqlStructType, + ) klass = self._get_target_class() TypeNames = bq.StandardSqlTypeNames @@ -584,8 +589,10 @@ def test_from_api_repr_missing_columns(self): assert result.columns == [] def test_from_api_repr_with_incomplete_columns(self): - from google.cloud.bigquery.standard_sql import StandardSqlDataType - from google.cloud.bigquery.standard_sql import StandardSqlField + from google.cloud.bigquery.standard_sql import ( + StandardSqlDataType, + StandardSqlField, + ) resource = { "columns": [ 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() diff --git a/packages/google-cloud-bigquery/tests/unit/test_table_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_table_arrow.py index fdd1b7b78c60..b1798e76a5a5 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_table_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_table_arrow.py @@ -12,11 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import google.cloud.bigquery.table import pytest - from google.cloud import bigquery -import google.cloud.bigquery.table - pytest.importorskip("numpy") pytest.importorskip("pyarrow", minversion="3.0.0") diff --git a/packages/google-cloud-bigquery/tests/unit/test_table_pandas.py b/packages/google-cloud-bigquery/tests/unit/test_table_pandas.py index 94bd901c3275..1ecc2265b81a 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_table_pandas.py +++ b/packages/google-cloud-bigquery/tests/unit/test_table_pandas.py @@ -12,12 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import mock import datetime import decimal +from unittest import mock import pytest - from google.cloud import bigquery pandas = pytest.importorskip("pandas")