feat: prepare client-side delegation and telemetry for pandas-gbq

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 .*:

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

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]
 

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

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

packages/google-cloud-bigquery/google/cloud/bigquery/_string_references.py

@@ -21,7 +21,6 @@
 import re
 from typing import TypedDict, Union
 
-
 ParsedDatasetReference = TypedDict(
     "ParsedDatasetReference",
     {

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()

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,

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: