diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml
index 9ee60f7e4..a9fcd07cc 100644
--- a/.github/.OwlBot.lock.yaml
+++ b/.github/.OwlBot.lock.yaml
@@ -1,3 +1,3 @@
docker:
image: gcr.io/repo-automation-bots/owlbot-python:latest
- digest: sha256:aea14a583128771ae8aefa364e1652f3c56070168ef31beb203534222d842b8b
+ digest: sha256:9743664022bd63a8084be67f144898314c7ca12f0a03e422ac17c733c129d803
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0c08e7910..8a21df6fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,50 @@
[1]: https://pypi.org/project/google-cloud-bigquery/#history
+### [2.25.1](https://www.github.com/googleapis/python-bigquery/compare/v2.25.0...v2.25.1) (2021-08-25)
+
+
+### Bug Fixes
+
+* populate default `timeout` and retry after client-side timeout ([#896](https://www.github.com/googleapis/python-bigquery/issues/896)) ([b508809](https://www.github.com/googleapis/python-bigquery/commit/b508809c0f887575274309a463e763c56ddd017d))
+* use REST API in cell magic when requested ([#892](https://www.github.com/googleapis/python-bigquery/issues/892)) ([1cb3e55](https://www.github.com/googleapis/python-bigquery/commit/1cb3e55253e824e3a1da5201f6ec09065fb6b627))
+
+## [2.25.0](https://www.github.com/googleapis/python-bigquery/compare/v2.24.1...v2.25.0) (2021-08-24)
+
+
+### Features
+
+* Support using GeoPandas for GEOGRAPHY columns ([#848](https://www.github.com/googleapis/python-bigquery/issues/848)) ([16f65e6](https://www.github.com/googleapis/python-bigquery/commit/16f65e6ae15979217ceea6c6d398c9057a363a13))
+
+### [2.24.1](https://www.github.com/googleapis/python-bigquery/compare/v2.24.0...v2.24.1) (2021-08-13)
+
+
+### Bug Fixes
+
+* remove pytz dependency and require pyarrow>=3.0.0 ([#875](https://www.github.com/googleapis/python-bigquery/issues/875)) ([2cb3563](https://www.github.com/googleapis/python-bigquery/commit/2cb3563ee863edef7eaf5d04d739bcfe7bc6438e))
+
+## [2.24.0](https://www.github.com/googleapis/python-bigquery/compare/v2.23.3...v2.24.0) (2021-08-11)
+
+
+### Features
+
+* add support for transaction statistics ([#849](https://www.github.com/googleapis/python-bigquery/issues/849)) ([7f7b1a8](https://www.github.com/googleapis/python-bigquery/commit/7f7b1a808d50558772a0deb534ca654da65d629e))
+* make the same `Table*` instances equal to each other ([#867](https://www.github.com/googleapis/python-bigquery/issues/867)) ([c1a3d44](https://www.github.com/googleapis/python-bigquery/commit/c1a3d4435739a21d25aa154145e36d3a7c42eeb6))
+* retry failed query jobs in `result()` ([#837](https://www.github.com/googleapis/python-bigquery/issues/837)) ([519d99c](https://www.github.com/googleapis/python-bigquery/commit/519d99c20e7d1101f76981f3de036fdf3c7a4ecc))
+* support `ScalarQueryParameterType` for `type_` argument in `ScalarQueryParameter` constructor ([#850](https://www.github.com/googleapis/python-bigquery/issues/850)) ([93d15e2](https://www.github.com/googleapis/python-bigquery/commit/93d15e2e5405c2cc6d158c4e5737361344193dbc))
+
+
+### Bug Fixes
+
+* make unicode characters working well in load_table_from_json ([#865](https://www.github.com/googleapis/python-bigquery/issues/865)) ([ad9c802](https://www.github.com/googleapis/python-bigquery/commit/ad9c8026f0e667f13dd754279f9dc40d06f4fa78))
+
+### [2.23.3](https://www.github.com/googleapis/python-bigquery/compare/v2.23.2...v2.23.3) (2021-08-06)
+
+
+### Bug Fixes
+
+* increase default retry deadline to 10 minutes ([#859](https://www.github.com/googleapis/python-bigquery/issues/859)) ([30770fd](https://www.github.com/googleapis/python-bigquery/commit/30770fd0575fbd5aaa70c14196a4cc54627aecd2))
+
### [2.23.2](https://www.github.com/googleapis/python-bigquery/compare/v2.23.1...v2.23.2) (2021-07-29)
diff --git a/docs/conf.py b/docs/conf.py
index cb347160d..59a2d8fb3 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -110,6 +110,7 @@
# directories to ignore when looking for source files.
exclude_patterns = [
"_build",
+ "**/.nox/**/*",
"samples/AUTHORING_GUIDE.md",
"samples/CONTRIBUTING.md",
"samples/snippets/README.rst",
@@ -365,6 +366,8 @@
"grpc": ("https://grpc.github.io/grpc/python/", None),
"proto-plus": ("https://proto-plus-python.readthedocs.io/en/latest/", None),
"protobuf": ("https://googleapis.dev/python/protobuf/latest/", None),
+ "pandas": ("http://pandas.pydata.org/pandas-docs/dev", None),
+ "geopandas": ("https://geopandas.org/", None),
}
diff --git a/docs/reference.rst b/docs/reference.rst
index 8a5bff9a4..d8738e67b 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -68,6 +68,7 @@ Job-Related Types
job.SourceFormat
job.WriteDisposition
job.SchemaUpdateOption
+ job.TransactionInfo
Dataset
@@ -137,6 +138,7 @@ Query
query.ArrayQueryParameter
query.ScalarQueryParameter
+ query.ScalarQueryParameterType
query.StructQueryParameter
query.UDFResource
diff --git a/docs/snippets.py b/docs/snippets.py
index 3f9b9a88c..c62001fc0 100644
--- a/docs/snippets.py
+++ b/docs/snippets.py
@@ -363,7 +363,6 @@ def test_update_table_expiration(client, to_delete):
# [START bigquery_update_table_expiration]
import datetime
- import pytz
# from google.cloud import bigquery
# client = bigquery.Client()
@@ -375,7 +374,9 @@ def test_update_table_expiration(client, to_delete):
assert table.expires is None
# set table to expire 5 days from now
- expiration = datetime.datetime.now(pytz.utc) + datetime.timedelta(days=5)
+ expiration = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
+ days=5
+ )
table.expires = expiration
table = client.update_table(table, ["expires"]) # API request
diff --git a/docs/usage/pandas.rst b/docs/usage/pandas.rst
index 9db98dfbb..92eee67cf 100644
--- a/docs/usage/pandas.rst
+++ b/docs/usage/pandas.rst
@@ -37,6 +37,21 @@ To retrieve table rows as a :class:`pandas.DataFrame`:
:start-after: [START bigquery_list_rows_dataframe]
:end-before: [END bigquery_list_rows_dataframe]
+
+Retrieve BigQuery GEOGRAPHY data as a GeoPandas GeoDataFrame
+------------------------------------------------------------
+
+`GeoPandas `_ adds geospatial analytics
+capabilities to Pandas. To retrieve query results containing
+GEOGRAPHY data as a :class:`geopandas.GeoDataFrame`:
+
+.. literalinclude:: ../samples/geography/to_geodataframe.py
+ :language: python
+ :dedent: 4
+ :start-after: [START bigquery_query_results_geodataframe]
+ :end-before: [END bigquery_query_results_geodataframe]
+
+
Load a Pandas DataFrame to a BigQuery Table
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/google/cloud/bigquery/__init__.py b/google/cloud/bigquery/__init__.py
index 222aadcc9..a7a0da3dd 100644
--- a/google/cloud/bigquery/__init__.py
+++ b/google/cloud/bigquery/__init__.py
@@ -70,6 +70,7 @@
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
@@ -149,6 +150,7 @@
"GoogleSheetsOptions",
"ParquetOptions",
"ScriptOptions",
+ "TransactionInfo",
"DEFAULT_RETRY",
# Enum Constants
"enums",
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index b381fa5f7..ab58b1729 100644
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -20,12 +20,40 @@
import queue
import warnings
-from packaging import version
-
try:
import pandas
except ImportError: # pragma: NO COVER
pandas = None
+else:
+ import numpy
+
+try:
+ # _BaseGeometry is used to detect shapely objevys in `bq_to_arrow_array`
+ from shapely.geometry.base import BaseGeometry as _BaseGeometry
+except ImportError: # pragma: NO COVER
+ # No shapely, use NoneType for _BaseGeometry as a placeholder.
+ _BaseGeometry = type(None)
+else:
+ if pandas is not None: # pragma: NO COVER
+
+ def _to_wkb():
+ # Create a closure that:
+ # - Adds a not-null check. This allows the returned function to
+ # be used directly with apply, unlike `shapely.wkb.dumps`.
+ # - Avoid extra work done by `shapely.wkb.dumps` that we don't need.
+ # - Caches the WKBWriter (and write method lookup :) )
+ # - Avoids adding WKBWriter, lgeos, and notnull to the module namespace.
+ from shapely.geos import WKBWriter, lgeos
+
+ write = WKBWriter(lgeos).write
+ notnull = pandas.notnull
+
+ def _to_wkb(v):
+ return write(v) if notnull(v) else v
+
+ return _to_wkb
+
+ _to_wkb = _to_wkb()
try:
import pyarrow
@@ -71,6 +99,7 @@
"uint8": "INTEGER",
"uint16": "INTEGER",
"uint32": "INTEGER",
+ "geometry": "GEOGRAPHY",
}
@@ -110,6 +139,7 @@ def pyarrow_timestamp():
# This dictionary is duplicated in bigquery_storage/test/unite/test_reader.py
# When modifying it be sure to update it there as well.
BQ_TO_ARROW_SCALARS = {
+ "BIGNUMERIC": pyarrow_bignumeric,
"BOOL": pyarrow.bool_,
"BOOLEAN": pyarrow.bool_,
"BYTES": pyarrow.binary,
@@ -146,23 +176,15 @@ def pyarrow_timestamp():
pyarrow.date64().id: "DATETIME", # because millisecond resolution
pyarrow.binary().id: "BYTES",
pyarrow.string().id: "STRING", # also alias for pyarrow.utf8()
- # The exact scale and precision don't matter, see below.
- pyarrow.decimal128(38, scale=9).id: "NUMERIC",
- }
-
- if version.parse(pyarrow.__version__) >= version.parse("3.0.0"):
- BQ_TO_ARROW_SCALARS["BIGNUMERIC"] = pyarrow_bignumeric
# The exact decimal's scale and precision are not important, as only
# the type ID matters, and it's the same for all decimal256 instances.
- ARROW_SCALAR_IDS_TO_BQ[pyarrow.decimal256(76, scale=38).id] = "BIGNUMERIC"
- _BIGNUMERIC_SUPPORT = True
- else:
- _BIGNUMERIC_SUPPORT = False
+ pyarrow.decimal128(38, scale=9).id: "NUMERIC",
+ pyarrow.decimal256(76, scale=38).id: "BIGNUMERIC",
+ }
else: # pragma: NO COVER
BQ_TO_ARROW_SCALARS = {} # pragma: NO COVER
ARROW_SCALAR_IDS_TO_BQ = {} # pragma: NO_COVER
- _BIGNUMERIC_SUPPORT = False # pragma: NO COVER
def bq_to_arrow_struct_data_type(field):
@@ -202,14 +224,16 @@ def bq_to_arrow_data_type(field):
return data_type_constructor()
-def bq_to_arrow_field(bq_field):
+def bq_to_arrow_field(bq_field, array_type=None):
"""Return the Arrow field, corresponding to a given BigQuery column.
Returns:
None: if the Arrow type cannot be determined.
"""
arrow_type = bq_to_arrow_data_type(bq_field)
- if arrow_type:
+ if arrow_type is not None:
+ if array_type is not None:
+ arrow_type = array_type # For GEOGRAPHY, at least initially
is_nullable = bq_field.mode.upper() == "NULLABLE"
return pyarrow.field(bq_field.name, arrow_type, nullable=is_nullable)
@@ -234,7 +258,24 @@ def bq_to_arrow_schema(bq_schema):
def bq_to_arrow_array(series, bq_field):
- arrow_type = bq_to_arrow_data_type(bq_field)
+ if bq_field.field_type.upper() == "GEOGRAPHY":
+ arrow_type = None
+ first = _first_valid(series)
+ if first is not None:
+ if series.dtype.name == "geometry" or isinstance(first, _BaseGeometry):
+ arrow_type = pyarrow.binary()
+ # Convert shapey geometry to WKB binary format:
+ series = series.apply(_to_wkb)
+ elif isinstance(first, bytes):
+ arrow_type = pyarrow.binary()
+ elif series.dtype.name == "geometry":
+ # We have a GeoSeries containing all nulls, convert it to a pandas series
+ series = pandas.Series(numpy.array(series))
+
+ if arrow_type is None:
+ arrow_type = bq_to_arrow_data_type(bq_field)
+ else:
+ arrow_type = bq_to_arrow_data_type(bq_field)
field_type_upper = bq_field.field_type.upper() if bq_field.field_type else ""
@@ -288,6 +329,12 @@ def list_columns_and_indexes(dataframe):
return columns_and_indexes
+def _first_valid(series):
+ first_valid_index = series.first_valid_index()
+ if first_valid_index is not None:
+ return series.at[first_valid_index]
+
+
def dataframe_to_bq_schema(dataframe, bq_schema):
"""Convert a pandas DataFrame schema to a BigQuery schema.
@@ -328,6 +375,13 @@ def dataframe_to_bq_schema(dataframe, bq_schema):
# Otherwise, try to automatically determine the type based on the
# pandas dtype.
bq_type = _PANDAS_DTYPE_TO_BQ.get(dtype.name)
+ if bq_type is None:
+ sample_data = _first_valid(dataframe[column])
+ if (
+ isinstance(sample_data, _BaseGeometry)
+ and sample_data is not None # Paranoia
+ ):
+ bq_type = "GEOGRAPHY"
bq_field = schema.SchemaField(column, bq_type)
bq_schema_out.append(bq_field)
@@ -459,11 +513,11 @@ def dataframe_to_arrow(dataframe, bq_schema):
arrow_names = []
arrow_fields = []
for bq_field in bq_schema:
- arrow_fields.append(bq_to_arrow_field(bq_field))
arrow_names.append(bq_field.name)
arrow_arrays.append(
bq_to_arrow_array(get_column_or_index(dataframe, bq_field.name), bq_field)
)
+ arrow_fields.append(bq_to_arrow_field(bq_field, arrow_arrays[-1].type))
if all((field is not None for field in arrow_fields)):
return pyarrow.Table.from_arrays(
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
index 742ecac2e..023346ffa 100644
--- a/google/cloud/bigquery/client.py
+++ b/google/cloud/bigquery/client.py
@@ -76,17 +76,24 @@
from google.cloud.bigquery.opentelemetry_tracing import create_span
from google.cloud.bigquery import job
from google.cloud.bigquery.job import (
+ CopyJob,
+ CopyJobConfig,
+ ExtractJob,
+ ExtractJobConfig,
+ LoadJob,
LoadJobConfig,
QueryJob,
QueryJobConfig,
- CopyJobConfig,
- ExtractJobConfig,
)
from google.cloud.bigquery.model import Model
from google.cloud.bigquery.model import ModelReference
from google.cloud.bigquery.model import _model_arg_to_model_ref
from google.cloud.bigquery.query import _QueryResults
-from google.cloud.bigquery.retry import DEFAULT_RETRY
+from google.cloud.bigquery.retry import (
+ DEFAULT_JOB_RETRY,
+ DEFAULT_RETRY,
+ DEFAULT_TIMEOUT,
+)
from google.cloud.bigquery.routine import Routine
from google.cloud.bigquery.routine import RoutineReference
from google.cloud.bigquery.schema import SchemaField
@@ -245,7 +252,7 @@ def get_service_account_email(
self,
project: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> str:
"""Get the email address of the project's BigQuery service account
@@ -292,7 +299,7 @@ def list_projects(
max_results: int = None,
page_token: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
page_size: int = None,
) -> page_iterator.Iterator:
"""List projects for the project associated with this client.
@@ -358,7 +365,7 @@ def list_datasets(
max_results: int = None,
page_token: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
page_size: int = None,
) -> page_iterator.Iterator:
"""List datasets for the project associated with this client.
@@ -549,7 +556,7 @@ def create_dataset(
dataset: Union[str, Dataset, DatasetReference],
exists_ok: bool = False,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Dataset:
"""API call: create the dataset via a POST request.
@@ -624,7 +631,7 @@ def create_routine(
routine: Routine,
exists_ok: bool = False,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Routine:
"""[Beta] Create a routine via a POST request.
@@ -679,7 +686,7 @@ def create_table(
table: Union[str, Table, TableReference],
exists_ok: bool = False,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Table:
"""API call: create a table via a PUT request
@@ -751,7 +758,7 @@ def get_dataset(
self,
dataset_ref: Union[DatasetReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Dataset:
"""Fetch the dataset referenced by ``dataset_ref``
@@ -795,7 +802,7 @@ def get_iam_policy(
table: Union[Table, TableReference],
requested_policy_version: int = 1,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Policy:
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")
@@ -825,7 +832,7 @@ def set_iam_policy(
policy: Policy,
updateMask: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Policy:
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")
@@ -858,7 +865,7 @@ def test_iam_permissions(
table: Union[Table, TableReference],
permissions: Sequence[str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Dict[str, Any]:
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")
@@ -883,7 +890,7 @@ def get_model(
self,
model_ref: Union[ModelReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Model:
"""[Beta] Fetch the model referenced by ``model_ref``.
@@ -926,7 +933,7 @@ def get_routine(
self,
routine_ref: Union[Routine, RoutineReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Routine:
"""[Beta] Get the routine referenced by ``routine_ref``.
@@ -970,7 +977,7 @@ def get_table(
self,
table: Union[Table, TableReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Table:
"""Fetch the table referenced by ``table``.
@@ -1012,7 +1019,7 @@ def update_dataset(
dataset: Dataset,
fields: Sequence[str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Dataset:
"""Change some fields of a dataset.
@@ -1082,7 +1089,7 @@ def update_model(
model: Model,
fields: Sequence[str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Model:
"""[Beta] Change some fields of a model.
@@ -1146,7 +1153,7 @@ def update_routine(
routine: Routine,
fields: Sequence[str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Routine:
"""[Beta] Change some fields of a routine.
@@ -1220,7 +1227,7 @@ def update_table(
table: Table,
fields: Sequence[str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Table:
"""Change some fields of a table.
@@ -1286,7 +1293,7 @@ def list_models(
max_results: int = None,
page_token: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
page_size: int = None,
) -> page_iterator.Iterator:
"""[Beta] List models in the dataset.
@@ -1363,7 +1370,7 @@ def list_routines(
max_results: int = None,
page_token: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
page_size: int = None,
) -> page_iterator.Iterator:
"""[Beta] List routines in the dataset.
@@ -1440,7 +1447,7 @@ def list_tables(
max_results: int = None,
page_token: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
page_size: int = None,
) -> page_iterator.Iterator:
"""List tables in the dataset.
@@ -1515,7 +1522,7 @@ def delete_dataset(
dataset: Union[Dataset, DatasetReference, str],
delete_contents: bool = False,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
not_found_ok: bool = False,
) -> None:
"""Delete a dataset.
@@ -1574,7 +1581,7 @@ def delete_model(
self,
model: Union[Model, ModelReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
not_found_ok: bool = False,
) -> None:
"""[Beta] Delete a model
@@ -1624,12 +1631,12 @@ def delete_model(
def delete_job_metadata(
self,
- job_id,
- project=None,
- location=None,
- retry=DEFAULT_RETRY,
- timeout=None,
- not_found_ok=False,
+ job_id: Union[str, LoadJob, CopyJob, ExtractJob, QueryJob],
+ project: Optional[str] = None,
+ location: Optional[str] = None,
+ retry: retries.Retry = DEFAULT_RETRY,
+ timeout: float = DEFAULT_TIMEOUT,
+ not_found_ok: bool = False,
):
"""[Beta] Delete job metadata from job history.
@@ -1637,26 +1644,20 @@ def delete_job_metadata(
:func:`~google.cloud.bigquery.client.Client.cancel_job` instead.
Args:
- job_id (Union[ \
- str, \
- google.cloud.bigquery.job.LoadJob, \
- google.cloud.bigquery.job.CopyJob, \
- google.cloud.bigquery.job.ExtractJob, \
- google.cloud.bigquery.job.QueryJob \
- ]): Job identifier.
+ job_id: Job or job identifier.
Keyword Arguments:
- project (Optional[str]):
+ project:
ID of the project which owns the job (defaults to the client's project).
- location (Optional[str]):
+ location:
Location where the job was run. Ignored if ``job_id`` is a job
object.
- retry (Optional[google.api_core.retry.Retry]):
+ retry:
How to retry the RPC.
- timeout (Optional[float]):
+ timeout:
The number of seconds to wait for the underlying HTTP transport
before using ``retry``.
- not_found_ok (Optional[bool]):
+ not_found_ok:
Defaults to ``False``. If ``True``, ignore "not found" errors
when deleting the job.
"""
@@ -1697,7 +1698,7 @@ def delete_routine(
self,
routine: Union[Routine, RoutineReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
not_found_ok: bool = False,
) -> None:
"""[Beta] Delete a routine.
@@ -1751,7 +1752,7 @@ def delete_table(
self,
table: Union[Table, TableReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
not_found_ok: bool = False,
) -> None:
"""Delete a table
@@ -1804,7 +1805,7 @@ def _get_query_results(
project: str = None,
timeout_ms: int = None,
location: str = None,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> _QueryResults:
"""Get the query results object for a query job.
@@ -1893,7 +1894,7 @@ def create_job(
self,
job_config: dict,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Union[job.LoadJob, job.CopyJob, job.ExtractJob, job.QueryJob]:
"""Create a new job.
Args:
@@ -1990,7 +1991,7 @@ def get_job(
project: str = None,
location: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Union[job.LoadJob, job.CopyJob, job.ExtractJob, job.QueryJob]:
"""Fetch a job for the project associated with this client.
@@ -2064,7 +2065,7 @@ def cancel_job(
project: str = None,
location: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Union[job.LoadJob, job.CopyJob, job.ExtractJob, job.QueryJob]:
"""Attempt to cancel a job from a job ID.
@@ -2141,7 +2142,7 @@ def list_jobs(
all_users: bool = None,
state_filter: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
min_creation_time: datetime.datetime = None,
max_creation_time: datetime.datetime = None,
page_size: int = None,
@@ -2256,7 +2257,7 @@ def load_table_from_uri(
project: str = None,
job_config: LoadJobConfig = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> job.LoadJob:
"""Starts a job for loading data into a table from CloudStorage.
@@ -2340,7 +2341,7 @@ def load_table_from_file(
location: str = None,
project: str = None,
job_config: LoadJobConfig = None,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> job.LoadJob:
"""Upload the contents of this table from a file-like object.
@@ -2443,7 +2444,7 @@ def load_table_from_dataframe(
project: str = None,
job_config: LoadJobConfig = None,
parquet_compression: str = "snappy",
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> job.LoadJob:
"""Upload the contents of a table from a pandas DataFrame.
@@ -2678,7 +2679,7 @@ def load_table_from_json(
location: str = None,
project: str = None,
job_config: LoadJobConfig = None,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> job.LoadJob:
"""Upload the contents of a table from a JSON string or dict.
@@ -2762,7 +2763,7 @@ def load_table_from_json(
destination = _table_arg_to_table_ref(destination, default_project=self.project)
- data_str = "\n".join(json.dumps(item) for item in json_rows)
+ data_str = "\n".join(json.dumps(item, ensure_ascii=False) for item in json_rows)
encoded_str = data_str.encode()
data_file = io.BytesIO(encoded_str)
return self.load_table_from_file(
@@ -2961,7 +2962,7 @@ def copy_table(
project: str = None,
job_config: CopyJobConfig = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> job.CopyJob:
"""Copy one or more tables to another table.
@@ -3064,7 +3065,7 @@ def extract_table(
project: str = None,
job_config: ExtractJobConfig = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
source_type: str = "Table",
) -> job.ExtractJob:
"""Start a job to extract a table into Cloud Storage files.
@@ -3162,7 +3163,8 @@ def query(
location: str = None,
project: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
+ job_retry: retries.Retry = DEFAULT_JOB_RETRY,
) -> job.QueryJob:
"""Run a SQL query.
@@ -3192,21 +3194,52 @@ def query(
Project ID of the project of where to run the job. Defaults
to the client's project.
retry (Optional[google.api_core.retry.Retry]):
- How to retry the RPC.
+ How to retry the RPC. This only applies to making RPC
+ calls. It isn't used to retry failed jobs. This has
+ a reasonable default that should only be overridden
+ with care.
timeout (Optional[float]):
The number of seconds to wait for the underlying HTTP transport
before using ``retry``.
+ job_retry (Optional[google.api_core.retry.Retry]):
+ How to retry failed jobs. The default retries
+ rate-limit-exceeded errors. Passing ``None`` disables
+ job retry.
+
+ Not all jobs can be retried. If ``job_id`` is
+ provided, then the job returned by the query will not
+ be retryable, and an exception will be raised if a
+ non-``None`` (and non-default) value for ``job_retry``
+ is also provided.
+
+ Note that errors aren't detected until ``result()`` is
+ called on the job returned. The ``job_retry``
+ specified here becomes the default ``job_retry`` for
+ ``result()``, where it can also be specified.
Returns:
google.cloud.bigquery.job.QueryJob: A new query job instance.
Raises:
TypeError:
- If ``job_config`` is not an instance of :class:`~google.cloud.bigquery.job.QueryJobConfig`
- class.
+ If ``job_config`` is not an instance of
+ :class:`~google.cloud.bigquery.job.QueryJobConfig`
+ class, or if both ``job_id`` and non-``None`` non-default
+ ``job_retry`` are provided.
"""
job_id_given = job_id is not None
- job_id = _make_job_id(job_id, job_id_prefix)
+ if (
+ job_id_given
+ and job_retry is not None
+ and job_retry is not DEFAULT_JOB_RETRY
+ ):
+ raise TypeError(
+ "`job_retry` was provided, but the returned job is"
+ " not retryable, because a custom `job_id` was"
+ " provided."
+ )
+
+ job_id_save = job_id
if project is None:
project = self.project
@@ -3214,8 +3247,6 @@ def query(
if location is None:
location = self.location
- job_config = copy.deepcopy(job_config)
-
if self._default_query_job_config:
if job_config:
_verify_job_config_type(
@@ -3225,6 +3256,8 @@ def query(
# that is in the default,
# should be filled in with the default
# the incoming therefore has precedence
+ #
+ # Note that _fill_from_default doesn't mutate the receiver
job_config = job_config._fill_from_default(
self._default_query_job_config
)
@@ -3233,34 +3266,54 @@ def query(
self._default_query_job_config,
google.cloud.bigquery.job.QueryJobConfig,
)
- job_config = copy.deepcopy(self._default_query_job_config)
+ job_config = self._default_query_job_config
- job_ref = job._JobReference(job_id, project=project, location=location)
- query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config)
+ # Note that we haven't modified the original job_config (or
+ # _default_query_job_config) up to this point.
+ job_config_save = job_config
- try:
- query_job._begin(retry=retry, timeout=timeout)
- except core_exceptions.Conflict as create_exc:
- # The thought is if someone is providing their own job IDs and they get
- # their job ID generation wrong, this could end up returning results for
- # the wrong query. We thus only try to recover if job ID was not given.
- if job_id_given:
- raise create_exc
+ def do_query():
+ # Make a copy now, so that original doesn't get changed by the process
+ # below and to facilitate retry
+ job_config = copy.deepcopy(job_config_save)
+
+ job_id = _make_job_id(job_id_save, job_id_prefix)
+ job_ref = job._JobReference(job_id, project=project, location=location)
+ query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config)
try:
- query_job = self.get_job(
- job_id,
- project=project,
- location=location,
- retry=retry,
- timeout=timeout,
- )
- except core_exceptions.GoogleAPIError: # (includes RetryError)
- raise create_exc
+ query_job._begin(retry=retry, timeout=timeout)
+ except core_exceptions.Conflict as create_exc:
+ # The thought is if someone is providing their own job IDs and they get
+ # their job ID generation wrong, this could end up returning results for
+ # the wrong query. We thus only try to recover if job ID was not given.
+ if job_id_given:
+ raise create_exc
+
+ try:
+ query_job = self.get_job(
+ job_id,
+ project=project,
+ location=location,
+ retry=retry,
+ timeout=timeout,
+ )
+ except core_exceptions.GoogleAPIError: # (includes RetryError)
+ raise create_exc
+ else:
+ return query_job
else:
return query_job
- else:
- return query_job
+
+ future = do_query()
+ # The future might be in a failed state now, but if it's
+ # unrecoverable, we'll find out when we ask for it's result, at which
+ # point, we may retry.
+ if not job_id_given:
+ future._retry_do_query = do_query # in case we have to retry later
+ future._job_retry = job_retry
+
+ return future
def insert_rows(
self,
@@ -3392,7 +3445,7 @@ def insert_rows_json(
ignore_unknown_values: bool = None,
template_suffix: str = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Sequence[dict]:
"""Insert rows into a table without applying local type conversions.
@@ -3527,7 +3580,7 @@ def list_partitions(
self,
table: Union[Table, TableReference, str],
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> Sequence[str]:
"""List the partitions in a table.
@@ -3577,7 +3630,7 @@ def list_rows(
start_index: int = None,
page_size: int = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> RowIterator:
"""List the rows of the table.
@@ -3689,7 +3742,7 @@ def _list_rows_from_query_results(
start_index: int = None,
page_size: int = None,
retry: retries.Retry = DEFAULT_RETRY,
- timeout: float = None,
+ timeout: float = DEFAULT_TIMEOUT,
) -> RowIterator:
"""List the rows of a completed query.
See
diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py
index 0da01d665..d67cebd4c 100644
--- a/google/cloud/bigquery/enums.py
+++ b/google/cloud/bigquery/enums.py
@@ -259,23 +259,23 @@ class SqlTypeNames(str, enum.Enum):
class SqlParameterScalarTypes:
"""Supported scalar SQL query parameter types as type objects."""
- STRING = ScalarQueryParameterType("STRING")
+ BOOL = ScalarQueryParameterType("BOOL")
+ BOOLEAN = ScalarQueryParameterType("BOOL")
+ BIGDECIMAL = ScalarQueryParameterType("BIGNUMERIC")
+ BIGNUMERIC = ScalarQueryParameterType("BIGNUMERIC")
BYTES = ScalarQueryParameterType("BYTES")
- INTEGER = ScalarQueryParameterType("INT64")
- INT64 = ScalarQueryParameterType("INT64")
+ DATE = ScalarQueryParameterType("DATE")
+ DATETIME = ScalarQueryParameterType("DATETIME")
+ DECIMAL = ScalarQueryParameterType("NUMERIC")
FLOAT = ScalarQueryParameterType("FLOAT64")
FLOAT64 = ScalarQueryParameterType("FLOAT64")
- NUMERIC = ScalarQueryParameterType("NUMERIC")
- BIGNUMERIC = ScalarQueryParameterType("BIGNUMERIC")
- DECIMAL = ScalarQueryParameterType("NUMERIC")
- BIGDECIMAL = ScalarQueryParameterType("BIGNUMERIC")
- BOOLEAN = ScalarQueryParameterType("BOOL")
- BOOL = ScalarQueryParameterType("BOOL")
GEOGRAPHY = ScalarQueryParameterType("GEOGRAPHY")
- TIMESTAMP = ScalarQueryParameterType("TIMESTAMP")
- DATE = ScalarQueryParameterType("DATE")
+ INT64 = ScalarQueryParameterType("INT64")
+ INTEGER = ScalarQueryParameterType("INT64")
+ NUMERIC = ScalarQueryParameterType("NUMERIC")
+ STRING = ScalarQueryParameterType("STRING")
TIME = ScalarQueryParameterType("TIME")
- DATETIME = ScalarQueryParameterType("DATETIME")
+ TIMESTAMP = ScalarQueryParameterType("TIMESTAMP")
class WriteDisposition(object):
diff --git a/google/cloud/bigquery/job/__init__.py b/google/cloud/bigquery/job/__init__.py
index 4c16d0e20..f51311b0b 100644
--- a/google/cloud/bigquery/job/__init__.py
+++ b/google/cloud/bigquery/job/__init__.py
@@ -22,6 +22,7 @@
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
@@ -81,5 +82,6 @@
"QueryPriority",
"SchemaUpdateOption",
"SourceFormat",
+ "TransactionInfo",
"WriteDisposition",
]
diff --git a/google/cloud/bigquery/job/base.py b/google/cloud/bigquery/job/base.py
index 20ad81c0b..e5fc592a6 100644
--- a/google/cloud/bigquery/job/base.py
+++ b/google/cloud/bigquery/job/base.py
@@ -19,6 +19,7 @@
import http
import threading
import typing
+from typing import Dict, Optional
from google.api_core import exceptions
import google.api_core.future.polling
@@ -88,6 +89,22 @@ def _error_result_to_exception(error_result):
)
+class TransactionInfo(typing.NamedTuple):
+ """[Alpha] Information of a multi-statement transaction.
+
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#TransactionInfo
+
+ .. versionadded:: 2.24.0
+ """
+
+ transaction_id: str
+ """Output only. ID of the transaction."""
+
+ @classmethod
+ def from_api_repr(cls, transaction_info: Dict[str, str]) -> "TransactionInfo":
+ return cls(transaction_info["transactionId"])
+
+
class _JobReference(object):
"""A reference to a job.
@@ -336,6 +353,18 @@ def reservation_usage(self):
for usage in usage_stats_raw
]
+ @property
+ def transaction_info(self) -> Optional[TransactionInfo]:
+ """Information of the multi-statement transaction if this job is part of one.
+
+ .. versionadded:: 2.24.0
+ """
+ info = self._properties.get("statistics", {}).get("transactionInfo")
+ if info is None:
+ return None
+ else:
+ return TransactionInfo.from_api_repr(info)
+
@property
def error_result(self):
"""Error information about the job as a whole.
diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
index 2cb7ee28e..0cb4798be 100644
--- a/google/cloud/bigquery/job/query.py
+++ b/google/cloud/bigquery/job/query.py
@@ -36,7 +36,7 @@
from google.cloud.bigquery.query import ScalarQueryParameter
from google.cloud.bigquery.query import StructQueryParameter
from google.cloud.bigquery.query import UDFResource
-from google.cloud.bigquery.retry import DEFAULT_RETRY
+from google.cloud.bigquery.retry import DEFAULT_RETRY, DEFAULT_JOB_RETRY
from google.cloud.bigquery.routine import RoutineReference
from google.cloud.bigquery.table import _EmptyRowIterator
from google.cloud.bigquery.table import RangePartitioning
@@ -53,6 +53,7 @@
# Assumption: type checks are only used by library developers and CI environments
# that have all optional dependencies installed, thus no conditional imports.
import pandas
+ import geopandas
import pyarrow
from google.api_core import retry as retries
from google.cloud import bigquery_storage
@@ -1260,6 +1261,7 @@ def result(
retry: "retries.Retry" = DEFAULT_RETRY,
timeout: float = None,
start_index: int = None,
+ job_retry: "retries.Retry" = DEFAULT_JOB_RETRY,
) -> Union["RowIterator", _EmptyRowIterator]:
"""Start the job and wait for it to complete and get the result.
@@ -1270,9 +1272,13 @@ def result(
max_results (Optional[int]):
The maximum total number of rows from this request.
retry (Optional[google.api_core.retry.Retry]):
- How to retry the call that retrieves rows. If the job state is
- ``DONE``, retrying is aborted early even if the results are not
- available, as this will not change anymore.
+ How to retry the call that retrieves rows. This only
+ applies to making RPC calls. It isn't used to retry
+ failed jobs. This has a reasonable default that
+ should only be overridden with care. If the job state
+ is ``DONE``, retrying is aborted early even if the
+ results are not available, as this will not change
+ anymore.
timeout (Optional[float]):
The number of seconds to wait for the underlying HTTP transport
before using ``retry``.
@@ -1280,6 +1286,16 @@ def result(
applies to each individual request.
start_index (Optional[int]):
The zero-based index of the starting row to read.
+ job_retry (Optional[google.api_core.retry.Retry]):
+ How to retry failed jobs. The default retries
+ rate-limit-exceeded errors. Passing ``None`` disables
+ job retry.
+
+ Not all jobs can be retried. If ``job_id`` was
+ provided to the query that created this job, then the
+ job returned by the query will not be retryable, and
+ an exception will be raised if non-``None``
+ non-default ``job_retry`` is also provided.
Returns:
google.cloud.bigquery.table.RowIterator:
@@ -1295,17 +1311,66 @@ def result(
Raises:
google.cloud.exceptions.GoogleAPICallError:
- If the job failed.
+ If the job failed and retries aren't successful.
concurrent.futures.TimeoutError:
If the job did not complete in the given timeout.
+ TypeError:
+ If Non-``None`` and non-default ``job_retry`` is
+ provided and the job is not retryable.
"""
try:
- super(QueryJob, self).result(retry=retry, timeout=timeout)
+ retry_do_query = getattr(self, "_retry_do_query", None)
+ if retry_do_query is not None:
+ if job_retry is DEFAULT_JOB_RETRY:
+ job_retry = self._job_retry
+ else:
+ if job_retry is not None and job_retry is not DEFAULT_JOB_RETRY:
+ raise TypeError(
+ "`job_retry` was provided, but this job is"
+ " not retryable, because a custom `job_id` was"
+ " provided to the query that created this job."
+ )
+
+ first = True
+
+ def do_get_result():
+ nonlocal first
+
+ if first:
+ first = False
+ else:
+ # Note that we won't get here if retry_do_query is
+ # None, because we won't use a retry.
+
+ # The orinal job is failed. Create a new one.
+ job = retry_do_query()
+
+ # If it's already failed, we might as well stop:
+ if job.done() and job.exception() is not None:
+ raise job.exception()
+
+ # Become the new job:
+ self.__dict__.clear()
+ self.__dict__.update(job.__dict__)
+
+ # This shouldn't be necessary, because once we have a good
+ # job, it should stay good,and we shouldn't have to retry.
+ # But let's be paranoid. :)
+ self._retry_do_query = retry_do_query
+ self._job_retry = job_retry
+
+ super(QueryJob, self).result(retry=retry, timeout=timeout)
+
+ # Since the job could already be "done" (e.g. got a finished job
+ # via client.get_job), the superclass call to done() might not
+ # set the self._query_results cache.
+ self._reload_query_results(retry=retry, timeout=timeout)
+
+ if retry_do_query is not None and job_retry is not None:
+ do_get_result = job_retry(do_get_result)
+
+ do_get_result()
- # Since the job could already be "done" (e.g. got a finished job
- # via client.get_job), the superclass call to done() might not
- # set the self._query_results cache.
- self._reload_query_results(retry=retry, timeout=timeout)
except exceptions.GoogleAPICallError as exc:
exc.message += self._format_for_exception(self.query, self.job_id)
exc.query_job = self
@@ -1423,6 +1488,7 @@ def to_dataframe(
create_bqstorage_client: bool = True,
date_as_object: bool = True,
max_results: Optional[int] = None,
+ geography_as_object: bool = False,
) -> "pandas.DataFrame":
"""Return a pandas DataFrame from a QueryJob
@@ -1474,13 +1540,27 @@ def to_dataframe(
.. versionadded:: 2.21.0
+ geography_as_object (Optional[bool]):
+ If ``True``, convert GEOGRAPHY data to :mod:`shapely`
+ geometry objects. If ``False`` (default), don't cast
+ geography data to :mod:`shapely` geometry objects.
+
+ .. versionadded:: 2.24.0
+
Returns:
- A :class:`~pandas.DataFrame` populated with row data and column
- headers from the query results. The column headers are derived
- from the destination table's schema.
+ pandas.DataFrame:
+ A :class:`~pandas.DataFrame` populated with row data
+ and column headers from the query results. The column
+ headers are derived from the destination table's
+ schema.
Raises:
- ValueError: If the `pandas` library cannot be imported.
+ ValueError:
+ If the :mod:`pandas` library cannot be imported, or
+ the :mod:`google.cloud.bigquery_storage_v1` module is
+ required but cannot be imported. Also if
+ `geography_as_object` is `True`, but the
+ :mod:`shapely` library cannot be imported.
"""
query_result = wait_for_query(self, progress_bar_type, max_results=max_results)
return query_result.to_dataframe(
@@ -1489,6 +1569,101 @@ def to_dataframe(
progress_bar_type=progress_bar_type,
create_bqstorage_client=create_bqstorage_client,
date_as_object=date_as_object,
+ geography_as_object=geography_as_object,
+ )
+
+ # If changing the signature of this method, make sure to apply the same
+ # changes to table.RowIterator.to_dataframe(), except for the max_results parameter
+ # that should only exist here in the QueryJob method.
+ def to_geodataframe(
+ self,
+ bqstorage_client: "bigquery_storage.BigQueryReadClient" = None,
+ dtypes: Dict[str, Any] = None,
+ progress_bar_type: str = None,
+ create_bqstorage_client: bool = True,
+ date_as_object: bool = True,
+ max_results: Optional[int] = None,
+ geography_column: Optional[str] = None,
+ ) -> "geopandas.GeoDataFrame":
+ """Return a GeoPandas GeoDataFrame from a QueryJob
+
+ Args:
+ bqstorage_client (Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient]):
+ A BigQuery Storage API client. If supplied, use the faster
+ BigQuery Storage API to fetch rows from BigQuery. This
+ API is a billable API.
+
+ This method requires the ``fastavro`` and
+ ``google-cloud-bigquery-storage`` libraries.
+
+ Reading from a specific partition or snapshot is not
+ currently supported by this method.
+
+ dtypes (Optional[Map[str, Union[str, pandas.Series.dtype]]]):
+ A dictionary of column names pandas ``dtype``s. The provided
+ ``dtype`` is used when constructing the series for the column
+ specified. Otherwise, the default pandas behavior is used.
+
+ progress_bar_type (Optional[str]):
+ If set, use the `tqdm `_ library to
+ display a progress bar while the data downloads. Install the
+ ``tqdm`` package to use this feature.
+
+ See
+ :func:`~google.cloud.bigquery.table.RowIterator.to_dataframe`
+ for details.
+
+ .. versionadded:: 1.11.0
+ create_bqstorage_client (Optional[bool]):
+ If ``True`` (default), create a BigQuery Storage API client
+ using the default API settings. The BigQuery Storage API
+ is a faster way to fetch rows from BigQuery. See the
+ ``bqstorage_client`` parameter for more information.
+
+ This argument does nothing if ``bqstorage_client`` is supplied.
+
+ .. versionadded:: 1.24.0
+
+ date_as_object (Optional[bool]):
+ If ``True`` (default), cast dates to objects. If ``False``, convert
+ to datetime64[ns] dtype.
+
+ .. versionadded:: 1.26.0
+
+ max_results (Optional[int]):
+ Maximum number of rows to include in the result. No limit by default.
+
+ .. versionadded:: 2.21.0
+
+ geography_column (Optional[str]):
+ If there are more than one GEOGRAPHY column,
+ identifies which one to use to construct a GeoPandas
+ GeoDataFrame. This option can be ommitted if there's
+ only one GEOGRAPHY column.
+
+ Returns:
+ geopandas.GeoDataFrame:
+ A :class:`geopandas.GeoDataFrame` populated with row
+ data and column headers from the query results. The
+ column headers are derived from the destination
+ table's schema.
+
+ Raises:
+ ValueError:
+ If the :mod:`geopandas` library cannot be imported, or the
+ :mod:`google.cloud.bigquery_storage_v1` module is
+ required but cannot be imported.
+
+ .. versionadded:: 2.24.0
+ """
+ query_result = wait_for_query(self, progress_bar_type, max_results=max_results)
+ return query_result.to_geodataframe(
+ bqstorage_client=bqstorage_client,
+ dtypes=dtypes,
+ progress_bar_type=progress_bar_type,
+ create_bqstorage_client=create_bqstorage_client,
+ date_as_object=date_as_object,
+ geography_column=geography_column,
)
def __iter__(self):
diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py
index 2b8c2928e..d368bbeaa 100644
--- a/google/cloud/bigquery/magics/magics.py
+++ b/google/cloud/bigquery/magics/magics.py
@@ -671,7 +671,9 @@ def _cell_magic(line, query):
_handle_error(ex, args.destination_var)
return
- result = rows.to_dataframe(bqstorage_client=bqstorage_client)
+ result = rows.to_dataframe(
+ bqstorage_client=bqstorage_client, create_bqstorage_client=False,
+ )
if args.destination_var:
IPython.get_ipython().push({args.destination_var: result})
return
@@ -728,11 +730,15 @@ def _cell_magic(line, query):
if max_results:
result = query_job.result(max_results=max_results).to_dataframe(
- bqstorage_client=bqstorage_client, progress_bar_type=progress_bar
+ bqstorage_client=None,
+ create_bqstorage_client=False,
+ progress_bar_type=progress_bar,
)
else:
result = query_job.to_dataframe(
- bqstorage_client=bqstorage_client, progress_bar_type=progress_bar
+ bqstorage_client=bqstorage_client,
+ create_bqstorage_client=False,
+ progress_bar_type=progress_bar,
)
if args.destination_var:
diff --git a/google/cloud/bigquery/query.py b/google/cloud/bigquery/query.py
index d1e9a45a5..1f449f189 100644
--- a/google/cloud/bigquery/query.py
+++ b/google/cloud/bigquery/query.py
@@ -16,7 +16,9 @@
from collections import OrderedDict
import copy
-from typing import Union
+import datetime
+import decimal
+from typing import Optional, Union
from google.cloud.bigquery.table import _parse_schema_resource
from google.cloud.bigquery._helpers import _rows_from_json
@@ -24,6 +26,11 @@
from google.cloud.bigquery._helpers import _SCALAR_VALUE_TO_JSON_PARAM
+_SCALAR_VALUE_TYPE = Optional[
+ Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]
+]
+
+
class UDFResource(object):
"""Describe a single user-defined function (UDF) resource.
@@ -325,35 +332,46 @@ class ScalarQueryParameter(_AbstractQueryParameter):
"""Named / positional query parameters for scalar values.
Args:
- name (Optional[str]):
+ name:
Parameter name, used via ``@foo`` syntax. If None, the
parameter can only be addressed via position (``?``).
- type_ (str):
- Name of parameter type. One of 'STRING', 'INT64',
- 'FLOAT64', 'NUMERIC', 'BIGNUMERIC', 'BOOL', 'TIMESTAMP', 'DATETIME', or
- 'DATE'.
+ type_:
+ Name of parameter type. See
+ :class:`google.cloud.bigquery.enums.SqlTypeNames` and
+ :class:`google.cloud.bigquery.enums.SqlParameterScalarTypes` for
+ supported types.
- value (Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]):
+ value:
The scalar parameter value.
"""
- def __init__(self, name, type_, value):
+ def __init__(
+ self,
+ name: Optional[str],
+ type_: Optional[Union[str, ScalarQueryParameterType]],
+ value: _SCALAR_VALUE_TYPE,
+ ):
self.name = name
- self.type_ = type_
+ if isinstance(type_, ScalarQueryParameterType):
+ self.type_ = type_._type
+ else:
+ self.type_ = type_
self.value = value
@classmethod
- def positional(cls, type_: str, value) -> "ScalarQueryParameter":
+ def positional(
+ cls, type_: Union[str, ScalarQueryParameterType], value: _SCALAR_VALUE_TYPE
+ ) -> "ScalarQueryParameter":
"""Factory for positional paramater.
Args:
- type_ (str):
+ type_:
Name of parameter type. One of 'STRING', 'INT64',
'FLOAT64', 'NUMERIC', 'BIGNUMERIC', 'BOOL', 'TIMESTAMP', 'DATETIME', or
'DATE'.
- value (Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]):
+ value:
The scalar parameter value.
Returns:
diff --git a/google/cloud/bigquery/retry.py b/google/cloud/bigquery/retry.py
index bab28aacb..830582322 100644
--- a/google/cloud/bigquery/retry.py
+++ b/google/cloud/bigquery/retry.py
@@ -29,9 +29,12 @@
exceptions.BadGateway,
requests.exceptions.ChunkedEncodingError,
requests.exceptions.ConnectionError,
+ requests.exceptions.Timeout,
auth_exceptions.TransportError,
)
+_DEFAULT_JOB_DEADLINE = 60.0 * 10.0 # seconds
+
def _should_retry(exc):
"""Predicate for determining when to retry.
@@ -56,3 +59,28 @@ def _should_retry(exc):
on ``DEFAULT_RETRY``. For example, to change the deadline to 30 seconds,
pass ``retry=bigquery.DEFAULT_RETRY.with_deadline(30)``.
"""
+
+DEFAULT_TIMEOUT = 5.0 * 60.0
+"""The default API timeout.
+
+This is the time to wait per request. To adjust the total wait time, set a
+deadline on the retry object.
+"""
+
+job_retry_reasons = "rateLimitExceeded", "backendError"
+
+
+def _job_should_retry(exc):
+ if not hasattr(exc, "errors") or len(exc.errors) == 0:
+ return False
+
+ reason = exc.errors[0]["reason"]
+ return reason in job_retry_reasons
+
+
+DEFAULT_JOB_RETRY = retry.Retry(
+ predicate=_job_should_retry, deadline=_DEFAULT_JOB_DEADLINE
+)
+"""
+The default job retry object.
+"""
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index daade1ac6..609c0b57e 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -20,7 +20,6 @@
import datetime
import functools
import operator
-import pytz
import typing
from typing import Any, Dict, Iterable, Iterator, Optional, Tuple
import warnings
@@ -30,6 +29,20 @@
except ImportError: # pragma: NO COVER
pandas = None
+try:
+ import geopandas
+except ImportError:
+ geopandas = None
+else:
+ _COORDINATE_REFERENCE_SYSTEM = "EPSG:4326"
+
+try:
+ import shapely.geos
+except ImportError:
+ shapely = None
+else:
+ _read_wkt = shapely.geos.WKTReader(shapely.geos.lgeos).read
+
try:
import pyarrow
except ImportError: # pragma: NO COVER
@@ -53,6 +66,7 @@
# Unconditionally import optional dependencies again to tell pytype that
# they are not None, avoiding false "no attribute" errors.
import pandas
+ import geopandas
import pyarrow
from google.cloud import bigquery_storage
@@ -61,6 +75,14 @@
"The pandas library is not installed, please install "
"pandas to use the to_dataframe() function."
)
+_NO_GEOPANDAS_ERROR = (
+ "The geopandas library is not installed, please install "
+ "geopandas to use the to_geodataframe() function."
+)
+_NO_SHAPELY_ERROR = (
+ "The shapely library is not installed, please install "
+ "shapely to use the geography_as_object option."
+)
_NO_PYARROW_ERROR = (
"The pyarrow library is not installed, please install "
"pyarrow to use the to_arrow() function."
@@ -255,9 +277,16 @@ def _key(self):
return (self._project, self._dataset_id, self._table_id)
def __eq__(self, other):
- if not isinstance(other, TableReference):
+ if isinstance(other, (Table, TableListItem)):
+ return (
+ self.project == other.project
+ and self.dataset_id == other.dataset_id
+ and self.table_id == other.table_id
+ )
+ elif isinstance(other, TableReference):
+ return self._key() == other._key()
+ else:
return NotImplemented
- return self._key() == other._key()
def __ne__(self, other):
return not self == other
@@ -1011,6 +1040,24 @@ def _build_resource(self, filter_fields):
"""Generate a resource for ``update``."""
return _helpers._build_resource_from_properties(self, filter_fields)
+ def __eq__(self, other):
+ if isinstance(other, Table):
+ return (
+ self._properties["tableReference"]
+ == other._properties["tableReference"]
+ )
+ elif isinstance(other, (TableReference, TableListItem)):
+ return (
+ self.project == other.project
+ and self.dataset_id == other.dataset_id
+ and self.table_id == other.table_id
+ )
+ else:
+ return NotImplemented
+
+ def __hash__(self):
+ return hash((self.project, self.dataset_id, self.table_id))
+
def __repr__(self):
return "Table({})".format(repr(self.reference))
@@ -1229,6 +1276,19 @@ def to_api_repr(self) -> dict:
"""
return copy.deepcopy(self._properties)
+ def __eq__(self, other):
+ if isinstance(other, (Table, TableReference, TableListItem)):
+ return (
+ self.project == other.project
+ and self.dataset_id == other.dataset_id
+ and self.table_id == other.table_id
+ )
+ else:
+ return NotImplemented
+
+ def __hash__(self):
+ return hash((self.project, self.dataset_id, self.table_id))
+
def _row_from_mapping(mapping, schema):
"""Convert a mapping to a row tuple using the schema.
@@ -1841,6 +1901,7 @@ def to_dataframe(
progress_bar_type: str = None,
create_bqstorage_client: bool = True,
date_as_object: bool = True,
+ geography_as_object: bool = False,
) -> "pandas.DataFrame":
"""Create a pandas DataFrame by loading all pages of a query.
@@ -1896,6 +1957,13 @@ def to_dataframe(
.. versionadded:: 1.26.0
+ geography_as_object (Optional[bool]):
+ If ``True``, convert GEOGRAPHY data to :mod:`shapely`
+ geometry objects. If ``False`` (default), don't cast
+ geography data to :mod:`shapely` geometry objects.
+
+ .. versionadded:: 2.24.0
+
Returns:
pandas.DataFrame:
A :class:`~pandas.DataFrame` populated with row data and column
@@ -1904,13 +1972,18 @@ def to_dataframe(
Raises:
ValueError:
- If the :mod:`pandas` library cannot be imported, or the
- :mod:`google.cloud.bigquery_storage_v1` module is
- required but cannot be imported.
+ If the :mod:`pandas` library cannot be imported, or
+ the :mod:`google.cloud.bigquery_storage_v1` module is
+ required but cannot be imported. Also if
+ `geography_as_object` is `True`, but the
+ :mod:`shapely` library cannot be imported.
"""
if pandas is None:
raise ValueError(_NO_PANDAS_ERROR)
+ if geography_as_object and shapely is None:
+ raise ValueError(_NO_SHAPELY_ERROR)
+
if dtypes is None:
dtypes = {}
@@ -1931,7 +2004,7 @@ def to_dataframe(
# Pandas, we set the timestamp_as_object parameter to True, if necessary.
types_to_check = {
pyarrow.timestamp("us"),
- pyarrow.timestamp("us", tz=pytz.UTC),
+ pyarrow.timestamp("us", tz=datetime.timezone.utc),
}
for column in record_batch:
@@ -1951,8 +2024,136 @@ def to_dataframe(
for column in dtypes:
df[column] = pandas.Series(df[column], dtype=dtypes[column])
+ if geography_as_object:
+ for field in self.schema:
+ if field.field_type.upper() == "GEOGRAPHY":
+ df[field.name] = df[field.name].dropna().apply(_read_wkt)
+
return df
+ # If changing the signature of this method, make sure to apply the same
+ # changes to job.QueryJob.to_geodataframe()
+ def to_geodataframe(
+ self,
+ bqstorage_client: "bigquery_storage.BigQueryReadClient" = None,
+ dtypes: Dict[str, Any] = None,
+ progress_bar_type: str = None,
+ create_bqstorage_client: bool = True,
+ date_as_object: bool = True,
+ geography_column: Optional[str] = None,
+ ) -> "geopandas.GeoDataFrame":
+ """Create a GeoPandas GeoDataFrame by loading all pages of a query.
+
+ Args:
+ bqstorage_client (Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient]):
+ A BigQuery Storage API client. If supplied, use the faster
+ BigQuery Storage API to fetch rows from BigQuery.
+
+ This method requires the ``pyarrow`` and
+ ``google-cloud-bigquery-storage`` libraries.
+
+ This method only exposes a subset of the capabilities of the
+ BigQuery Storage API. For full access to all features
+ (projections, filters, snapshots) use the Storage API directly.
+
+ dtypes (Optional[Map[str, Union[str, pandas.Series.dtype]]]):
+ A dictionary of column names pandas ``dtype``s. The provided
+ ``dtype`` is used when constructing the series for the column
+ specified. Otherwise, the default pandas behavior is used.
+ progress_bar_type (Optional[str]):
+ If set, use the `tqdm `_ library to
+ display a progress bar while the data downloads. Install the
+ ``tqdm`` package to use this feature.
+
+ Possible values of ``progress_bar_type`` include:
+
+ ``None``
+ No progress bar.
+ ``'tqdm'``
+ Use the :func:`tqdm.tqdm` function to print a progress bar
+ to :data:`sys.stderr`.
+ ``'tqdm_notebook'``
+ Use the :func:`tqdm.tqdm_notebook` function to display a
+ progress bar as a Jupyter notebook widget.
+ ``'tqdm_gui'``
+ Use the :func:`tqdm.tqdm_gui` function to display a
+ progress bar as a graphical dialog box.
+
+ create_bqstorage_client (Optional[bool]):
+ If ``True`` (default), create a BigQuery Storage API client
+ using the default API settings. The BigQuery Storage API
+ is a faster way to fetch rows from BigQuery. See the
+ ``bqstorage_client`` parameter for more information.
+
+ This argument does nothing if ``bqstorage_client`` is supplied.
+
+ date_as_object (Optional[bool]):
+ If ``True`` (default), cast dates to objects. If ``False``, convert
+ to datetime64[ns] dtype.
+
+ geography_column (Optional[str]):
+ If there are more than one GEOGRAPHY column,
+ identifies which one to use to construct a geopandas
+ GeoDataFrame. This option can be ommitted if there's
+ only one GEOGRAPHY column.
+
+ Returns:
+ geopandas.GeoDataFrame:
+ A :class:`geopandas.GeoDataFrame` populated with row
+ data and column headers from the query results. The
+ column headers are derived from the destination
+ table's schema.
+
+ Raises:
+ ValueError:
+ If the :mod:`geopandas` library cannot be imported, or the
+ :mod:`google.cloud.bigquery_storage_v1` module is
+ required but cannot be imported.
+
+ .. versionadded:: 2.24.0
+ """
+ if geopandas is None:
+ raise ValueError(_NO_GEOPANDAS_ERROR)
+
+ geography_columns = set(
+ field.name
+ for field in self.schema
+ if field.field_type.upper() == "GEOGRAPHY"
+ )
+ if not geography_columns:
+ raise TypeError(
+ "There must be at least one GEOGRAPHY column"
+ " to create a GeoDataFrame"
+ )
+
+ if geography_column:
+ if geography_column not in geography_columns:
+ raise ValueError(
+ f"The given geography column, {geography_column}, doesn't name"
+ f" a GEOGRAPHY column in the result."
+ )
+ elif len(geography_columns) == 1:
+ [geography_column] = geography_columns
+ else:
+ raise ValueError(
+ "There is more than one GEOGRAPHY column in the result. "
+ "The geography_column argument must be used to specify which "
+ "one to use to create a GeoDataFrame"
+ )
+
+ df = self.to_dataframe(
+ bqstorage_client,
+ dtypes,
+ progress_bar_type,
+ create_bqstorage_client,
+ date_as_object,
+ geography_as_object=True,
+ )
+
+ return geopandas.GeoDataFrame(
+ df, crs=_COORDINATE_REFERENCE_SYSTEM, geometry=geography_column
+ )
+
class _EmptyRowIterator(RowIterator):
"""An empty row iterator.
@@ -2005,6 +2206,7 @@ def to_dataframe(
progress_bar_type=None,
create_bqstorage_client=True,
date_as_object=True,
+ geography_as_object=False,
) -> "pandas.DataFrame":
"""Create an empty dataframe.
@@ -2022,6 +2224,31 @@ def to_dataframe(
raise ValueError(_NO_PANDAS_ERROR)
return pandas.DataFrame()
+ def to_geodataframe(
+ self,
+ bqstorage_client=None,
+ dtypes=None,
+ progress_bar_type=None,
+ create_bqstorage_client=True,
+ date_as_object=True,
+ geography_column: Optional[str] = None,
+ ) -> "pandas.DataFrame":
+ """Create an empty dataframe.
+
+ Args:
+ bqstorage_client (Any): Ignored. Added for compatibility with RowIterator.
+ dtypes (Any): Ignored. Added for compatibility with RowIterator.
+ progress_bar_type (Any): Ignored. Added for compatibility with RowIterator.
+ create_bqstorage_client (bool): Ignored. Added for compatibility with RowIterator.
+ date_as_object (bool): Ignored. Added for compatibility with RowIterator.
+
+ Returns:
+ pandas.DataFrame: An empty :class:`~pandas.DataFrame`.
+ """
+ if geopandas is None:
+ raise ValueError(_NO_GEOPANDAS_ERROR)
+ return geopandas.GeoDataFrame(crs=_COORDINATE_REFERENCE_SYSTEM)
+
def to_dataframe_iterable(
self,
bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None,
diff --git a/google/cloud/bigquery/version.py b/google/cloud/bigquery/version.py
index 0460e7bb9..21cbec9fe 100644
--- a/google/cloud/bigquery/version.py
+++ b/google/cloud/bigquery/version.py
@@ -12,4 +12,4 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-__version__ = "2.23.2"
+__version__ = "2.25.1"
diff --git a/noxfile.py b/noxfile.py
index 0dfe7bf93..9077924e9 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -160,10 +160,6 @@ def snippets(session):
if os.environ.get("RUN_SNIPPETS_TESTS", "true") == "false":
session.skip("RUN_SNIPPETS_TESTS is set to false, skipping")
- # Sanity check: Only run snippets tests if the environment variable is set.
- if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", ""):
- session.skip("Credentials must be set via environment variable.")
-
constraints_path = str(
CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt"
)
diff --git a/owlbot.py b/owlbot.py
index 09845480a..8664b658a 100644
--- a/owlbot.py
+++ b/owlbot.py
@@ -63,7 +63,7 @@
s.replace(
library / f"google/cloud/bigquery_{library.name}/types/standard_sql.py",
r"type_ ",
- "type "
+ "type ",
)
s.move(
@@ -78,8 +78,8 @@
"noxfile.py",
"setup.py",
f"scripts/fixup_bigquery_{library.name}_keywords.py",
- f"google/cloud/bigquery/__init__.py",
- f"google/cloud/bigquery/py.typed",
+ "google/cloud/bigquery/__init__.py",
+ "google/cloud/bigquery/py.typed",
# There are no public API endpoints for the generated ModelServiceClient,
# thus there's no point in generating it and its tests.
f"google/cloud/bigquery_{library.name}/services/**",
@@ -97,6 +97,10 @@
samples=True,
microgenerator=True,
split_system_tests=True,
+ intersphinx_dependencies={
+ "pandas": "http://pandas.pydata.org/pandas-docs/dev",
+ "geopandas": "https://geopandas.org/",
+ },
)
# BigQuery has a custom multiprocessing note
@@ -109,7 +113,11 @@
# Include custom SNIPPETS_TESTS job for performance.
# https://github.com/googleapis/python-bigquery/issues/191
".kokoro/presubmit/presubmit.cfg",
- ]
+ # Group all renovate PRs together. If this works well, remove this and
+ # update the shared templates (possibly with configuration option to
+ # py_library.)
+ "renovate.json",
+ ],
)
# ----------------------------------------------------------------------------
@@ -121,14 +129,14 @@
s.replace(
"docs/conf.py",
r'\{"members": True\}',
- '{"members": True, "inherited-members": True}'
+ '{"members": True, "inherited-members": True}',
)
# Tell Sphinx to ingore autogenerated docs files.
s.replace(
"docs/conf.py",
r'"samples/snippets/README\.rst",',
- '\g<0>\n "bigquery_v2/services.rst", # generated by the code generator',
+ '\\g<0>\n "bigquery_v2/services.rst", # generated by the code generator',
)
# ----------------------------------------------------------------------------
@@ -136,13 +144,14 @@
# ----------------------------------------------------------------------------
# Add .pytype to .gitignore
-s.replace(".gitignore", r"\.pytest_cache", "\g<0>\n.pytype")
+s.replace(".gitignore", r"\.pytest_cache", "\\g<0>\n.pytype")
# Add pytype config to setup.cfg
s.replace(
"setup.cfg",
r"universal = 1",
- textwrap.dedent(""" \g<0>
+ textwrap.dedent(
+ """ \\g<0>
[pytype]
python_version = 3.8
@@ -156,7 +165,8 @@
# There's some issue with finding some pyi files, thus disabling.
# The issue https://github.com/google/pytype/issues/150 is closed, but the
# error still occurs for some reason.
- pyi-error""")
+ pyi-error"""
+ ),
)
s.shell.run(["nox", "-s", "blacken"], hide_output=False)
diff --git a/renovate.json b/renovate.json
index c04895563..713c60bb4 100644
--- a/renovate.json
+++ b/renovate.json
@@ -1,6 +1,6 @@
{
"extends": [
- "config:base", ":preserveSemverRanges"
+ "config:base", "group:all", ":preserveSemverRanges"
],
"ignorePaths": [".pre-commit-config.yaml"],
"pip_requirements": {
diff --git a/samples/client_query_w_timestamp_params.py b/samples/client_query_w_timestamp_params.py
index ca8eec0b5..41a27770e 100644
--- a/samples/client_query_w_timestamp_params.py
+++ b/samples/client_query_w_timestamp_params.py
@@ -18,7 +18,6 @@ def client_query_w_timestamp_params():
# [START bigquery_query_params_timestamps]
import datetime
- import pytz
from google.cloud import bigquery
# Construct a BigQuery client object.
@@ -30,7 +29,7 @@ def client_query_w_timestamp_params():
bigquery.ScalarQueryParameter(
"ts_value",
"TIMESTAMP",
- datetime.datetime(2016, 12, 7, 8, 0, tzinfo=pytz.UTC),
+ datetime.datetime(2016, 12, 7, 8, 0, tzinfo=datetime.timezone.utc),
)
]
)
diff --git a/samples/create_routine.py b/samples/create_routine.py
index 012c7927a..1cb4a80b4 100644
--- a/samples/create_routine.py
+++ b/samples/create_routine.py
@@ -22,7 +22,7 @@ def create_routine(routine_id):
# Construct a BigQuery client object.
client = bigquery.Client()
- # TODO(developer): Choose a fully-qualified ID for the routine.
+ # TODO(developer): Choose a fully qualified ID for the routine.
# routine_id = "my-project.my_dataset.my_routine"
routine = bigquery.Routine(
diff --git a/samples/geography/noxfile.py b/samples/geography/noxfile.py
index 9fc7f1782..b008613f0 100644
--- a/samples/geography/noxfile.py
+++ b/samples/geography/noxfile.py
@@ -39,7 +39,7 @@
TEST_CONFIG = {
# You can opt out from the test for specific Python versions.
- "ignored_versions": ["2.7"],
+ "ignored_versions": [],
# Old samples are opted out of enforcing Python type hints
# All new samples should feature them
"enforce_type_hints": False,
@@ -86,15 +86,18 @@ def get_pytest_env_vars() -> Dict[str, str]:
# DO NOT EDIT - automatically generated.
-# All versions used to tested samples.
-ALL_VERSIONS = ["2.7", "3.6", "3.7", "3.8", "3.9"]
+# All versions used to test samples.
+ALL_VERSIONS = ["3.6", "3.7", "3.8", "3.9"]
# Any default versions that should be ignored.
IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"]
TESTED_VERSIONS = sorted([v for v in ALL_VERSIONS if v not in IGNORED_VERSIONS])
-INSTALL_LIBRARY_FROM_SOURCE = bool(os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False))
+INSTALL_LIBRARY_FROM_SOURCE = os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False) in (
+ "True",
+ "true",
+)
#
# Style Checks
#
diff --git a/samples/geography/requirements.txt b/samples/geography/requirements.txt
index 5aa967b24..b5fe247cb 100644
--- a/samples/geography/requirements.txt
+++ b/samples/geography/requirements.txt
@@ -1,4 +1,50 @@
+attrs==21.2.0
+cachetools==4.2.2
+certifi==2021.5.30
+cffi==1.14.6
+charset-normalizer==2.0.4
+click==8.0.1
+click-plugins==1.1.1
+cligj==0.7.2
+dataclasses==0.6; python_version < '3.7'
+Fiona==1.8.20
geojson==2.5.0
-google-cloud-bigquery==2.23.2
-google-cloud-bigquery-storage==2.6.2
+geopandas==0.9.0
+google-api-core==1.31.2
+google-auth==1.35.0
+google-cloud-bigquery==2.25.0
+google-cloud-bigquery-storage==2.6.3
+google-cloud-core==2.0.0
+google-crc32c==1.1.2
+google-resumable-media==1.3.3
+googleapis-common-protos==1.53.0
+grpcio==1.39.0
+idna==3.2
+importlib-metadata==4.6.4
+libcst==0.3.20
+munch==2.5.0
+mypy-extensions==0.4.3
+numpy==1.19.5; python_version < "3.7"
+numpy==1.21.2; python_version > "3.6"
+packaging==21.0
+pandas==1.1.5; python_version < '3.7'
+pandas==1.3.2; python_version >= '3.7'
+proto-plus==1.19.0
+protobuf==3.17.3
+pyarrow==5.0.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pycparser==2.20
+pyparsing==2.4.7
+pyproj==3.0.1
+python-dateutil==2.8.2
+pytz==2021.1
+PyYAML==5.4.1
+requests==2.26.0
+rsa==4.7.2
Shapely==1.7.1
+six==1.16.0
+typing-extensions==3.10.0.0
+typing-inspect==0.7.1
+urllib3==1.26.6
+zipp==3.5.0
diff --git a/samples/geography/to_geodataframe.py b/samples/geography/to_geodataframe.py
new file mode 100644
index 000000000..fa8073fef
--- /dev/null
+++ b/samples/geography/to_geodataframe.py
@@ -0,0 +1,32 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from google.cloud import bigquery
+
+client = bigquery.Client()
+
+
+def get_austin_service_requests_as_geography():
+ # [START bigquery_query_results_geodataframe]
+
+ sql = """
+ SELECT created_date, complaint_description,
+ ST_GEOGPOINT(longitude, latitude) as location
+ FROM bigquery-public-data.austin_311.311_service_requests
+ LIMIT 10
+ """
+
+ df = client.query(sql).to_geodataframe()
+ # [END bigquery_query_results_geodataframe]
+ return df
diff --git a/samples/geography/to_geodataframe_test.py b/samples/geography/to_geodataframe_test.py
new file mode 100644
index 000000000..7a2ba6937
--- /dev/null
+++ b/samples/geography/to_geodataframe_test.py
@@ -0,0 +1,25 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 pytest
+
+from .to_geodataframe import get_austin_service_requests_as_geography
+
+
+def test_get_austin_service_requests_as_geography():
+ geopandas = pytest.importorskip("geopandas")
+ df = get_austin_service_requests_as_geography()
+ assert isinstance(df, geopandas.GeoDataFrame)
+ assert len(list(df)) == 3 # verify the number of columns
+ assert len(df) == 10 # verify the number of rows
diff --git a/samples/snippets/noxfile.py b/samples/snippets/noxfile.py
index 9fc7f1782..b008613f0 100644
--- a/samples/snippets/noxfile.py
+++ b/samples/snippets/noxfile.py
@@ -39,7 +39,7 @@
TEST_CONFIG = {
# You can opt out from the test for specific Python versions.
- "ignored_versions": ["2.7"],
+ "ignored_versions": [],
# Old samples are opted out of enforcing Python type hints
# All new samples should feature them
"enforce_type_hints": False,
@@ -86,15 +86,18 @@ def get_pytest_env_vars() -> Dict[str, str]:
# DO NOT EDIT - automatically generated.
-# All versions used to tested samples.
-ALL_VERSIONS = ["2.7", "3.6", "3.7", "3.8", "3.9"]
+# All versions used to test samples.
+ALL_VERSIONS = ["3.6", "3.7", "3.8", "3.9"]
# Any default versions that should be ignored.
IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"]
TESTED_VERSIONS = sorted([v for v in ALL_VERSIONS if v not in IGNORED_VERSIONS])
-INSTALL_LIBRARY_FROM_SOURCE = bool(os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False))
+INSTALL_LIBRARY_FROM_SOURCE = os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False) in (
+ "True",
+ "true",
+)
#
# Style Checks
#
diff --git a/samples/snippets/requirements.txt b/samples/snippets/requirements.txt
index 4f2eaf90b..d75c747fb 100644
--- a/samples/snippets/requirements.txt
+++ b/samples/snippets/requirements.txt
@@ -1,5 +1,5 @@
-google-cloud-bigquery==2.23.2
-google-cloud-bigquery-storage==2.6.2
+google-cloud-bigquery==2.25.0
+google-cloud-bigquery-storage==2.6.3
google-auth-oauthlib==0.4.5
grpcio==1.39.0
ipython==7.16.1; python_version < '3.7'
@@ -7,6 +7,6 @@ ipython==7.17.0; python_version >= '3.7'
matplotlib==3.3.4; python_version < '3.7'
matplotlib==3.4.1; python_version >= '3.7'
pandas==1.1.5; python_version < '3.7'
-pandas==1.2.0; python_version >= '3.7'
+pandas==1.3.2; python_version >= '3.7'
pyarrow==5.0.0
pytz==2021.1
diff --git a/scripts/readme-gen/templates/install_deps.tmpl.rst b/scripts/readme-gen/templates/install_deps.tmpl.rst
index a0406dba8..275d64989 100644
--- a/scripts/readme-gen/templates/install_deps.tmpl.rst
+++ b/scripts/readme-gen/templates/install_deps.tmpl.rst
@@ -12,7 +12,7 @@ Install Dependencies
.. _Python Development Environment Setup Guide:
https://cloud.google.com/python/setup
-#. Create a virtualenv. Samples are compatible with Python 2.7 and 3.4+.
+#. Create a virtualenv. Samples are compatible with Python 3.6+.
.. code-block:: bash
diff --git a/setup.py b/setup.py
index e9deaf117..e7515493d 100644
--- a/setup.py
+++ b/setup.py
@@ -54,9 +54,10 @@
# grpc.Channel.close() method isn't added until 1.32.0.
# https://github.com/grpc/grpc/pull/15254
"grpcio >= 1.38.1, < 2.0dev",
- "pyarrow >= 1.0.0, < 6.0dev",
+ "pyarrow >= 3.0.0, < 6.0dev",
],
- "pandas": ["pandas>=0.23.0", "pyarrow >= 1.0.0, < 6.0dev"],
+ "geopandas": ["geopandas>=0.9.0, <1.0dev", "Shapely>=1.6.0, <2.0dev"],
+ "pandas": ["pandas>=0.23.0", "pyarrow >= 3.0.0, < 6.0dev"],
"bignumeric_type": ["pyarrow >= 3.0.0, < 6.0dev"],
"tqdm": ["tqdm >= 4.7.4, <5.0.0dev"],
"opentelemetry": [
diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt
index af6e82efd..be1a992fa 100644
--- a/testing/constraints-3.6.txt
+++ b/testing/constraints-3.6.txt
@@ -5,6 +5,7 @@
#
# e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev",
# Then this file should have foo==1.14.0
+geopandas==0.9.0
google-api-core==1.29.0
google-cloud-bigquery-storage==2.0.0
google-cloud-core==1.4.1
@@ -13,10 +14,11 @@ grpcio==1.38.1
opentelemetry-api==0.11b0
opentelemetry-instrumentation==0.11b0
opentelemetry-sdk==0.11b0
-pandas==0.23.0
+pandas==0.24.2
proto-plus==1.10.0
protobuf==3.12.0
-pyarrow==1.0.0
+pyarrow==3.0.0
requests==2.18.0
+shapely==1.6.0
six==1.13.0
tqdm==4.7.4
diff --git a/tests/data/scalars.jsonl b/tests/data/scalars.jsonl
index 4419a6e9a..e06139e5c 100644
--- a/tests/data/scalars.jsonl
+++ b/tests/data/scalars.jsonl
@@ -1,2 +1,2 @@
-{"bool_col": true, "bytes_col": "abcd", "date_col": "2021-07-21", "datetime_col": "2021-07-21 11:39:45", "geography_col": "POINT(-122.0838511 37.3860517)", "int64_col": "123456789", "numeric_col": "1.23456789", "bignumeric_col": "10.111213141516171819", "float64_col": "1.25", "string_col": "Hello, World", "time_col": "11:41:43.07616", "timestamp_col": "2021-07-21T17:43:43.945289Z"}
-{"bool_col": null, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": null, "numeric_col": null, "bignumeric_col": null, "float64_col": null, "string_col": null, "time_col": null, "timestamp_col": null}
+{"bool_col": true, "bytes_col": "SGVsbG8sIFdvcmxkIQ==", "date_col": "2021-07-21", "datetime_col": "2021-07-21 11:39:45", "geography_col": "POINT(-122.0838511 37.3860517)", "int64_col": "123456789", "interval_col": "P7Y11M9DT4H15M37.123456S", "numeric_col": "1.23456789", "bignumeric_col": "10.111213141516171819", "float64_col": "1.25", "rowindex": 0, "string_col": "Hello, World!", "time_col": "11:41:43.07616", "timestamp_col": "2021-07-21T17:43:43.945289Z"}
+{"bool_col": null, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": null, "interval_col": null, "numeric_col": null, "bignumeric_col": null, "float64_col": null, "rowindex": 1, "string_col": null, "time_col": null, "timestamp_col": null}
diff --git a/tests/data/scalars_extreme.jsonl b/tests/data/scalars_extreme.jsonl
index ceccd8dbc..d0a33fdba 100644
--- a/tests/data/scalars_extreme.jsonl
+++ b/tests/data/scalars_extreme.jsonl
@@ -1,5 +1,5 @@
-{"bool_col": true, "bytes_col": "DQo=\n", "date_col": "9999-12-31", "datetime_col": "9999-12-31 23:59:59.999999", "geography_col": "POINT(-135.0000 90.0000)", "int64_col": "9223372036854775807", "numeric_col": "9.9999999999999999999999999999999999999E+28", "bignumeric_col": "9.999999999999999999999999999999999999999999999999999999999999999999999999999E+37", "float64_col": "+inf", "string_col": "Hello, World", "time_col": "23:59:59.99999", "timestamp_col": "9999-12-31T23:59:59.999999Z"}
-{"bool_col": false, "bytes_col": "8J+Zgw==\n", "date_col": "0001-01-01", "datetime_col": "0001-01-01 00:00:00", "geography_col": "POINT(45.0000 -90.0000)", "int64_col": "-9223372036854775808", "numeric_col": "-9.9999999999999999999999999999999999999E+28", "bignumeric_col": "-9.999999999999999999999999999999999999999999999999999999999999999999999999999E+37", "float64_col": "-inf", "string_col": "Hello, World", "time_col": "00:00:00", "timestamp_col": "0001-01-01T00:00:00.000000Z"}
-{"bool_col": true, "bytes_col": "AA==\n", "date_col": "1900-01-01", "datetime_col": "1900-01-01 00:00:00", "geography_col": "POINT(-180.0000 0.0000)", "int64_col": "-1", "numeric_col": "0.000000001", "bignumeric_col": "-0.00000000000000000000000000000000000001", "float64_col": "nan", "string_col": "こんにちは", "time_col": "00:00:00.000001", "timestamp_col": "1900-01-01T00:00:00.000000Z"}
-{"bool_col": false, "bytes_col": "", "date_col": "1970-01-01", "datetime_col": "1970-01-01 00:00:00", "geography_col": "POINT(0 0)", "int64_col": "0", "numeric_col": "0.0", "bignumeric_col": "0.0", "float64_col": 0.0, "string_col": "", "time_col": "12:00:00", "timestamp_col": "1970-01-01T00:00:00.000000Z"}
-{"bool_col": null, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": null, "numeric_col": null, "bignumeric_col": null, "float64_col": null, "string_col": null, "time_col": null, "timestamp_col": null}
+{"bool_col": true, "bytes_col": "DQo=\n", "date_col": "9999-12-31", "datetime_col": "9999-12-31 23:59:59.999999", "geography_col": "POINT(-135.0000 90.0000)", "int64_col": "9223372036854775807", "interval_col": "P-10000Y0M-3660000DT-87840000H0M0S", "numeric_col": "9.9999999999999999999999999999999999999E+28", "bignumeric_col": "9.999999999999999999999999999999999999999999999999999999999999999999999999999E+37", "float64_col": "+inf", "rowindex": 0, "string_col": "Hello, World", "time_col": "23:59:59.999999", "timestamp_col": "9999-12-31T23:59:59.999999Z"}
+{"bool_col": false, "bytes_col": "8J+Zgw==\n", "date_col": "0001-01-01", "datetime_col": "0001-01-01 00:00:00", "geography_col": "POINT(45.0000 -90.0000)", "int64_col": "-9223372036854775808", "interval_col": "P10000Y0M3660000DT87840000H0M0S", "numeric_col": "-9.9999999999999999999999999999999999999E+28", "bignumeric_col": "-9.999999999999999999999999999999999999999999999999999999999999999999999999999E+37", "float64_col": "-inf", "rowindex": 1, "string_col": "Hello, World", "time_col": "00:00:00", "timestamp_col": "0001-01-01T00:00:00.000000Z"}
+{"bool_col": true, "bytes_col": "AA==\n", "date_col": "1900-01-01", "datetime_col": "1900-01-01 00:00:00", "geography_col": "POINT(-180.0000 0.0000)", "int64_col": "-1", "interval_col": "P0Y0M0DT0H0M0.000001S", "numeric_col": "0.000000001", "bignumeric_col": "-0.00000000000000000000000000000000000001", "float64_col": "nan", "rowindex": 2, "string_col": "こんにちは", "time_col": "00:00:00.000001", "timestamp_col": "1900-01-01T00:00:00.000000Z"}
+{"bool_col": false, "bytes_col": "", "date_col": "1970-01-01", "datetime_col": "1970-01-01 00:00:00", "geography_col": "POINT(0 0)", "int64_col": "0", "interval_col": "P0Y0M0DT0H0M0S", "numeric_col": "0.0", "bignumeric_col": "0.0", "float64_col": 0.0, "rowindex": 3, "string_col": "", "time_col": "12:00:00", "timestamp_col": "1970-01-01T00:00:00.000000Z"}
+{"bool_col": null, "bytes_col": null, "date_col": null, "datetime_col": null, "geography_col": null, "int64_col": null, "interval_col": null, "numeric_col": null, "bignumeric_col": null, "float64_col": null, "rowindex": 4, "string_col": null, "time_col": null, "timestamp_col": null}
diff --git a/tests/data/scalars_schema.json b/tests/data/scalars_schema.json
index 00bd150fd..676d37d56 100644
--- a/tests/data/scalars_schema.json
+++ b/tests/data/scalars_schema.json
@@ -1,33 +1,33 @@
[
{
"mode": "NULLABLE",
- "name": "timestamp_col",
- "type": "TIMESTAMP"
+ "name": "bool_col",
+ "type": "BOOLEAN"
},
{
"mode": "NULLABLE",
- "name": "time_col",
- "type": "TIME"
+ "name": "bignumeric_col",
+ "type": "BIGNUMERIC"
},
{
"mode": "NULLABLE",
- "name": "float64_col",
- "type": "FLOAT"
+ "name": "bytes_col",
+ "type": "BYTES"
},
{
"mode": "NULLABLE",
- "name": "datetime_col",
- "type": "DATETIME"
+ "name": "date_col",
+ "type": "DATE"
},
{
"mode": "NULLABLE",
- "name": "bignumeric_col",
- "type": "BIGNUMERIC"
+ "name": "datetime_col",
+ "type": "DATETIME"
},
{
"mode": "NULLABLE",
- "name": "numeric_col",
- "type": "NUMERIC"
+ "name": "float64_col",
+ "type": "FLOAT"
},
{
"mode": "NULLABLE",
@@ -36,27 +36,37 @@
},
{
"mode": "NULLABLE",
- "name": "date_col",
- "type": "DATE"
+ "name": "int64_col",
+ "type": "INTEGER"
},
{
"mode": "NULLABLE",
- "name": "string_col",
- "type": "STRING"
+ "name": "interval_col",
+ "type": "INTERVAL"
},
{
"mode": "NULLABLE",
- "name": "bool_col",
- "type": "BOOLEAN"
+ "name": "numeric_col",
+ "type": "NUMERIC"
+ },
+ {
+ "mode": "REQUIRED",
+ "name": "rowindex",
+ "type": "INTEGER"
},
{
"mode": "NULLABLE",
- "name": "bytes_col",
- "type": "BYTES"
+ "name": "string_col",
+ "type": "STRING"
},
{
"mode": "NULLABLE",
- "name": "int64_col",
- "type": "INTEGER"
+ "name": "time_col",
+ "type": "TIME"
+ },
+ {
+ "mode": "NULLABLE",
+ "name": "timestamp_col",
+ "type": "TIMESTAMP"
}
]
diff --git a/tests/system/test_arrow.py b/tests/system/test_arrow.py
index f97488e39..12f7af9cb 100644
--- a/tests/system/test_arrow.py
+++ b/tests/system/test_arrow.py
@@ -14,8 +14,14 @@
"""System tests for Arrow connector."""
+from typing import Optional
+
import pytest
+from google.cloud import bigquery
+from google.cloud.bigquery import enums
+
+
pyarrow = pytest.importorskip(
"pyarrow", minversion="3.0.0"
) # Needs decimal256 for BIGNUMERIC columns.
@@ -31,17 +37,35 @@
),
)
def test_list_rows_nullable_scalars_dtypes(
- bigquery_client,
- scalars_table,
- scalars_extreme_table,
- max_results,
- scalars_table_name,
+ bigquery_client: bigquery.Client,
+ scalars_table: str,
+ scalars_extreme_table: str,
+ max_results: Optional[int],
+ scalars_table_name: str,
):
table_id = scalars_table
if scalars_table_name == "scalars_extreme_table":
table_id = scalars_extreme_table
+
+ # TODO(GH#836): Avoid INTERVAL columns until they are supported by the
+ # BigQuery Storage API and pyarrow.
+ schema = [
+ bigquery.SchemaField("bool_col", enums.SqlTypeNames.BOOLEAN),
+ bigquery.SchemaField("bignumeric_col", enums.SqlTypeNames.BIGNUMERIC),
+ bigquery.SchemaField("bytes_col", enums.SqlTypeNames.BYTES),
+ bigquery.SchemaField("date_col", enums.SqlTypeNames.DATE),
+ bigquery.SchemaField("datetime_col", enums.SqlTypeNames.DATETIME),
+ bigquery.SchemaField("float64_col", enums.SqlTypeNames.FLOAT64),
+ bigquery.SchemaField("geography_col", enums.SqlTypeNames.GEOGRAPHY),
+ bigquery.SchemaField("int64_col", enums.SqlTypeNames.INT64),
+ bigquery.SchemaField("numeric_col", enums.SqlTypeNames.NUMERIC),
+ bigquery.SchemaField("string_col", enums.SqlTypeNames.STRING),
+ bigquery.SchemaField("time_col", enums.SqlTypeNames.TIME),
+ bigquery.SchemaField("timestamp_col", enums.SqlTypeNames.TIMESTAMP),
+ ]
+
arrow_table = bigquery_client.list_rows(
- table_id, max_results=max_results,
+ table_id, max_results=max_results, selected_fields=schema,
).to_arrow()
schema = arrow_table.schema
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
index baa2b6ad8..9da45ee6e 100644
--- a/tests/system/test_client.py
+++ b/tests/system/test_client.py
@@ -30,7 +30,6 @@
import psutil
import pytest
-from google.cloud.bigquery._pandas_helpers import _BIGNUMERIC_SUPPORT
from . import helpers
try:
@@ -1557,6 +1556,40 @@ def test_dml_statistics(self):
assert query_job.dml_stats.updated_row_count == 0
assert query_job.dml_stats.deleted_row_count == 3
+ def test_transaction_info(self):
+ table_schema = (
+ bigquery.SchemaField("foo", "STRING"),
+ bigquery.SchemaField("bar", "INTEGER"),
+ )
+
+ dataset_id = _make_dataset_id("bq_system_test")
+ self.temp_dataset(dataset_id)
+ table_id = f"{Config.CLIENT.project}.{dataset_id}.test_dml_statistics"
+
+ # Create the table before loading so that the column order is deterministic.
+ table = helpers.retry_403(Config.CLIENT.create_table)(
+ Table(table_id, schema=table_schema)
+ )
+ self.to_delete.insert(0, table)
+
+ # Insert a few rows and check the stats.
+ sql = f"""
+ BEGIN TRANSACTION;
+ INSERT INTO `{table_id}`
+ VALUES ("one", 1), ("two", 2), ("three", 3), ("four", 4);
+
+ UPDATE `{table_id}`
+ SET bar = bar + 1
+ WHERE bar > 2;
+ COMMIT TRANSACTION;
+ """
+ query_job = Config.CLIENT.query(sql)
+ query_job.result()
+
+ # Transaction ID set by the server should be accessible
+ assert query_job.transaction_info is not None
+ assert query_job.transaction_info.transaction_id != ""
+
def test_dbapi_w_standard_sql_types(self):
for sql, expected in helpers.STANDARD_SQL_EXAMPLES:
Config.CURSOR.execute(sql)
@@ -1938,15 +1971,12 @@ def test_query_w_query_params(self):
"expected": {"friends": [phred_name, bharney_name]},
"query_parameters": [with_friends_param],
},
+ {
+ "sql": "SELECT @bignum_param",
+ "expected": bignum,
+ "query_parameters": [bignum_param],
+ },
]
- if _BIGNUMERIC_SUPPORT:
- examples.append(
- {
- "sql": "SELECT @bignum_param",
- "expected": bignum,
- "query_parameters": [bignum_param],
- }
- )
for example in examples:
jconfig = QueryJobConfig()
@@ -2330,9 +2360,6 @@ def test_create_table_rows_fetch_nested_schema(self):
self.assertEqual(found[7], e_favtime)
self.assertEqual(found[8], decimal.Decimal(expected["FavoriteNumber"]))
- def _fetch_dataframe(self, query):
- return Config.CLIENT.query(query).result().to_dataframe()
-
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
@unittest.skipIf(
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
@@ -2394,54 +2421,6 @@ def test_nested_table_to_arrow(self):
self.assertTrue(pyarrow.types.is_list(record_col[1].type))
self.assertTrue(pyarrow.types.is_int64(record_col[1].type.value_type))
- def test_list_rows_empty_table(self):
- from google.cloud.bigquery.table import RowIterator
-
- dataset_id = _make_dataset_id("empty_table")
- dataset = self.temp_dataset(dataset_id)
- table_ref = dataset.table("empty_table")
- table = Config.CLIENT.create_table(bigquery.Table(table_ref))
-
- # It's a bit silly to list rows for an empty table, but this does
- # happen as the result of a DDL query from an IPython magic command.
- rows = Config.CLIENT.list_rows(table)
- self.assertIsInstance(rows, RowIterator)
- self.assertEqual(tuple(rows), ())
-
- def test_list_rows_page_size(self):
- from google.cloud.bigquery.job import SourceFormat
- from google.cloud.bigquery.job import WriteDisposition
-
- num_items = 7
- page_size = 3
- num_pages, num_last_page = divmod(num_items, page_size)
-
- SF = bigquery.SchemaField
- schema = [SF("string_col", "STRING", mode="NULLABLE")]
- to_insert = [{"string_col": "item%d" % i} for i in range(num_items)]
- rows = [json.dumps(row) for row in to_insert]
- body = io.BytesIO("{}\n".format("\n".join(rows)).encode("ascii"))
-
- table_id = "test_table"
- dataset = self.temp_dataset(_make_dataset_id("nested_df"))
- table = dataset.table(table_id)
- self.to_delete.insert(0, table)
- job_config = bigquery.LoadJobConfig()
- job_config.write_disposition = WriteDisposition.WRITE_TRUNCATE
- job_config.source_format = SourceFormat.NEWLINE_DELIMITED_JSON
- job_config.schema = schema
- # Load a table using a local JSON file from memory.
- Config.CLIENT.load_table_from_file(body, table, job_config=job_config).result()
-
- df = Config.CLIENT.list_rows(table, selected_fields=schema, page_size=page_size)
- pages = df.pages
-
- for i in range(num_pages):
- page = next(pages)
- self.assertEqual(page.num_items, page_size)
- page = next(pages)
- self.assertEqual(page.num_items, num_last_page)
-
def temp_dataset(self, dataset_id, location=None):
project = Config.CLIENT.project
dataset_ref = bigquery.DatasetReference(project, dataset_id)
diff --git a/tests/system/test_job_retry.py b/tests/system/test_job_retry.py
new file mode 100644
index 000000000..520545493
--- /dev/null
+++ b/tests/system/test_job_retry.py
@@ -0,0 +1,72 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 contextlib
+import threading
+import time
+
+import google.api_core.exceptions
+import google.cloud.bigquery
+import pytest
+
+
+def thread(func):
+ thread = threading.Thread(target=func, daemon=True)
+ thread.start()
+ return thread
+
+
+@pytest.mark.parametrize("job_retry_on_query", [True, False])
+def test_query_retry_539(bigquery_client, dataset_id, job_retry_on_query):
+ """
+ Test job_retry
+
+ 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
+
+ table_name = f"{dataset_id}.t539"
+
+ # Without a custom retry, we fail:
+ with pytest.raises(google.api_core.exceptions.NotFound):
+ bigquery_client.query(f"select count(*) from {table_name}").result()
+
+ retry_notfound = Retry(predicate=if_exception_type(exceptions.NotFound))
+
+ job_retry = dict(job_retry=retry_notfound) if job_retry_on_query else {}
+ job = bigquery_client.query(f"select count(*) from {table_name}", **job_retry)
+ job_id = job.job_id
+
+ # We can already know that the job failed, but we're not supposed
+ # to find out until we call result, which is where retry happend
+ assert job.done()
+ assert job.exception() is not None
+
+ @thread
+ def create_table():
+ time.sleep(1) # Give the first retry attempt time to fail.
+ with contextlib.closing(google.cloud.bigquery.Client()) as client:
+ client.query(f"create table {table_name} (id int64)").result()
+
+ job_retry = {} if job_retry_on_query else dict(job_retry=retry_notfound)
+ [[count]] = list(job.result(**job_retry))
+ assert count == 0
+
+ # The job was retried, and thus got a new job id
+ assert job.job_id != job_id
+
+ # Make sure we don't leave a thread behind:
+ create_table.join()
+ bigquery_client.query(f"drop table {table_name}").result()
diff --git a/tests/system/test_list_rows.py b/tests/system/test_list_rows.py
new file mode 100644
index 000000000..70388059e
--- /dev/null
+++ b/tests/system/test_list_rows.py
@@ -0,0 +1,112 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 datetime
+import decimal
+
+from google.cloud import bigquery
+from google.cloud.bigquery import enums
+
+
+def test_list_rows_empty_table(bigquery_client: bigquery.Client, table_id: str):
+ from google.cloud.bigquery.table import RowIterator
+
+ table = bigquery_client.create_table(table_id)
+
+ # It's a bit silly to list rows for an empty table, but this does
+ # happen as the result of a DDL query from an IPython magic command.
+ rows = bigquery_client.list_rows(table)
+ assert isinstance(rows, RowIterator)
+ assert tuple(rows) == ()
+
+
+def test_list_rows_page_size(bigquery_client: bigquery.Client, table_id: str):
+ num_items = 7
+ page_size = 3
+ num_pages, num_last_page = divmod(num_items, page_size)
+
+ to_insert = [{"string_col": "item%d" % i, "rowindex": i} for i in range(num_items)]
+ bigquery_client.load_table_from_json(to_insert, table_id).result()
+
+ df = bigquery_client.list_rows(
+ table_id,
+ selected_fields=[bigquery.SchemaField("string_col", enums.SqlTypeNames.STRING)],
+ page_size=page_size,
+ )
+ pages = df.pages
+
+ for i in range(num_pages):
+ page = next(pages)
+ assert page.num_items == page_size
+ page = next(pages)
+ assert page.num_items == num_last_page
+
+
+def test_list_rows_scalars(bigquery_client: bigquery.Client, scalars_table: str):
+ rows = sorted(
+ bigquery_client.list_rows(scalars_table), key=lambda row: row["rowindex"]
+ )
+ row = rows[0]
+ assert row["bool_col"] # True
+ assert row["bytes_col"] == b"Hello, World!"
+ assert row["date_col"] == datetime.date(2021, 7, 21)
+ assert row["datetime_col"] == datetime.datetime(2021, 7, 21, 11, 39, 45)
+ assert row["geography_col"] == "POINT(-122.0838511 37.3860517)"
+ assert row["int64_col"] == 123456789
+ assert row["numeric_col"] == decimal.Decimal("1.23456789")
+ assert row["bignumeric_col"] == decimal.Decimal("10.111213141516171819")
+ assert row["float64_col"] == 1.25
+ assert row["string_col"] == "Hello, World!"
+ assert row["time_col"] == datetime.time(11, 41, 43, 76160)
+ assert row["timestamp_col"] == datetime.datetime(
+ 2021, 7, 21, 17, 43, 43, 945289, tzinfo=datetime.timezone.utc
+ )
+
+ nullrow = rows[1]
+ for column, value in nullrow.items():
+ if column == "rowindex":
+ assert value == 1
+ else:
+ assert value is None
+
+
+def test_list_rows_scalars_extreme(
+ bigquery_client: bigquery.Client, scalars_extreme_table: str
+):
+ rows = sorted(
+ bigquery_client.list_rows(scalars_extreme_table),
+ key=lambda row: row["rowindex"],
+ )
+ row = rows[0]
+ assert row["bool_col"] # True
+ assert row["bytes_col"] == b"\r\n"
+ assert row["date_col"] == datetime.date(9999, 12, 31)
+ assert row["datetime_col"] == datetime.datetime(9999, 12, 31, 23, 59, 59, 999999)
+ assert row["geography_col"] == "POINT(-135 90)"
+ assert row["int64_col"] == 9223372036854775807
+ assert row["numeric_col"] == decimal.Decimal(f"9.{'9' * 37}E+28")
+ assert row["bignumeric_col"] == decimal.Decimal(f"9.{'9' * 75}E+37")
+ assert row["float64_col"] == float("Inf")
+ assert row["string_col"] == "Hello, World"
+ assert row["time_col"] == datetime.time(23, 59, 59, 999999)
+ assert row["timestamp_col"] == datetime.datetime(
+ 9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc
+ )
+
+ nullrow = rows[4]
+ for column, value in nullrow.items():
+ if column == "rowindex":
+ assert value == 4
+ else:
+ assert value is None
diff --git a/tests/system/test_pandas.py b/tests/system/test_pandas.py
index 821b375e1..93ce23481 100644
--- a/tests/system/test_pandas.py
+++ b/tests/system/test_pandas.py
@@ -24,10 +24,8 @@
import google.api_core.retry
import pkg_resources
import pytest
-import pytz
from google.cloud import bigquery
-from google.cloud.bigquery._pandas_helpers import _BIGNUMERIC_SUPPORT
from . import helpers
@@ -64,7 +62,7 @@ def test_load_table_from_dataframe_w_automatic_schema(bigquery_client, dataset_i
datetime.datetime(2012, 3, 14, 15, 16),
],
dtype="datetime64[ns]",
- ).dt.tz_localize(pytz.utc),
+ ).dt.tz_localize(datetime.timezone.utc),
),
(
"dt_col",
@@ -189,12 +187,11 @@ def test_load_table_from_dataframe_w_nulls(bigquery_client, dataset_id):
bigquery.SchemaField("geo_col", "GEOGRAPHY"),
bigquery.SchemaField("int_col", "INTEGER"),
bigquery.SchemaField("num_col", "NUMERIC"),
+ bigquery.SchemaField("bignum_col", "BIGNUMERIC"),
bigquery.SchemaField("str_col", "STRING"),
bigquery.SchemaField("time_col", "TIME"),
bigquery.SchemaField("ts_col", "TIMESTAMP"),
)
- if _BIGNUMERIC_SUPPORT:
- scalars_schema += (bigquery.SchemaField("bignum_col", "BIGNUMERIC"),)
table_schema = scalars_schema + (
# TODO: Array columns can't be read due to NULLABLE versus REPEATED
@@ -216,12 +213,11 @@ def test_load_table_from_dataframe_w_nulls(bigquery_client, dataset_id):
("geo_col", nulls),
("int_col", nulls),
("num_col", nulls),
+ ("bignum_col", nulls),
("str_col", nulls),
("time_col", nulls),
("ts_col", nulls),
]
- if _BIGNUMERIC_SUPPORT:
- df_data.append(("bignum_col", nulls))
df_data = collections.OrderedDict(df_data)
dataframe = pandas.DataFrame(df_data, columns=df_data.keys())
@@ -283,8 +279,6 @@ def test_load_table_from_dataframe_w_required(bigquery_client, dataset_id):
def test_load_table_from_dataframe_w_explicit_schema(bigquery_client, dataset_id):
# Schema with all scalar types.
- # TODO: Uploading DATETIME columns currently fails, thus that field type
- # is temporarily removed from the test.
# See:
# https://github.com/googleapis/python-bigquery/issues/61
# https://issuetracker.google.com/issues/151765076
@@ -292,17 +286,16 @@ def test_load_table_from_dataframe_w_explicit_schema(bigquery_client, dataset_id
bigquery.SchemaField("bool_col", "BOOLEAN"),
bigquery.SchemaField("bytes_col", "BYTES"),
bigquery.SchemaField("date_col", "DATE"),
- # bigquery.SchemaField("dt_col", "DATETIME"),
+ bigquery.SchemaField("dt_col", "DATETIME"),
bigquery.SchemaField("float_col", "FLOAT"),
bigquery.SchemaField("geo_col", "GEOGRAPHY"),
bigquery.SchemaField("int_col", "INTEGER"),
bigquery.SchemaField("num_col", "NUMERIC"),
+ bigquery.SchemaField("bignum_col", "BIGNUMERIC"),
bigquery.SchemaField("str_col", "STRING"),
bigquery.SchemaField("time_col", "TIME"),
bigquery.SchemaField("ts_col", "TIMESTAMP"),
)
- if _BIGNUMERIC_SUPPORT:
- scalars_schema += (bigquery.SchemaField("bignum_col", "BIGNUMERIC"),)
table_schema = scalars_schema + (
# TODO: Array columns can't be read due to NULLABLE versus REPEATED
@@ -318,14 +311,14 @@ def test_load_table_from_dataframe_w_explicit_schema(bigquery_client, dataset_id
("bool_col", [True, None, False]),
("bytes_col", [b"abc", None, b"def"]),
("date_col", [datetime.date(1, 1, 1), None, datetime.date(9999, 12, 31)]),
- # (
- # "dt_col",
- # [
- # datetime.datetime(1, 1, 1, 0, 0, 0),
- # None,
- # datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
- # ],
- # ),
+ (
+ "dt_col",
+ [
+ datetime.datetime(1, 1, 1, 0, 0, 0),
+ None,
+ datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
+ ],
+ ),
("float_col", [float("-inf"), float("nan"), float("inf")]),
(
"geo_col",
@@ -340,6 +333,14 @@ def test_load_table_from_dataframe_w_explicit_schema(bigquery_client, dataset_id
decimal.Decimal("99999999999999999999999999999.999999999"),
],
),
+ (
+ "bignum_col",
+ [
+ decimal.Decimal("-{d38}.{d38}".format(d38="9" * 38)),
+ None,
+ decimal.Decimal("{d38}.{d38}".format(d38="9" * 38)),
+ ],
+ ),
("str_col", ["abc", None, "def"]),
(
"time_col",
@@ -348,23 +349,14 @@ def test_load_table_from_dataframe_w_explicit_schema(bigquery_client, dataset_id
(
"ts_col",
[
- datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
None,
- datetime.datetime(9999, 12, 31, 23, 59, 59, 999999, tzinfo=pytz.utc),
+ datetime.datetime(
+ 9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc
+ ),
],
),
]
- if _BIGNUMERIC_SUPPORT:
- df_data.append(
- (
- "bignum_col",
- [
- decimal.Decimal("-{d38}.{d38}".format(d38="9" * 38)),
- None,
- decimal.Decimal("{d38}.{d38}".format(d38="9" * 38)),
- ],
- )
- )
df_data = collections.OrderedDict(df_data)
dataframe = pandas.DataFrame(df_data, dtype="object", columns=df_data.keys())
@@ -484,10 +476,10 @@ def test_load_table_from_dataframe_w_explicit_schema_source_format_csv(
(
"ts_col",
[
- datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
None,
datetime.datetime(
- 9999, 12, 31, 23, 59, 59, 999999, tzinfo=pytz.utc
+ 9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc
),
],
),
@@ -804,3 +796,190 @@ def test_list_rows_max_results_w_bqstorage(bigquery_client):
dataframe = row_iterator.to_dataframe(bqstorage_client=bqstorage_client)
assert len(dataframe.index) == 100
+
+
+def test_upload_time_and_datetime_56(bigquery_client, dataset_id):
+ df = pandas.DataFrame(
+ dict(
+ dt=[
+ datetime.datetime(2020, 1, 8, 8, 0, 0),
+ datetime.datetime(
+ 2020,
+ 1,
+ 8,
+ 8,
+ 0,
+ 0,
+ tzinfo=datetime.timezone(datetime.timedelta(hours=-7)),
+ ),
+ ],
+ t=[datetime.time(0, 0, 10, 100001), None],
+ )
+ )
+ table = f"{dataset_id}.test_upload_time_and_datetime"
+ bigquery_client.load_table_from_dataframe(df, table).result()
+ data = list(map(list, bigquery_client.list_rows(table)))
+ assert data == [
+ [
+ datetime.datetime(2020, 1, 8, 8, 0, tzinfo=datetime.timezone.utc),
+ datetime.time(0, 0, 10, 100001),
+ ],
+ [datetime.datetime(2020, 1, 8, 15, 0, tzinfo=datetime.timezone.utc), None],
+ ]
+
+ from google.cloud.bigquery import job, schema
+
+ table = f"{dataset_id}.test_upload_time_and_datetime_dt"
+ config = job.LoadJobConfig(
+ schema=[schema.SchemaField("dt", "DATETIME"), schema.SchemaField("t", "TIME")]
+ )
+
+ bigquery_client.load_table_from_dataframe(df, table, job_config=config).result()
+ data = list(map(list, bigquery_client.list_rows(table)))
+ assert data == [
+ [datetime.datetime(2020, 1, 8, 8, 0), datetime.time(0, 0, 10, 100001)],
+ [datetime.datetime(2020, 1, 8, 15, 0), None],
+ ]
+
+
+def test_to_dataframe_geography_as_objects(bigquery_client, dataset_id):
+ wkt = pytest.importorskip("shapely.wkt")
+ bigquery_client.query(
+ f"create table {dataset_id}.lake (name string, geog geography)"
+ ).result()
+ bigquery_client.query(
+ f"""
+ insert into {dataset_id}.lake (name, geog) values
+ ('foo', st_geogfromtext('point(0 0)')),
+ ('bar', st_geogfromtext('point(0 1)')),
+ ('baz', null)
+ """
+ ).result()
+ df = bigquery_client.query(
+ f"select * from {dataset_id}.lake order by name"
+ ).to_dataframe(geography_as_object=True)
+ assert list(df["name"]) == ["bar", "baz", "foo"]
+ assert df["geog"][0] == wkt.loads("point(0 1)")
+ assert pandas.isna(df["geog"][1])
+ assert df["geog"][2] == wkt.loads("point(0 0)")
+
+
+def test_to_geodataframe(bigquery_client, dataset_id):
+ geopandas = pytest.importorskip("geopandas")
+ from shapely import wkt
+
+ bigquery_client.query(
+ f"create table {dataset_id}.geolake (name string, geog geography)"
+ ).result()
+ bigquery_client.query(
+ f"""
+ insert into {dataset_id}.geolake (name, geog) values
+ ('foo', st_geogfromtext('point(0 0)')),
+ ('bar', st_geogfromtext('polygon((0 0, 1 0, 1 1, 0 0))')),
+ ('baz', null)
+ """
+ ).result()
+ df = bigquery_client.query(
+ f"select * from {dataset_id}.geolake order by name"
+ ).to_geodataframe()
+ assert df["geog"][0] == wkt.loads("polygon((0 0, 1 0, 1 1, 0 0))")
+ assert pandas.isna(df["geog"][1])
+ assert df["geog"][2] == wkt.loads("point(0 0)")
+ assert isinstance(df, geopandas.GeoDataFrame)
+ assert isinstance(df["geog"], geopandas.GeoSeries)
+ assert df.area[0] == 0.5
+ assert pandas.isna(df.area[1])
+ assert df.area[2] == 0.0
+ assert df.crs.srs == "EPSG:4326"
+ assert df.crs.name == "WGS 84"
+ assert df.geog.crs.srs == "EPSG:4326"
+ assert df.geog.crs.name == "WGS 84"
+
+
+def test_load_geodataframe(bigquery_client, dataset_id):
+ geopandas = pytest.importorskip("geopandas")
+ import pandas
+ from shapely import wkt
+ from google.cloud.bigquery.schema import SchemaField
+
+ df = geopandas.GeoDataFrame(
+ pandas.DataFrame(
+ dict(
+ name=["foo", "bar"],
+ geo1=[None, None],
+ geo2=[None, wkt.loads("Point(1 1)")],
+ )
+ ),
+ geometry="geo1",
+ )
+
+ table_id = f"{dataset_id}.lake_from_gp"
+ bigquery_client.load_table_from_dataframe(df, table_id).result()
+
+ table = bigquery_client.get_table(table_id)
+ assert table.schema == [
+ SchemaField("name", "STRING", "NULLABLE"),
+ SchemaField("geo1", "GEOGRAPHY", "NULLABLE"),
+ SchemaField("geo2", "GEOGRAPHY", "NULLABLE"),
+ ]
+ assert sorted(map(list, bigquery_client.list_rows(table_id))) == [
+ ["bar", None, "POINT(1 1)"],
+ ["foo", None, None],
+ ]
+
+
+def test_load_dataframe_w_shapely(bigquery_client, dataset_id):
+ wkt = pytest.importorskip("shapely.wkt")
+ from google.cloud.bigquery.schema import SchemaField
+
+ df = pandas.DataFrame(
+ dict(name=["foo", "bar"], geo=[None, wkt.loads("Point(1 1)")])
+ )
+
+ table_id = f"{dataset_id}.lake_from_shapes"
+ bigquery_client.load_table_from_dataframe(df, table_id).result()
+
+ table = bigquery_client.get_table(table_id)
+ assert table.schema == [
+ SchemaField("name", "STRING", "NULLABLE"),
+ SchemaField("geo", "GEOGRAPHY", "NULLABLE"),
+ ]
+ assert sorted(map(list, bigquery_client.list_rows(table_id))) == [
+ ["bar", "POINT(1 1)"],
+ ["foo", None],
+ ]
+
+ bigquery_client.load_table_from_dataframe(df, table_id).result()
+ assert sorted(map(list, bigquery_client.list_rows(table_id))) == [
+ ["bar", "POINT(1 1)"],
+ ["bar", "POINT(1 1)"],
+ ["foo", None],
+ ["foo", None],
+ ]
+
+
+def test_load_dataframe_w_wkb(bigquery_client, dataset_id):
+ wkt = pytest.importorskip("shapely.wkt")
+ from shapely import wkb
+ from google.cloud.bigquery.schema import SchemaField
+
+ df = pandas.DataFrame(
+ dict(name=["foo", "bar"], geo=[None, wkb.dumps(wkt.loads("Point(1 1)"))])
+ )
+
+ table_id = f"{dataset_id}.lake_from_wkb"
+ # We create the table first, to inform the interpretation of the wkb data
+ bigquery_client.query(
+ f"create table {table_id} (name string, geo GEOGRAPHY)"
+ ).result()
+ bigquery_client.load_table_from_dataframe(df, table_id).result()
+
+ table = bigquery_client.get_table(table_id)
+ assert table.schema == [
+ SchemaField("name", "STRING", "NULLABLE"),
+ SchemaField("geo", "GEOGRAPHY", "NULLABLE"),
+ ]
+ assert sorted(map(list, bigquery_client.list_rows(table_id))) == [
+ ["bar", "POINT(1 1)"],
+ ["foo", None],
+ ]
diff --git a/tests/unit/job/helpers.py b/tests/unit/job/helpers.py
index ea071c5ac..c792214e7 100644
--- a/tests/unit/job/helpers.py
+++ b/tests/unit/job/helpers.py
@@ -162,6 +162,7 @@ def _verifyInitialReadonlyProperties(self, job):
self.assertIsNone(job.created)
self.assertIsNone(job.started)
self.assertIsNone(job.ended)
+ self.assertIsNone(job.transaction_info)
# derived from resource['status']
self.assertIsNone(job.error_result)
diff --git a/tests/unit/job/test_base.py b/tests/unit/job/test_base.py
index 405ad6ee5..c3f7854e3 100644
--- a/tests/unit/job/test_base.py
+++ b/tests/unit/job/test_base.py
@@ -227,6 +227,20 @@ def test_script_statistics(self):
self.assertEqual(stack_frame.end_column, 14)
self.assertEqual(stack_frame.text, "QUERY TEXT")
+ def test_transaction_info(self):
+ from google.cloud.bigquery.job.base import TransactionInfo
+
+ client = _make_client(project=self.PROJECT)
+ job = self._make_one(self.JOB_ID, client)
+ assert job.transaction_info is None
+
+ statistics = job._properties["statistics"] = {}
+ assert job.transaction_info is None
+
+ statistics["transactionInfo"] = {"transactionId": "123-abc-xyz"}
+ assert isinstance(job.transaction_info, TransactionInfo)
+ assert job.transaction_info.transaction_id == "123-abc-xyz"
+
def test_num_child_jobs(self):
client = _make_client(project=self.PROJECT)
job = self._make_one(self.JOB_ID, client)
@@ -281,11 +295,11 @@ def test_user_email(self):
@staticmethod
def _datetime_and_millis():
import datetime
- import pytz
from google.cloud._helpers import _millis
now = datetime.datetime.utcnow().replace(
- microsecond=123000, tzinfo=pytz.UTC # stats timestamps have ms precision
+ microsecond=123000,
+ tzinfo=datetime.timezone.utc, # stats timestamps have ms precision
)
return now, _millis(now)
diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py
index 482f7f3af..d41370520 100644
--- a/tests/unit/job/test_query.py
+++ b/tests/unit/job/test_query.py
@@ -128,6 +128,18 @@ def _verify_dml_stats_resource_properties(self, job, resource):
else:
assert job.dml_stats is None
+ def _verify_transaction_info_resource_properties(self, job, resource):
+ resource_stats = resource.get("statistics", {})
+
+ if "transactionInfo" in resource_stats:
+ resource_transaction_info = resource_stats["transactionInfo"]
+ job_transaction_info = job.transaction_info
+ assert job_transaction_info.transaction_id == resource_transaction_info.get(
+ "transactionId"
+ )
+ else:
+ assert job.transaction_info is None
+
def _verify_configuration_properties(self, job, configuration):
if "dryRun" in configuration:
self.assertEqual(job.dry_run, configuration["dryRun"])
@@ -137,6 +149,7 @@ def _verify_configuration_properties(self, job, configuration):
def _verifyResourceProperties(self, job, resource):
self._verifyReadonlyResourceProperties(job, resource)
self._verify_dml_stats_resource_properties(job, resource)
+ self._verify_transaction_info_resource_properties(job, resource)
configuration = resource.get("configuration", {})
self._verify_configuration_properties(job, configuration)
@@ -325,6 +338,22 @@ def test_from_api_repr_with_dml_stats(self):
self.assertIs(job._client, client)
self._verifyResourceProperties(job, RESOURCE)
+ def test_from_api_repr_with_transaction_info(self):
+ self._setUpConstants()
+ client = _make_client(project=self.PROJECT)
+ RESOURCE = {
+ "id": self.JOB_ID,
+ "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+ "configuration": {"query": {"query": self.QUERY}},
+ "statistics": {"transactionInfo": {"transactionId": "1a2b-3c4d"}},
+ }
+ klass = self._get_target_class()
+
+ job = klass.from_api_repr(RESOURCE, client=client)
+
+ self.assertIs(job._client, client)
+ 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
diff --git a/tests/unit/job/test_query_pandas.py b/tests/unit/job/test_query_pandas.py
index c537802f4..b5af90c0b 100644
--- a/tests/unit/job/test_query_pandas.py
+++ b/tests/unit/job/test_query_pandas.py
@@ -23,6 +23,14 @@
import pandas
except (ImportError, AttributeError): # pragma: NO COVER
pandas = None
+try:
+ import shapely
+except (ImportError, AttributeError): # pragma: NO COVER
+ shapely = None
+try:
+ import geopandas
+except (ImportError, AttributeError): # pragma: NO COVER
+ geopandas = None
try:
import pyarrow
except (ImportError, AttributeError): # pragma: NO COVER
@@ -425,38 +433,41 @@ def test_to_arrow_w_tqdm_wo_query_plan():
result_patch_tqdm.assert_called()
-@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
-def test_to_dataframe():
+def _make_job(schema=(), rows=()):
from google.cloud.bigquery.job import QueryJob as target_class
begun_resource = _make_job_resource(job_type="query")
query_resource = {
"jobComplete": True,
"jobReference": begun_resource["jobReference"],
- "totalRows": "4",
+ "totalRows": str(len(rows)),
"schema": {
"fields": [
- {"name": "name", "type": "STRING", "mode": "NULLABLE"},
- {"name": "age", "type": "INTEGER", "mode": "NULLABLE"},
+ dict(name=field[0], type=field[1], mode=field[2]) for field in schema
]
},
}
- tabledata_resource = {
- "rows": [
- {"f": [{"v": "Phred Phlyntstone"}, {"v": "32"}]},
- {"f": [{"v": "Bharney Rhubble"}, {"v": "33"}]},
- {"f": [{"v": "Wylma Phlyntstone"}, {"v": "29"}]},
- {"f": [{"v": "Bhettye Rhubble"}, {"v": "27"}]},
- ]
- }
+ tabledata_resource = {"rows": [{"f": [{"v": v} for v in row]} for row in rows]}
done_resource = copy.deepcopy(begun_resource)
done_resource["status"] = {"state": "DONE"}
connection = _make_connection(
begun_resource, query_resource, done_resource, tabledata_resource
)
client = _make_client(connection=connection)
- job = target_class.from_api_repr(begun_resource, client)
+ return target_class.from_api_repr(begun_resource, client)
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_to_dataframe():
+ job = _make_job(
+ (("name", "STRING", "NULLABLE"), ("age", "INTEGER", "NULLABLE")),
+ (
+ ("Phred Phlyntstone", "32"),
+ ("Bharney Rhubble", "33"),
+ ("Wylma Phlyntstone", "29"),
+ ("Bhettye Rhubble", "27"),
+ ),
+ )
df = job.to_dataframe(create_bqstorage_client=False)
assert isinstance(df, pandas.DataFrame)
@@ -868,3 +879,94 @@ def test_to_dataframe_w_tqdm_max_results():
result_patch_tqdm.assert_called_with(
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=3
)
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@pytest.mark.skipif(shapely is None, reason="Requires `shapely`")
+def test_to_dataframe_geography_as_object():
+ job = _make_job(
+ (("name", "STRING", "NULLABLE"), ("geog", "GEOGRAPHY", "NULLABLE")),
+ (
+ ("Phred Phlyntstone", "Point(0 0)"),
+ ("Bharney Rhubble", "Point(0 1)"),
+ ("Wylma Phlyntstone", None),
+ ),
+ )
+ df = job.to_dataframe(create_bqstorage_client=False, geography_as_object=True)
+
+ assert isinstance(df, pandas.DataFrame)
+ assert len(df) == 3 # verify the number of rows
+ assert list(df) == ["name", "geog"] # verify the column names
+ assert [v.__class__.__name__ for v in df.geog] == [
+ "Point",
+ "Point",
+ "float",
+ ] # float because nan
+
+
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+def test_to_geodataframe():
+ job = _make_job(
+ (("name", "STRING", "NULLABLE"), ("geog", "GEOGRAPHY", "NULLABLE")),
+ (
+ ("Phred Phlyntstone", "Point(0 0)"),
+ ("Bharney Rhubble", "Point(0 1)"),
+ ("Wylma Phlyntstone", None),
+ ),
+ )
+ df = job.to_geodataframe(create_bqstorage_client=False)
+
+ assert isinstance(df, geopandas.GeoDataFrame)
+ assert len(df) == 3 # verify the number of rows
+ assert list(df) == ["name", "geog"] # verify the column names
+ assert [v.__class__.__name__ for v in df.geog] == [
+ "Point",
+ "Point",
+ "NoneType",
+ ] # float because nan
+ assert isinstance(df.geog, geopandas.GeoSeries)
+
+
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+@mock.patch("google.cloud.bigquery.job.query.wait_for_query")
+def test_query_job_to_geodataframe_delegation(wait_for_query):
+ """
+ QueryJob.to_geodataframe just delegates to RowIterator.to_geodataframe.
+
+ This test just demonstrates that. We don't need to test all the
+ variations, which are tested for RowIterator.
+ """
+ import numpy
+
+ job = _make_job()
+ bqstorage_client = object()
+ dtypes = dict(xxx=numpy.dtype("int64"))
+ progress_bar_type = "normal"
+ create_bqstorage_client = False
+ date_as_object = False
+ max_results = 42
+ geography_column = "g"
+
+ df = job.to_geodataframe(
+ bqstorage_client=bqstorage_client,
+ dtypes=dtypes,
+ progress_bar_type=progress_bar_type,
+ create_bqstorage_client=create_bqstorage_client,
+ date_as_object=date_as_object,
+ max_results=max_results,
+ geography_column=geography_column,
+ )
+
+ wait_for_query.assert_called_once_with(
+ job, progress_bar_type, max_results=max_results
+ )
+ row_iterator = wait_for_query.return_value
+ row_iterator.to_geodataframe.assert_called_once_with(
+ bqstorage_client=bqstorage_client,
+ dtypes=dtypes,
+ progress_bar_type=progress_bar_type,
+ create_bqstorage_client=create_bqstorage_client,
+ date_as_object=date_as_object,
+ geography_column=geography_column,
+ )
+ assert df is row_iterator.to_geodataframe.return_value
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 0ba671cd9..a9b0ae21f 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -36,13 +36,16 @@
# Mock out pyarrow when missing, because methods from pyarrow.types are
# used in test parameterization.
pyarrow = mock.Mock()
+try:
+ import geopandas
+except ImportError: # pragma: NO COVER
+ geopandas = None
+
import pytest
-import pytz
from google import api_core
from google.cloud.bigquery import _helpers
from google.cloud.bigquery import schema
-from google.cloud.bigquery._pandas_helpers import _BIGNUMERIC_SUPPORT
try:
from google.cloud import bigquery_storage
@@ -60,11 +63,6 @@
PANDAS_INSTALLED_VERSION = pkg_resources.parse_version("0.0.0")
-skip_if_no_bignumeric = pytest.mark.skipif(
- not _BIGNUMERIC_SUPPORT, reason="BIGNUMERIC support requires pyarrow>=3.0.0",
-)
-
-
@pytest.fixture
def module_under_test():
from google.cloud.bigquery import _pandas_helpers
@@ -153,9 +151,7 @@ def test_all_():
("FLOAT", "NULLABLE", pyarrow.types.is_float64),
("FLOAT64", "NULLABLE", pyarrow.types.is_float64),
("NUMERIC", "NULLABLE", is_numeric),
- pytest.param(
- "BIGNUMERIC", "NULLABLE", is_bignumeric, marks=skip_if_no_bignumeric,
- ),
+ ("BIGNUMERIC", "NULLABLE", is_bignumeric),
("BOOLEAN", "NULLABLE", pyarrow.types.is_boolean),
("BOOL", "NULLABLE", pyarrow.types.is_boolean),
("TIMESTAMP", "NULLABLE", is_timestamp),
@@ -234,11 +230,10 @@ def test_all_():
"REPEATED",
all_(pyarrow.types.is_list, lambda type_: is_numeric(type_.value_type)),
),
- pytest.param(
+ (
"BIGNUMERIC",
"REPEATED",
all_(pyarrow.types.is_list, lambda type_: is_bignumeric(type_.value_type)),
- marks=skip_if_no_bignumeric,
),
(
"BOOLEAN",
@@ -312,6 +307,7 @@ def test_bq_to_arrow_data_type_w_struct(module_under_test, bq_type):
schema.SchemaField("field05", "FLOAT"),
schema.SchemaField("field06", "FLOAT64"),
schema.SchemaField("field07", "NUMERIC"),
+ schema.SchemaField("field08", "BIGNUMERIC"),
schema.SchemaField("field09", "BOOLEAN"),
schema.SchemaField("field10", "BOOL"),
schema.SchemaField("field11", "TIMESTAMP"),
@@ -321,9 +317,6 @@ def test_bq_to_arrow_data_type_w_struct(module_under_test, bq_type):
schema.SchemaField("field15", "GEOGRAPHY"),
)
- if _BIGNUMERIC_SUPPORT:
- fields += (schema.SchemaField("field08", "BIGNUMERIC"),)
-
field = schema.SchemaField("ignored_name", bq_type, mode="NULLABLE", fields=fields)
actual = module_under_test.bq_to_arrow_data_type(field)
@@ -335,6 +328,7 @@ def test_bq_to_arrow_data_type_w_struct(module_under_test, bq_type):
pyarrow.field("field05", pyarrow.float64()),
pyarrow.field("field06", pyarrow.float64()),
pyarrow.field("field07", module_under_test.pyarrow_numeric()),
+ pyarrow.field("field08", module_under_test.pyarrow_bignumeric()),
pyarrow.field("field09", pyarrow.bool_()),
pyarrow.field("field10", pyarrow.bool_()),
pyarrow.field("field11", module_under_test.pyarrow_timestamp()),
@@ -343,8 +337,6 @@ def test_bq_to_arrow_data_type_w_struct(module_under_test, bq_type):
pyarrow.field("field14", module_under_test.pyarrow_datetime()),
pyarrow.field("field15", pyarrow.string()),
)
- if _BIGNUMERIC_SUPPORT:
- expected += (pyarrow.field("field08", module_under_test.pyarrow_bignumeric()),)
expected = pyarrow.struct(expected)
assert pyarrow.types.is_struct(actual)
@@ -363,6 +355,7 @@ def test_bq_to_arrow_data_type_w_array_struct(module_under_test, bq_type):
schema.SchemaField("field05", "FLOAT"),
schema.SchemaField("field06", "FLOAT64"),
schema.SchemaField("field07", "NUMERIC"),
+ schema.SchemaField("field08", "BIGNUMERIC"),
schema.SchemaField("field09", "BOOLEAN"),
schema.SchemaField("field10", "BOOL"),
schema.SchemaField("field11", "TIMESTAMP"),
@@ -372,9 +365,6 @@ def test_bq_to_arrow_data_type_w_array_struct(module_under_test, bq_type):
schema.SchemaField("field15", "GEOGRAPHY"),
)
- if _BIGNUMERIC_SUPPORT:
- fields += (schema.SchemaField("field08", "BIGNUMERIC"),)
-
field = schema.SchemaField("ignored_name", bq_type, mode="REPEATED", fields=fields)
actual = module_under_test.bq_to_arrow_data_type(field)
@@ -386,6 +376,7 @@ def test_bq_to_arrow_data_type_w_array_struct(module_under_test, bq_type):
pyarrow.field("field05", pyarrow.float64()),
pyarrow.field("field06", pyarrow.float64()),
pyarrow.field("field07", module_under_test.pyarrow_numeric()),
+ pyarrow.field("field08", module_under_test.pyarrow_bignumeric()),
pyarrow.field("field09", pyarrow.bool_()),
pyarrow.field("field10", pyarrow.bool_()),
pyarrow.field("field11", module_under_test.pyarrow_timestamp()),
@@ -394,8 +385,6 @@ def test_bq_to_arrow_data_type_w_array_struct(module_under_test, bq_type):
pyarrow.field("field14", module_under_test.pyarrow_datetime()),
pyarrow.field("field15", pyarrow.string()),
)
- if _BIGNUMERIC_SUPPORT:
- expected += (pyarrow.field("field08", module_under_test.pyarrow_bignumeric()),)
expected_value_type = pyarrow.struct(expected)
assert pyarrow.types.is_list(actual)
@@ -441,7 +430,7 @@ def test_bq_to_arrow_data_type_w_struct_unknown_subfield(module_under_test):
decimal.Decimal("999.123456789"),
],
),
- pytest.param(
+ (
"BIGNUMERIC",
[
decimal.Decimal("-{d38}.{d38}".format(d38="9" * 38)),
@@ -449,17 +438,18 @@ def test_bq_to_arrow_data_type_w_struct_unknown_subfield(module_under_test):
decimal.Decimal("{d38}.{d38}".format(d38="9" * 38)),
decimal.Decimal("3.141592653589793238462643383279"),
],
- marks=skip_if_no_bignumeric,
),
("BOOLEAN", [True, None, False, None]),
("BOOL", [False, None, True, None]),
(
"TIMESTAMP",
[
- datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
None,
- datetime.datetime(9999, 12, 31, 23, 59, 59, 999999, tzinfo=pytz.utc),
- datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(
+ 9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc
+ ),
+ datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
],
),
(
@@ -599,6 +589,60 @@ def test_bq_to_arrow_array_w_special_floats(module_under_test):
assert roundtrip[3] is None
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_bq_to_arrow_array_w_geography_dtype(module_under_test):
+ from shapely import wkb, wkt
+
+ bq_field = schema.SchemaField("field_name", "GEOGRAPHY")
+
+ series = geopandas.GeoSeries([None, wkt.loads("point(0 0)")])
+ array = module_under_test.bq_to_arrow_array(series, bq_field)
+ # The result is binary, because we use wkb format
+ assert array.type == pyarrow.binary()
+ assert array.to_pylist() == [None, wkb.dumps(series[1])]
+
+ # All na:
+ series = geopandas.GeoSeries([None, None])
+ array = module_under_test.bq_to_arrow_array(series, bq_field)
+ assert array.type == pyarrow.string()
+ assert array.to_pylist() == list(series)
+
+
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_bq_to_arrow_array_w_geography_type_shapely_data(module_under_test):
+ from shapely import wkb, wkt
+
+ bq_field = schema.SchemaField("field_name", "GEOGRAPHY")
+
+ series = pandas.Series([None, wkt.loads("point(0 0)")])
+ array = module_under_test.bq_to_arrow_array(series, bq_field)
+ # The result is binary, because we use wkb format
+ assert array.type == pyarrow.binary()
+ assert array.to_pylist() == [None, wkb.dumps(series[1])]
+
+ # All na:
+ series = pandas.Series([None, None])
+ array = module_under_test.bq_to_arrow_array(series, bq_field)
+ assert array.type == pyarrow.string()
+ assert array.to_pylist() == list(series)
+
+
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_bq_to_arrow_array_w_geography_type_wkb_data(module_under_test):
+ from shapely import wkb, wkt
+
+ bq_field = schema.SchemaField("field_name", "GEOGRAPHY")
+
+ series = pandas.Series([None, wkb.dumps(wkt.loads("point(0 0)"))])
+ array = module_under_test.bq_to_arrow_array(series, bq_field)
+ # The result is binary, because we use wkb format
+ assert array.type == pyarrow.binary()
+ assert array.to_pylist() == list(series)
+
+
@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
def test_bq_to_arrow_schema_w_unknown_type(module_under_test):
fields = (
@@ -938,6 +982,7 @@ def test_dataframe_to_arrow_with_required_fields(module_under_test):
schema.SchemaField("field05", "FLOAT", mode="REQUIRED"),
schema.SchemaField("field06", "FLOAT64", mode="REQUIRED"),
schema.SchemaField("field07", "NUMERIC", mode="REQUIRED"),
+ schema.SchemaField("field08", "BIGNUMERIC", mode="REQUIRED"),
schema.SchemaField("field09", "BOOLEAN", mode="REQUIRED"),
schema.SchemaField("field10", "BOOL", mode="REQUIRED"),
schema.SchemaField("field11", "TIMESTAMP", mode="REQUIRED"),
@@ -946,8 +991,6 @@ def test_dataframe_to_arrow_with_required_fields(module_under_test):
schema.SchemaField("field14", "DATETIME", mode="REQUIRED"),
schema.SchemaField("field15", "GEOGRAPHY", mode="REQUIRED"),
)
- if _BIGNUMERIC_SUPPORT:
- bq_schema += (schema.SchemaField("field08", "BIGNUMERIC", mode="REQUIRED"),)
data = {
"field01": ["hello", "world"],
@@ -957,11 +1000,15 @@ def test_dataframe_to_arrow_with_required_fields(module_under_test):
"field05": [1.25, 9.75],
"field06": [-1.75, -3.5],
"field07": [decimal.Decimal("1.2345"), decimal.Decimal("6.7891")],
+ "field08": [
+ decimal.Decimal("-{d38}.{d38}".format(d38="9" * 38)),
+ decimal.Decimal("{d38}.{d38}".format(d38="9" * 38)),
+ ],
"field09": [True, False],
"field10": [False, True],
"field11": [
- datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
- datetime.datetime(2012, 12, 21, 9, 7, 42, tzinfo=pytz.utc),
+ datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2012, 12, 21, 9, 7, 42, tzinfo=datetime.timezone.utc),
],
"field12": [datetime.date(9999, 12, 31), datetime.date(1970, 1, 1)],
"field13": [datetime.time(23, 59, 59, 999999), datetime.time(12, 0, 0)],
@@ -971,11 +1018,6 @@ def test_dataframe_to_arrow_with_required_fields(module_under_test):
],
"field15": ["POINT(30 10)", "POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))"],
}
- if _BIGNUMERIC_SUPPORT:
- data["field08"] = [
- decimal.Decimal("-{d38}.{d38}".format(d38="9" * 38)),
- decimal.Decimal("{d38}.{d38}".format(d38="9" * 38)),
- ]
dataframe = pandas.DataFrame(data)
arrow_table = module_under_test.dataframe_to_arrow(dataframe, bq_schema)
@@ -1175,6 +1217,28 @@ def test_dataframe_to_bq_schema_pyarrow_fallback_fails(module_under_test):
assert "struct_field" in str(expected_warnings[0])
+@pytest.mark.skipif(geopandas is None, reason="Requires `geopandas`")
+def test_dataframe_to_bq_schema_geography(module_under_test):
+ from shapely import wkt
+
+ df = geopandas.GeoDataFrame(
+ pandas.DataFrame(
+ dict(
+ name=["foo", "bar"],
+ geo1=[None, None],
+ geo2=[None, wkt.loads("Point(1 1)")],
+ )
+ ),
+ geometry="geo1",
+ )
+ bq_schema = module_under_test.dataframe_to_bq_schema(df, [])
+ assert bq_schema == (
+ schema.SchemaField("name", "STRING"),
+ schema.SchemaField("geo1", "GEOGRAPHY"),
+ schema.SchemaField("geo2", "GEOGRAPHY"),
+ )
+
+
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
def test_augment_schema_type_detection_succeeds(module_under_test):
@@ -1210,11 +1274,8 @@ def test_augment_schema_type_detection_succeeds(module_under_test):
schema.SchemaField("bytes_field", field_type=None, mode="NULLABLE"),
schema.SchemaField("string_field", field_type=None, mode="NULLABLE"),
schema.SchemaField("numeric_field", field_type=None, mode="NULLABLE"),
+ schema.SchemaField("bignumeric_field", field_type=None, mode="NULLABLE"),
)
- if _BIGNUMERIC_SUPPORT:
- current_schema += (
- schema.SchemaField("bignumeric_field", field_type=None, mode="NULLABLE"),
- )
with warnings.catch_warnings(record=True) as warned:
augmented_schema = module_under_test.augment_schema(dataframe, current_schema)
@@ -1236,13 +1297,10 @@ def test_augment_schema_type_detection_succeeds(module_under_test):
schema.SchemaField("bytes_field", field_type="BYTES", mode="NULLABLE"),
schema.SchemaField("string_field", field_type="STRING", mode="NULLABLE"),
schema.SchemaField("numeric_field", field_type="NUMERIC", mode="NULLABLE"),
+ schema.SchemaField(
+ "bignumeric_field", field_type="BIGNUMERIC", mode="NULLABLE"
+ ),
)
- if _BIGNUMERIC_SUPPORT:
- expected_schema += (
- schema.SchemaField(
- "bignumeric_field", field_type="BIGNUMERIC", mode="NULLABLE"
- ),
- )
by_name = operator.attrgetter("name")
assert sorted(augmented_schema, key=by_name) == sorted(expected_schema, key=by_name)
@@ -1577,3 +1635,22 @@ def test_download_dataframe_row_iterator_dict_sequence_schema(module_under_test)
def test_table_data_listpage_to_dataframe_skips_stop_iteration(module_under_test):
dataframe = module_under_test._row_iterator_page_to_dataframe([], [], {})
assert isinstance(dataframe, pandas.DataFrame)
+
+
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_bq_to_arrow_field_type_override(module_under_test):
+ # When loading pandas data, we may need to override the type
+ # decision based on data contents, because GEOGRAPHY data can be
+ # stored as either text or binary.
+
+ assert (
+ module_under_test.bq_to_arrow_field(schema.SchemaField("g", "GEOGRAPHY")).type
+ == pyarrow.string()
+ )
+
+ assert (
+ module_under_test.bq_to_arrow_field(
+ schema.SchemaField("g", "GEOGRAPHY"), pyarrow.binary(),
+ ).type
+ == pyarrow.binary()
+ )
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 535685511..e9204f1de 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -30,7 +30,6 @@
import packaging
import requests
import pytest
-import pytz
import pkg_resources
try:
@@ -57,6 +56,7 @@
import google.cloud._helpers
from google.cloud import bigquery_v2
from google.cloud.bigquery.dataset import DatasetReference
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
try:
from google.cloud import bigquery_storage
@@ -368,7 +368,7 @@ def test__get_query_results_miss_w_client_location(self):
method="GET",
path="/projects/PROJECT/queries/nothere",
query_params={"maxResults": 0, "location": self.LOCATION},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test__get_query_results_hit(self):
@@ -429,7 +429,9 @@ def test_get_service_account_email_w_alternate_project(self):
service_account_email = client.get_service_account_email(project=project)
final_attributes.assert_called_once_with({"path": path}, client, None)
- conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None)
+ conn.api_request.assert_called_once_with(
+ method="GET", path=path, timeout=DEFAULT_TIMEOUT
+ )
self.assertEqual(service_account_email, email)
def test_get_service_account_email_w_custom_retry(self):
@@ -772,7 +774,7 @@ def test_create_routine_w_conflict(self):
}
}
conn.api_request.assert_called_once_with(
- method="POST", path=path, data=resource, timeout=None,
+ method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,
)
@unittest.skipIf(opentelemetry is None, "Requires `opentelemetry`")
@@ -808,7 +810,7 @@ def test_span_status_is_set(self):
}
}
conn.api_request.assert_called_once_with(
- method="POST", path=path, data=resource, timeout=None,
+ method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,
)
def test_create_routine_w_conflict_exists_ok(self):
@@ -844,11 +846,13 @@ def test_create_routine_w_conflict_exists_ok(self):
self.assertEqual(actual_routine.routine_id, "minimal_routine")
conn.api_request.assert_has_calls(
[
- mock.call(method="POST", path=path, data=resource, timeout=None,),
+ mock.call(
+ method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,
+ ),
mock.call(
method="GET",
path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine",
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
]
)
@@ -924,7 +928,7 @@ def test_create_table_w_custom_property(self):
"newAlphaProperty": "unreleased property",
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got._properties["newAlphaProperty"], "unreleased property")
self.assertEqual(got.table_id, self.TABLE_ID)
@@ -965,7 +969,7 @@ def test_create_table_w_encryption_configuration(self):
"labels": {},
"encryptionConfiguration": {"kmsKeyName": self.KMS_KEY_NAME},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
@@ -1001,7 +1005,7 @@ def test_create_table_w_day_partition_and_expire(self):
"timePartitioning": {"type": "DAY", "expirationMs": "100"},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(table.time_partitioning.type_, "DAY")
self.assertEqual(table.time_partitioning.expiration_ms, 100)
@@ -1082,7 +1086,7 @@ def test_create_table_w_schema_and_query(self):
"view": {"query": query, "useLegacySql": False},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
self.assertEqual(got.project, self.PROJECT)
@@ -1137,7 +1141,7 @@ def test_create_table_w_external(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
self.assertEqual(got.project, self.PROJECT)
@@ -1176,7 +1180,7 @@ def test_create_table_w_reference(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
@@ -1210,7 +1214,7 @@ def test_create_table_w_fully_qualified_string(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
@@ -1242,7 +1246,7 @@ def test_create_table_w_string(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(got.table_id, self.TABLE_ID)
@@ -1277,7 +1281,7 @@ def test_create_table_alreadyexists_w_exists_ok_false(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_create_table_alreadyexists_w_exists_ok_true(self):
@@ -1320,9 +1324,9 @@ def test_create_table_alreadyexists_w_exists_ok_true(self):
},
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
- mock.call(method="GET", path=get_path, timeout=None),
+ mock.call(method="GET", path=get_path, timeout=DEFAULT_TIMEOUT),
]
)
@@ -1395,7 +1399,7 @@ def test_get_model_w_string(self):
final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)
conn.api_request.assert_called_once_with(
- method="GET", path="/%s" % path, timeout=None
+ method="GET", path="/%s" % path, timeout=DEFAULT_TIMEOUT
)
self.assertEqual(got.model_id, self.MODEL_ID)
@@ -1504,7 +1508,7 @@ def test_get_table_sets_user_agent(self):
"User-Agent": expected_user_agent,
},
data=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertIn("my-application/1.2.3", expected_user_agent)
@@ -1847,7 +1851,7 @@ def test_update_dataset_w_custom_property(self):
data={"newAlphaProperty": "unreleased property"},
path=path,
headers=None,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(dataset.dataset_id, self.DS_ID)
@@ -2137,7 +2141,7 @@ def test_update_table_w_custom_property(self):
path="/%s" % path,
data={"newAlphaProperty": "unreleased property"},
headers=None,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(
updated_table._properties["newAlphaProperty"], "unreleased property"
@@ -2172,7 +2176,7 @@ def test_update_table_only_use_legacy_sql(self):
path="/%s" % path,
data={"view": {"useLegacySql": True}},
headers=None,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(updated_table.view_use_legacy_sql, table.view_use_legacy_sql)
@@ -2270,7 +2274,7 @@ def test_update_table_w_query(self):
"schema": schema_resource,
},
headers=None,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_update_table_w_schema_None(self):
@@ -2399,7 +2403,7 @@ def test_delete_job_metadata_not_found(self):
method="DELETE",
path="/projects/client-proj/jobs/my-job/delete",
query_params={"location": "client-loc"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_delete_job_metadata_with_id(self):
@@ -2413,7 +2417,7 @@ def test_delete_job_metadata_with_id(self):
method="DELETE",
path="/projects/param-proj/jobs/my-job/delete",
query_params={"location": "param-loc"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_delete_job_metadata_with_resource(self):
@@ -2438,7 +2442,7 @@ def test_delete_job_metadata_with_resource(self):
method="DELETE",
path="/projects/job-based-proj/jobs/query_job/delete",
query_params={"location": "us-east1"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_delete_model(self):
@@ -2493,7 +2497,9 @@ def test_delete_model_w_not_found_ok_false(self):
with self.assertRaises(google.api_core.exceptions.NotFound):
client.delete_model("{}.{}".format(self.DS_ID, self.MODEL_ID))
- conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
+ conn.api_request.assert_called_with(
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT
+ )
def test_delete_model_w_not_found_ok_true(self):
path = "/projects/{}/datasets/{}/models/{}".format(
@@ -2514,7 +2520,9 @@ def test_delete_model_w_not_found_ok_true(self):
final_attributes.assert_called_once_with({"path": path}, client, None)
- conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
+ conn.api_request.assert_called_with(
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT
+ )
def test_delete_routine(self):
from google.cloud.bigquery.routine import Routine
@@ -2568,7 +2576,7 @@ def test_delete_routine_w_not_found_ok_false(self):
final_attributes.assert_called_once_with({"path": path}, client, None)
conn.api_request.assert_called_with(
- method="DELETE", path=path, timeout=None,
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT,
)
def test_delete_routine_w_not_found_ok_true(self):
@@ -2590,7 +2598,7 @@ def test_delete_routine_w_not_found_ok_true(self):
final_attributes.assert_called_once_with({"path": path}, client, None)
conn.api_request.assert_called_with(
- method="DELETE", path=path, timeout=None,
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT,
)
def test_delete_table(self):
@@ -2654,7 +2662,9 @@ def test_delete_table_w_not_found_ok_false(self):
final_attributes.assert_called_once_with({"path": path}, client, None)
- conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
+ conn.api_request.assert_called_with(
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT
+ )
def test_delete_table_w_not_found_ok_true(self):
path = "/projects/{}/datasets/{}/tables/{}".format(
@@ -2676,7 +2686,9 @@ def test_delete_table_w_not_found_ok_true(self):
final_attributes.assert_called_once_with({"path": path}, client, None)
- conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
+ conn.api_request.assert_called_with(
+ method="DELETE", path=path, timeout=DEFAULT_TIMEOUT
+ )
def _create_job_helper(self, job_config):
from google.cloud.bigquery import _helpers
@@ -2698,7 +2710,7 @@ def _create_job_helper(self, job_config):
method="POST",
path="/projects/%s/jobs" % self.PROJECT,
data=RESOURCE,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_create_job_load_config(self):
@@ -2847,7 +2859,7 @@ def test_create_job_query_config_w_rateLimitExceeded_error(self):
method="POST",
path="/projects/PROJECT/jobs",
data=data_without_destination,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
)
@@ -2887,7 +2899,7 @@ def test_get_job_miss_w_explict_project(self):
method="GET",
path="/projects/OTHER_PROJECT/jobs/NONESUCH",
query_params={"projection": "full"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_get_job_miss_w_client_location(self):
@@ -2905,7 +2917,7 @@ def test_get_job_miss_w_client_location(self):
method="GET",
path="/projects/client-proj/jobs/NONESUCH",
query_params={"projection": "full", "location": "client-loc"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_get_job_hit_w_timeout(self):
@@ -2974,7 +2986,7 @@ def test_cancel_job_miss_w_explict_project(self):
method="POST",
path="/projects/OTHER_PROJECT/jobs/NONESUCH/cancel",
query_params={"projection": "full", "location": self.LOCATION},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_cancel_job_miss_w_client_location(self):
@@ -2993,7 +3005,7 @@ def test_cancel_job_miss_w_client_location(self):
method="POST",
path="/projects/OTHER_PROJECT/jobs/NONESUCH/cancel",
query_params={"projection": "full", "location": self.LOCATION},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_cancel_job_hit(self):
@@ -3029,7 +3041,7 @@ def test_cancel_job_hit(self):
method="POST",
path="/projects/job-based-proj/jobs/query_job/cancel",
query_params={"projection": "full", "location": "asia-northeast1"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_cancel_job_w_timeout(self):
@@ -3155,7 +3167,7 @@ def test_load_table_from_uri_w_explicit_project(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_load_table_from_uri_w_client_location(self):
@@ -3199,7 +3211,7 @@ def test_load_table_from_uri_w_client_location(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_load_table_from_uri_w_invalid_job_config(self):
@@ -3487,7 +3499,7 @@ def test_copy_table_w_multiple_sources(self):
method="POST",
path="/projects/%s/jobs" % self.PROJECT,
data=expected_resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertIsInstance(job, CopyJob)
self.assertIs(job._client, client)
@@ -3549,7 +3561,7 @@ def test_copy_table_w_explicit_project(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_copy_table_w_client_location(self):
@@ -3599,7 +3611,7 @@ def test_copy_table_w_client_location(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_copy_table_w_source_strings(self):
@@ -3692,7 +3704,7 @@ def test_copy_table_w_valid_job_config(self):
method="POST",
path="/projects/%s/jobs" % self.PROJECT,
data=RESOURCE,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertIsInstance(job._configuration, CopyJobConfig)
@@ -3798,7 +3810,7 @@ def test_extract_table_w_explicit_project(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_extract_table_w_client_location(self):
@@ -3842,7 +3854,7 @@ def test_extract_table_w_client_location(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_extract_table_generated_job_id(self):
@@ -3885,7 +3897,7 @@ def test_extract_table_generated_job_id(self):
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/projects/PROJECT/jobs")
self.assertIsInstance(req["data"]["jobReference"]["jobId"], str)
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
# Check the job resource.
self.assertIsInstance(job, ExtractJob)
@@ -3930,7 +3942,7 @@ def test_extract_table_w_destination_uris(self):
_, req = conn.api_request.call_args
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/projects/PROJECT/jobs")
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
# Check the job resource.
self.assertIsInstance(job, ExtractJob)
@@ -4100,7 +4112,7 @@ def test_query_defaults(self):
_, req = conn.api_request.call_args
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/projects/PROJECT/jobs")
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
sent = req["data"]
self.assertIsInstance(sent["jobReference"]["jobId"], str)
sent_config = sent["configuration"]["query"]
@@ -4153,7 +4165,7 @@ def test_query_w_explicit_project(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_query_w_explicit_job_config(self):
@@ -4209,7 +4221,10 @@ def test_query_w_explicit_job_config(self):
# Check that query actually starts the job.
conn.api_request.assert_called_once_with(
- method="POST", path="/projects/PROJECT/jobs", data=resource, timeout=None
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=resource,
+ timeout=DEFAULT_TIMEOUT,
)
# the original config object should not have been modified
@@ -4253,7 +4268,10 @@ def test_query_preserving_explicit_job_config(self):
# Check that query actually starts the job.
conn.api_request.assert_called_once_with(
- method="POST", path="/projects/PROJECT/jobs", data=resource, timeout=None
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=resource,
+ timeout=DEFAULT_TIMEOUT,
)
# the original config object should not have been modified
@@ -4305,7 +4323,10 @@ def test_query_preserving_explicit_default_job_config(self):
# Check that query actually starts the job.
conn.api_request.assert_called_once_with(
- method="POST", path="/projects/PROJECT/jobs", data=resource, timeout=None
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=resource,
+ timeout=DEFAULT_TIMEOUT,
)
# the original default config object should not have been modified
@@ -4390,7 +4411,10 @@ def test_query_w_explicit_job_config_override(self):
# Check that query actually starts the job.
conn.api_request.assert_called_once_with(
- method="POST", path="/projects/PROJECT/jobs", data=resource, timeout=None
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=resource,
+ timeout=DEFAULT_TIMEOUT,
)
def test_query_w_client_default_config_no_incoming(self):
@@ -4431,7 +4455,10 @@ def test_query_w_client_default_config_no_incoming(self):
# Check that query actually starts the job.
conn.api_request.assert_called_once_with(
- method="POST", path="/projects/PROJECT/jobs", data=resource, timeout=None
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=resource,
+ timeout=DEFAULT_TIMEOUT,
)
def test_query_w_invalid_default_job_config(self):
@@ -4476,7 +4503,7 @@ def test_query_w_client_location(self):
method="POST",
path="/projects/other-project/jobs",
data=resource,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_query_detect_location(self):
@@ -4547,7 +4574,7 @@ def test_query_w_udf_resources(self):
_, req = conn.api_request.call_args
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/projects/PROJECT/jobs")
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
sent = req["data"]
self.assertIsInstance(sent["jobReference"]["jobId"], str)
sent_config = sent["configuration"]["query"]
@@ -4603,7 +4630,7 @@ def test_query_w_query_parameters(self):
_, req = conn.api_request.call_args
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/projects/PROJECT/jobs")
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
sent = req["data"]
self.assertEqual(sent["jobReference"]["jobId"], JOB)
sent_config = sent["configuration"]["query"]
@@ -4795,7 +4822,7 @@ def _row_data(row):
self.assertEqual(req["method"], "POST")
self.assertEqual(req["path"], "/%s" % PATH)
self.assertEqual(req["data"], SENT)
- self.assertIsNone(req["timeout"])
+ self.assertEqual(req["timeout"], DEFAULT_TIMEOUT)
def test_insert_rows_w_list_of_dictionaries(self):
import datetime
@@ -4863,7 +4890,7 @@ def _row_data(row):
self.assertEqual(len(errors), 0)
conn.api_request.assert_called_once_with(
- method="POST", path="/%s" % PATH, data=SENT, timeout=None
+ method="POST", path="/%s" % PATH, data=SENT, timeout=DEFAULT_TIMEOUT
)
def test_insert_rows_w_list_of_Rows(self):
@@ -4908,7 +4935,7 @@ def _row_data(row):
self.assertEqual(len(errors), 0)
conn.api_request.assert_called_once_with(
- method="POST", path="/%s" % PATH, data=SENT, timeout=None
+ method="POST", path="/%s" % PATH, data=SENT, timeout=DEFAULT_TIMEOUT
)
def test_insert_rows_w_skip_invalid_and_ignore_unknown(self):
@@ -4985,7 +5012,7 @@ def _row_data(row):
errors[0]["errors"][0], RESPONSE["insertErrors"][0]["errors"][0]
)
conn.api_request.assert_called_once_with(
- method="POST", path="/%s" % PATH, data=SENT, timeout=None
+ method="POST", path="/%s" % PATH, data=SENT, timeout=DEFAULT_TIMEOUT
)
def test_insert_rows_w_repeated_fields(self):
@@ -5018,16 +5045,24 @@ def test_insert_rows_w_repeated_fields(self):
(
12,
[
- datetime.datetime(2018, 12, 1, 12, 0, 0, tzinfo=pytz.utc),
- datetime.datetime(2018, 12, 1, 13, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(
+ 2018, 12, 1, 12, 0, 0, tzinfo=datetime.timezone.utc
+ ),
+ datetime.datetime(
+ 2018, 12, 1, 13, 0, 0, tzinfo=datetime.timezone.utc
+ ),
],
[1.25, 2.5],
),
{
"score": 13,
"times": [
- datetime.datetime(2018, 12, 2, 12, 0, 0, tzinfo=pytz.utc),
- datetime.datetime(2018, 12, 2, 13, 0, 0, tzinfo=pytz.utc),
+ datetime.datetime(
+ 2018, 12, 2, 12, 0, 0, tzinfo=datetime.timezone.utc
+ ),
+ datetime.datetime(
+ 2018, 12, 2, 13, 0, 0, tzinfo=datetime.timezone.utc
+ ),
],
"distances": [-1.25, -2.5],
},
@@ -5078,7 +5113,7 @@ def test_insert_rows_w_repeated_fields(self):
self.assertEqual(len(errors), 0)
conn.api_request.assert_called_once_with(
- method="POST", path="/%s" % PATH, data=SENT, timeout=None,
+ method="POST", path="/%s" % PATH, data=SENT, timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_w_record_schema(self):
@@ -5144,7 +5179,7 @@ def test_insert_rows_w_record_schema(self):
self.assertEqual(len(errors), 0)
conn.api_request.assert_called_once_with(
- method="POST", path="/%s" % PATH, data=SENT, timeout=None
+ method="POST", path="/%s" % PATH, data=SENT, timeout=DEFAULT_TIMEOUT
)
def test_insert_rows_w_explicit_none_insert_ids(self):
@@ -5178,7 +5213,7 @@ def _row_data(row):
self.assertEqual(len(errors), 0)
conn.api_request.assert_called_once_with(
- method="POST", path="/{}".format(PATH), data=SENT, timeout=None,
+ method="POST", path="/{}".format(PATH), data=SENT, timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_errors(self):
@@ -5262,7 +5297,7 @@ def test_insert_rows_w_numeric(self):
project, ds_id, table_id
),
data=sent,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@unittest.skipIf(pandas is None, "Requires `pandas`")
@@ -5454,7 +5489,10 @@ def test_insert_rows_from_dataframe_many_columns(self):
]
}
expected_call = mock.call(
- method="POST", path=API_PATH, data=EXPECTED_SENT_DATA, timeout=None
+ method="POST",
+ path=API_PATH,
+ data=EXPECTED_SENT_DATA,
+ timeout=DEFAULT_TIMEOUT,
)
actual_calls = conn.api_request.call_args_list
@@ -5507,7 +5545,10 @@ def test_insert_rows_from_dataframe_w_explicit_none_insert_ids(self):
actual_calls = conn.api_request.call_args_list
assert len(actual_calls) == 1
assert actual_calls[0] == mock.call(
- method="POST", path=API_PATH, data=EXPECTED_SENT_DATA, timeout=None
+ method="POST",
+ path=API_PATH,
+ data=EXPECTED_SENT_DATA,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_json_default_behavior(self):
@@ -5587,7 +5628,7 @@ def test_insert_rows_json_w_explicitly_requested_autogenerated_insert_ids(self):
method="POST",
path="/projects/proj/datasets/dset/tables/tbl/insertAll",
data=expected_row_data,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_json_w_explicitly_disabled_insert_ids(self):
@@ -5617,7 +5658,7 @@ def test_insert_rows_json_w_explicitly_disabled_insert_ids(self):
method="POST",
path="/projects/proj/datasets/dset/tables/tbl/insertAll",
data=expected_row_data,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_json_with_iterator_row_ids(self):
@@ -5644,7 +5685,7 @@ def test_insert_rows_json_with_iterator_row_ids(self):
method="POST",
path="/projects/proj/datasets/dset/tables/tbl/insertAll",
data=expected_row_data,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_json_with_non_iterable_row_ids(self):
@@ -5697,7 +5738,7 @@ def test_insert_rows_json_w_explicit_none_insert_ids(self):
method="POST",
path="/projects/proj/datasets/dset/tables/tbl/insertAll",
data=expected,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_json_w_none_insert_ids_sequence(self):
@@ -5736,7 +5777,7 @@ def test_insert_rows_json_w_none_insert_ids_sequence(self):
method="POST",
path="/projects/proj/datasets/dset/tables/tbl/insertAll",
data=expected_row_data,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_insert_rows_w_wrong_arg(self):
@@ -5931,7 +5972,7 @@ def test_list_rows_w_start_index_w_page_size(self):
"maxResults": 2,
"formatOptions.useInt64Timestamp": True,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
mock.call(
method="GET",
@@ -5941,7 +5982,7 @@ def test_list_rows_w_start_index_w_page_size(self):
"maxResults": 2,
"formatOptions.useInt64Timestamp": True,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
]
)
@@ -6092,7 +6133,7 @@ def test_list_rows_repeated_fields(self):
"selectedFields": "color,struct",
"formatOptions.useInt64Timestamp": True,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_list_rows_w_record_schema(self):
@@ -6162,7 +6203,7 @@ def test_list_rows_w_record_schema(self):
method="GET",
path="/%s" % PATH,
query_params={"formatOptions.useInt64Timestamp": True},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
def test_list_rows_with_missing_schema(self):
@@ -6217,7 +6258,7 @@ def test_list_rows_with_missing_schema(self):
row_iter = client.list_rows(table)
conn.api_request.assert_called_once_with(
- method="GET", path=table_path, timeout=None
+ method="GET", path=table_path, timeout=DEFAULT_TIMEOUT
)
conn.api_request.reset_mock()
self.assertEqual(row_iter.total_rows, 2, msg=repr(table))
@@ -6227,7 +6268,7 @@ def test_list_rows_with_missing_schema(self):
method="GET",
path=tabledata_path,
query_params={"formatOptions.useInt64Timestamp": True},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
self.assertEqual(row_iter.total_rows, 3, msg=repr(table))
self.assertEqual(rows[0].name, "Phred Phlyntstone", msg=repr(table))
@@ -6400,7 +6441,7 @@ def test_load_table_from_file_resumable(self):
file_obj,
self.EXPECTED_CONFIGURATION,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project=self.EXPECTED_CONFIGURATION["jobReference"]["projectId"],
)
@@ -6433,7 +6474,7 @@ def test_load_table_from_file_w_explicit_project(self):
file_obj,
expected_resource,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project="other-project",
)
@@ -6467,7 +6508,7 @@ def test_load_table_from_file_w_client_location(self):
file_obj,
expected_resource,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project="other-project",
)
@@ -6529,7 +6570,7 @@ def test_load_table_from_file_resumable_metadata(self):
file_obj,
expected_config,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project=self.EXPECTED_CONFIGURATION["jobReference"]["projectId"],
)
@@ -6558,7 +6599,7 @@ def test_load_table_from_file_multipart(self):
self.EXPECTED_CONFIGURATION,
file_obj_size,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project=self.PROJECT,
)
@@ -6583,7 +6624,7 @@ def test_load_table_from_file_with_retries(self):
file_obj,
self.EXPECTED_CONFIGURATION,
num_retries,
- None,
+ DEFAULT_TIMEOUT,
project=self.EXPECTED_CONFIGURATION["jobReference"]["projectId"],
)
@@ -6620,7 +6661,7 @@ def test_load_table_from_file_with_readable_gzip(self):
gzip_file,
self.EXPECTED_CONFIGURATION,
_DEFAULT_NUM_RETRIES,
- None,
+ DEFAULT_TIMEOUT,
project=self.EXPECTED_CONFIGURATION["jobReference"]["projectId"],
)
@@ -6743,7 +6784,7 @@ def test_load_table_from_dataframe(self):
location=None,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_file = load_table_from_file.mock_calls[0][1][1]
@@ -6801,7 +6842,7 @@ def test_load_table_from_dataframe_w_client_location(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_file = load_table_from_file.mock_calls[0][1][1]
@@ -6855,7 +6896,7 @@ def test_load_table_from_dataframe_w_custom_job_config_wihtout_source_format(sel
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -6911,7 +6952,7 @@ def test_load_table_from_dataframe_w_custom_job_config_w_source_format(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -6974,7 +7015,7 @@ def test_load_table_from_dataframe_w_automatic_schema(self):
datetime.datetime(2012, 3, 14, 15, 16),
],
dtype="datetime64[ns]",
- ).dt.tz_localize(pytz.utc),
+ ).dt.tz_localize(datetime.timezone.utc),
),
]
)
@@ -7005,7 +7046,7 @@ def test_load_table_from_dataframe_w_automatic_schema(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7066,7 +7107,7 @@ def test_load_table_from_dataframe_w_index_and_auto_schema(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7113,7 +7154,7 @@ def test_load_table_from_dataframe_unknown_table(self):
location=None,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@unittest.skipIf(
@@ -7155,7 +7196,7 @@ def test_load_table_from_dataframe_w_nullable_int64_datatype(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7203,7 +7244,7 @@ def test_load_table_from_dataframe_w_nullable_int64_datatype_automatic_schema(se
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7265,7 +7306,7 @@ def test_load_table_from_dataframe_struct_fields(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7306,7 +7347,7 @@ def test_load_table_from_dataframe_w_partial_schema(self):
datetime.datetime(2012, 3, 14, 15, 16),
],
dtype="datetime64[ns]",
- ).dt.tz_localize(pytz.utc),
+ ).dt.tz_localize(datetime.timezone.utc),
),
("string_col", ["abc", None, "def"]),
("bytes_col", [b"abc", b"def", None]),
@@ -7340,7 +7381,7 @@ def test_load_table_from_dataframe_w_partial_schema(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7435,7 +7476,7 @@ def test_load_table_from_dataframe_w_partial_schema_missing_types(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
assert warned # there should be at least one warning
@@ -7585,7 +7626,7 @@ def test_load_table_from_dataframe_w_nulls(self):
location=self.LOCATION,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7651,7 +7692,7 @@ def test_load_table_from_dataframe_with_csv_source_format(self):
location=None,
project=None,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_file = load_table_from_file.mock_calls[0][1][1]
@@ -7689,7 +7730,7 @@ def test_load_table_from_json_basic_use(self):
location=client.location,
project=client.project,
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7742,7 +7783,7 @@ def test_load_table_from_json_non_default_args(self):
location="EU",
project="project-x",
job_config=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
@@ -7775,6 +7816,42 @@ def test_load_table_from_json_w_invalid_job_config(self):
err_msg = str(exc.value)
assert "Expected an instance of LoadJobConfig" in err_msg
+ def test_load_table_from_json_unicode_emoji_data_case(self):
+ from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
+
+ client = self._make_client()
+
+ emoji = "\U0001F3E6"
+ json_row = {"emoji": emoji}
+ json_rows = [json_row]
+
+ load_patch = mock.patch(
+ "google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
+ )
+
+ with load_patch as load_table_from_file:
+ client.load_table_from_json(json_rows, self.TABLE_REF)
+
+ load_table_from_file.assert_called_once_with(
+ client,
+ mock.ANY,
+ self.TABLE_REF,
+ size=mock.ANY,
+ num_retries=_DEFAULT_NUM_RETRIES,
+ job_id=mock.ANY,
+ job_id_prefix=None,
+ location=client.location,
+ project=client.project,
+ job_config=mock.ANY,
+ timeout=DEFAULT_TIMEOUT,
+ )
+
+ sent_data_file = load_table_from_file.mock_calls[0][1][1]
+
+ # make sure json_row's unicode characters are only encoded one time
+ expected_bytes = b'{"emoji": "' + emoji.encode("utf8") + b'"}'
+ assert sent_data_file.getvalue() == expected_bytes
+
# Low-level tests
@classmethod
diff --git a/tests/unit/test_create_dataset.py b/tests/unit/test_create_dataset.py
index d07aaed4f..67b21225d 100644
--- a/tests/unit/test_create_dataset.py
+++ b/tests/unit/test_create_dataset.py
@@ -15,6 +15,7 @@
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 mock
import pytest
@@ -111,7 +112,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
"access": [{"role": "OWNER", "userByEmail": USER_EMAIL}, {"view": VIEW}],
"labels": LABELS,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -143,7 +144,7 @@ def test_create_dataset_w_custom_property(client, PROJECT, DS_ID):
"newAlphaProperty": "unreleased property",
"labels": {},
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -176,7 +177,7 @@ def test_create_dataset_w_client_location_wo_dataset_location(PROJECT, DS_ID, LO
"labels": {},
"location": LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -211,7 +212,7 @@ def test_create_dataset_w_client_location_w_dataset_location(PROJECT, DS_ID, LOC
"labels": {},
"location": OTHER_LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -241,7 +242,7 @@ def test_create_dataset_w_reference(PROJECT, DS_ID, LOCATION):
"labels": {},
"location": LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -271,7 +272,7 @@ def test_create_dataset_w_fully_qualified_string(PROJECT, DS_ID, LOCATION):
"labels": {},
"location": LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -306,7 +307,7 @@ def test_create_dataset_w_string(PROJECT, DS_ID, LOCATION):
"labels": {},
"location": LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -356,8 +357,8 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(PROJECT, DS_ID, LOCATION)
"labels": {},
"location": LOCATION,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
),
- mock.call(method="GET", path=get_path, timeout=None),
+ mock.call(method="GET", path=get_path, timeout=DEFAULT_TIMEOUT),
]
)
diff --git a/tests/unit/test_delete_dataset.py b/tests/unit/test_delete_dataset.py
index 3a65e031c..b48beb147 100644
--- a/tests/unit/test_delete_dataset.py
+++ b/tests/unit/test_delete_dataset.py
@@ -14,6 +14,7 @@
from .helpers import make_connection, make_client, dataset_polymorphic
import google.api_core.exceptions
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
import pytest
@@ -40,7 +41,7 @@ def test_delete_dataset_delete_contents(
method="DELETE",
path="/%s" % PATH,
query_params={"deleteContents": "true"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -61,7 +62,7 @@ def test_delete_dataset_w_not_found_ok_false(PROJECT, DS_ID):
client.delete_dataset(DS_ID)
conn.api_request.assert_called_with(
- method="DELETE", path=path, query_params={}, timeout=None
+ method="DELETE", path=path, query_params={}, timeout=DEFAULT_TIMEOUT
)
@@ -74,5 +75,5 @@ def test_delete_dataset_w_not_found_ok_true(PROJECT, DS_ID):
)
client.delete_dataset(DS_ID, not_found_ok=True)
conn.api_request.assert_called_with(
- method="DELETE", path=path, query_params={}, timeout=None
+ method="DELETE", path=path, query_params={}, timeout=DEFAULT_TIMEOUT
)
diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py
new file mode 100644
index 000000000..b2095d2f2
--- /dev/null
+++ b/tests/unit/test_job_retry.py
@@ -0,0 +1,247 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 datetime
+import re
+
+import mock
+import pytest
+
+import google.api_core.exceptions
+import google.api_core.retry
+
+from .helpers import make_connection
+
+
+# With job_retry_on_query, we're testing 4 scenarios:
+# - No `job_retry` passed, retry on default rateLimitExceeded.
+# - Pass NotFound retry to `query`.
+# - Pass NotFound retry to `result`.
+# - Pass BadRequest retry to query, with the value passed to `result` overriding.
+@pytest.mark.parametrize("job_retry_on_query", [None, "Query", "Result", "Both"])
+@mock.patch("time.sleep")
+def test_retry_failed_jobs(sleep, client, job_retry_on_query):
+ """
+ Test retry of job failures, as opposed to API-invocation failures.
+ """
+
+ retry_notfound = google.api_core.retry.Retry(
+ predicate=google.api_core.retry.if_exception_type(
+ google.api_core.exceptions.NotFound
+ )
+ )
+ retry_badrequest = google.api_core.retry.Retry(
+ predicate=google.api_core.retry.if_exception_type(
+ google.api_core.exceptions.BadRequest
+ )
+ )
+
+ if job_retry_on_query is None:
+ reason = "rateLimitExceeded"
+ else:
+ reason = "notFound"
+
+ err = dict(reason=reason)
+ responses = [
+ dict(status=dict(state="DONE", errors=[err], errorResult=err)),
+ dict(status=dict(state="DONE", errors=[err], errorResult=err)),
+ dict(status=dict(state="DONE", errors=[err], errorResult=err)),
+ dict(status=dict(state="DONE")),
+ dict(rows=[{"f": [{"v": "1"}]}], totalRows="1"),
+ ]
+
+ def api_request(method, path, query_params=None, data=None, **kw):
+ response = responses.pop(0)
+ if data:
+ response["jobReference"] = data["jobReference"]
+ else:
+ response["jobReference"] = dict(
+ jobId=path.split("/")[-1], projectId="PROJECT"
+ )
+ return response
+
+ conn = client._connection = make_connection()
+ conn.api_request.side_effect = api_request
+
+ if job_retry_on_query == "Query":
+ job_retry = dict(job_retry=retry_notfound)
+ elif job_retry_on_query == "Both":
+ # This will be overridden in `result`
+ job_retry = dict(job_retry=retry_badrequest)
+ else:
+ job_retry = {}
+ job = client.query("select 1", **job_retry)
+
+ orig_job_id = job.job_id
+ job_retry = (
+ dict(job_retry=retry_notfound)
+ if job_retry_on_query in ("Result", "Both")
+ else {}
+ )
+ result = job.result(**job_retry)
+ assert result.total_rows == 1
+ assert not responses # We made all the calls we expected to.
+
+ # The job adjusts it's job id based on the id of the last attempt.
+ assert job.job_id != orig_job_id
+ assert job.job_id == conn.mock_calls[3][2]["data"]["jobReference"]["jobId"]
+
+ # We had to sleep three times
+ assert len(sleep.mock_calls) == 3
+
+ # Sleeps are random, however they're more than 0
+ assert min(c[1][0] for c in sleep.mock_calls) > 0
+
+ # They're at most 2 * (multiplier**(number of sleeps - 1)) * initial
+ # The default multiplier is 2
+ assert max(c[1][0] for c in sleep.mock_calls) <= 8
+
+ # We can ask for the result again:
+ responses = [
+ dict(rows=[{"f": [{"v": "1"}]}], totalRows="1"),
+ ]
+ orig_job_id = job.job_id
+ result = job.result()
+ assert result.total_rows == 1
+ assert not responses # We made all the calls we expected to.
+
+ # We wouldn't (and didn't) fail, because we're dealing with a successful job.
+ # So the job id hasn't changed.
+ assert job.job_id == orig_job_id
+
+
+# With job_retry_on_query, we're testing 4 scenarios:
+# - Pass None retry to `query`.
+# - Pass None retry to `result`.
+@pytest.mark.parametrize("job_retry_on_query", ["Query", "Result"])
+@mock.patch("time.sleep")
+def test_disable_retry_failed_jobs(sleep, client, job_retry_on_query):
+ """
+ Test retry of job failures, as opposed to API-invocation failures.
+ """
+ err = dict(reason="rateLimitExceeded")
+ responses = [dict(status=dict(state="DONE", errors=[err], errorResult=err))] * 3
+
+ def api_request(method, path, query_params=None, data=None, **kw):
+ response = responses.pop(0)
+ response["jobReference"] = data["jobReference"]
+ return response
+
+ conn = client._connection = make_connection()
+ conn.api_request.side_effect = api_request
+
+ if job_retry_on_query == "Query":
+ job_retry = dict(job_retry=None)
+ else:
+ job_retry = {}
+ job = client.query("select 1", **job_retry)
+
+ orig_job_id = job.job_id
+ job_retry = dict(job_retry=None) if job_retry_on_query == "Result" else {}
+ with pytest.raises(google.api_core.exceptions.Forbidden):
+ job.result(**job_retry)
+
+ assert job.job_id == orig_job_id
+ assert len(sleep.mock_calls) == 0
+
+
+@mock.patch("google.api_core.retry.datetime_helpers")
+@mock.patch("time.sleep")
+def test_retry_failed_jobs_after_retry_failed(sleep, datetime_helpers, client):
+ """
+ If at first you don't succeed, maybe you will later. :)
+ """
+ conn = client._connection = make_connection()
+
+ datetime_helpers.utcnow.return_value = datetime.datetime(2021, 7, 29, 10, 43, 2)
+
+ err = dict(reason="rateLimitExceeded")
+
+ def api_request(method, path, query_params=None, data=None, **kw):
+ calls = sleep.mock_calls
+ if calls:
+ datetime_helpers.utcnow.return_value += datetime.timedelta(
+ seconds=calls[-1][1][0]
+ )
+ response = dict(status=dict(state="DONE", errors=[err], errorResult=err))
+ response["jobReference"] = data["jobReference"]
+ return response
+
+ conn.api_request.side_effect = api_request
+
+ job = client.query("select 1")
+ orig_job_id = job.job_id
+
+ with pytest.raises(google.api_core.exceptions.RetryError):
+ job.result()
+
+ # We never got a successful job, so the job id never changed:
+ assert job.job_id == orig_job_id
+
+ # We failed because we couldn't succeed after 120 seconds.
+ # But we can try again:
+ err2 = dict(reason="backendError") # We also retry on this
+ responses = [
+ dict(status=dict(state="DONE", errors=[err2], errorResult=err2)),
+ dict(status=dict(state="DONE", errors=[err], errorResult=err)),
+ dict(status=dict(state="DONE", errors=[err2], errorResult=err2)),
+ dict(status=dict(state="DONE")),
+ dict(rows=[{"f": [{"v": "1"}]}], totalRows="1"),
+ ]
+
+ def api_request(method, path, query_params=None, data=None, **kw):
+ calls = sleep.mock_calls
+ datetime_helpers.utcnow.return_value += datetime.timedelta(
+ seconds=calls[-1][1][0]
+ )
+ response = responses.pop(0)
+ if data:
+ response["jobReference"] = data["jobReference"]
+ else:
+ response["jobReference"] = dict(
+ jobId=path.split("/")[-1], projectId="PROJECT"
+ )
+ return response
+
+ conn.api_request.side_effect = api_request
+ result = job.result()
+ assert result.total_rows == 1
+ assert not responses # We made all the calls we expected to.
+ assert job.job_id != orig_job_id
+
+
+def test_raises_on_job_retry_on_query_with_non_retryable_jobs(client):
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "`job_retry` was provided, but the returned job is"
+ " not retryable, because a custom `job_id` was"
+ " provided."
+ ),
+ ):
+ client.query("select 42", job_id=42, job_retry=google.api_core.retry.Retry())
+
+
+def test_raises_on_job_retry_on_result_with_non_retryable_jobs(client):
+ client._connection = make_connection({})
+ job = client.query("select 42", job_id=42)
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "`job_retry` was provided, but this job is"
+ " not retryable, because a custom `job_id` was"
+ " provided to the query that created this job."
+ ),
+ ):
+ job.result(job_retry=google.api_core.retry.Retry())
diff --git a/tests/unit/test_list_datasets.py b/tests/unit/test_list_datasets.py
index 7793a7ba6..6f0b55c5e 100644
--- a/tests/unit/test_list_datasets.py
+++ b/tests/unit/test_list_datasets.py
@@ -1,11 +1,11 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -15,6 +15,7 @@
import mock
import pytest
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
from .helpers import make_connection
@@ -65,7 +66,7 @@ def test_list_datasets_defaults(client, PROJECT, extra, query):
assert token == TOKEN
conn.api_request.assert_called_once_with(
- method="GET", path="/%s" % PATH, query_params=query, timeout=None
+ method="GET", path="/%s" % PATH, query_params=query, timeout=DEFAULT_TIMEOUT
)
@@ -120,5 +121,5 @@ def test_list_datasets_explicit_response_missing_datasets_key(client, PROJECT):
"maxResults": 3,
"pageToken": TOKEN,
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
diff --git a/tests/unit/test_list_jobs.py b/tests/unit/test_list_jobs.py
index f348be724..1fb40d446 100644
--- a/tests/unit/test_list_jobs.py
+++ b/tests/unit/test_list_jobs.py
@@ -1,11 +1,11 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -17,6 +17,7 @@
import mock
import pytest
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
from .helpers import make_connection
@@ -136,7 +137,7 @@ def test_list_jobs_defaults(client, PROJECT, DS_ID, extra, query):
method="GET",
path="/%s" % PATH,
query_params=dict({"projection": "full"}, **query),
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -185,7 +186,7 @@ def test_list_jobs_load_job_wo_sourceUris(client, PROJECT, DS_ID):
method="GET",
path="/%s" % PATH,
query_params={"projection": "full"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -220,7 +221,7 @@ def test_list_jobs_explicit_missing(client, PROJECT):
"allUsers": True,
"stateFilter": "done",
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -233,7 +234,7 @@ def test_list_jobs_w_project(client, PROJECT):
method="GET",
path="/projects/other-project/jobs",
query_params={"projection": "full"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -269,7 +270,7 @@ def test_list_jobs_w_time_filter(client, PROJECT):
"minCreationTime": "1",
"maxCreationTime": str(end_time_millis),
},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
@@ -286,6 +287,6 @@ def test_list_jobs_w_parent_job_filter(client, PROJECT):
method="GET",
path="/projects/%s/jobs" % PROJECT,
query_params={"projection": "full", "parentJobId": "parent-job-123"},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
conn.api_request.reset_mock()
diff --git a/tests/unit/test_list_models.py b/tests/unit/test_list_models.py
index 4ede9a7dd..b14852338 100644
--- a/tests/unit/test_list_models.py
+++ b/tests/unit/test_list_models.py
@@ -1,20 +1,22 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
-from .helpers import make_connection, dataset_polymorphic
import pytest
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
+from .helpers import make_connection, dataset_polymorphic
+
def test_list_models_empty_w_timeout(client, PROJECT, DS_ID):
path = "/projects/{}/datasets/{}/models".format(PROJECT, DS_ID)
@@ -82,7 +84,7 @@ def test_list_models_defaults(
assert token == TOKEN
conn.api_request.assert_called_once_with(
- method="GET", path="/%s" % PATH, query_params=query, timeout=None
+ method="GET", path="/%s" % PATH, query_params=query, timeout=DEFAULT_TIMEOUT
)
diff --git a/tests/unit/test_list_projects.py b/tests/unit/test_list_projects.py
index a88540dd5..190612b44 100644
--- a/tests/unit/test_list_projects.py
+++ b/tests/unit/test_list_projects.py
@@ -1,11 +1,11 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -15,6 +15,7 @@
import mock
import pytest
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
from .helpers import make_connection
@@ -66,7 +67,7 @@ def test_list_projects_defaults(client, PROJECT, extra, query):
assert token == TOKEN
conn.api_request.assert_called_once_with(
- method="GET", path="/projects", query_params=query, timeout=None
+ method="GET", path="/projects", query_params=query, timeout=DEFAULT_TIMEOUT
)
@@ -115,5 +116,5 @@ def test_list_projects_explicit_response_missing_projects_key(client):
method="GET",
path="/projects",
query_params={"maxResults": 3, "pageToken": TOKEN},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
diff --git a/tests/unit/test_list_routines.py b/tests/unit/test_list_routines.py
index 069966542..80e62d6bd 100644
--- a/tests/unit/test_list_routines.py
+++ b/tests/unit/test_list_routines.py
@@ -1,20 +1,22 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
-from .helpers import make_connection, dataset_polymorphic
import pytest
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
+from .helpers import make_connection, dataset_polymorphic
+
def test_list_routines_empty_w_timeout(client):
conn = client._connection = make_connection({})
@@ -85,7 +87,7 @@ def test_list_routines_defaults(
assert actual_token == token
conn.api_request.assert_called_once_with(
- method="GET", path=path, query_params=query, timeout=None
+ method="GET", path=path, query_params=query, timeout=DEFAULT_TIMEOUT
)
diff --git a/tests/unit/test_list_tables.py b/tests/unit/test_list_tables.py
index 45d15bed3..8360f6605 100644
--- a/tests/unit/test_list_tables.py
+++ b/tests/unit/test_list_tables.py
@@ -1,21 +1,23 @@
# Copyright 2021 Google LLC
-
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# https://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
-from .helpers import make_connection, dataset_polymorphic
-import google.cloud.bigquery.dataset
import pytest
+import google.cloud.bigquery.dataset
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
+from .helpers import make_connection, dataset_polymorphic
+
@dataset_polymorphic
def test_list_tables_empty_w_timeout(
@@ -89,7 +91,7 @@ def test_list_tables_defaults(make_dataset, get_reference, client, PROJECT, DS_I
assert token == TOKEN
conn.api_request.assert_called_once_with(
- method="GET", path="/%s" % PATH, query_params={}, timeout=None
+ method="GET", path="/%s" % PATH, query_params={}, timeout=DEFAULT_TIMEOUT
)
@@ -150,7 +152,7 @@ def test_list_tables_explicit(client, PROJECT, DS_ID):
method="GET",
path="/%s" % PATH,
query_params={"maxResults": 3, "pageToken": TOKEN},
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
index d030482cc..36cbf4993 100644
--- a/tests/unit/test_magics.py
+++ b/tests/unit/test_magics.py
@@ -32,6 +32,7 @@
from google.cloud.bigquery import job
from google.cloud.bigquery import table
from google.cloud.bigquery.magics import magics
+from google.cloud.bigquery.retry import DEFAULT_TIMEOUT
from tests.unit.helpers import make_connection
from test_utils.imports import maybe_fail_import
@@ -185,7 +186,7 @@ def test_context_with_default_connection():
method="POST",
path="/projects/project-from-env/jobs",
data=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
query_results_call = mock.call(
method="GET",
@@ -249,7 +250,7 @@ def test_context_with_custom_connection():
method="POST",
path="/projects/project-from-env/jobs",
data=mock.ANY,
- timeout=None,
+ timeout=DEFAULT_TIMEOUT,
)
query_results_call = mock.call(
method="GET",
@@ -660,7 +661,9 @@ def warning_match(warning):
assert client_info.user_agent == "ipython-" + IPython.__version__
query_job_mock.to_dataframe.assert_called_once_with(
- bqstorage_client=bqstorage_instance_mock, progress_bar_type="tqdm"
+ bqstorage_client=bqstorage_instance_mock,
+ create_bqstorage_client=mock.ANY,
+ progress_bar_type="tqdm",
)
assert isinstance(return_value, pandas.DataFrame)
@@ -703,7 +706,9 @@ def test_bigquery_magic_with_rest_client_requested(monkeypatch):
bqstorage_mock.assert_not_called()
query_job_mock.to_dataframe.assert_called_once_with(
- bqstorage_client=None, progress_bar_type="tqdm"
+ bqstorage_client=None,
+ create_bqstorage_client=False,
+ progress_bar_type="tqdm",
)
assert isinstance(return_value, pandas.DataFrame)
@@ -757,7 +762,12 @@ def test_bigquery_magic_w_max_results_valid_calls_queryjob_result():
client_query_mock.return_value = query_job_mock
ip.run_cell_magic("bigquery", "--max_results=5", sql)
- query_job_mock.result.assert_called_with(max_results=5)
+ query_job_mock.result.assert_called_with(max_results=5)
+ query_job_mock.result.return_value.to_dataframe.assert_called_once_with(
+ bqstorage_client=None,
+ create_bqstorage_client=False,
+ progress_bar_type=mock.ANY,
+ )
@pytest.mark.usefixtures("ipython_interactive")
@@ -929,7 +939,7 @@ def test_bigquery_magic_w_table_id_and_bqstorage_client():
ip.run_cell_magic("bigquery", "--max_results=5", table_id)
row_iterator_mock.to_dataframe.assert_called_once_with(
- bqstorage_client=bqstorage_instance_mock
+ bqstorage_client=bqstorage_instance_mock, create_bqstorage_client=mock.ANY,
)
@@ -1246,7 +1256,9 @@ def test_bigquery_magic_w_progress_bar_type_w_context_setter(monkeypatch):
bqstorage_mock.assert_not_called()
query_job_mock.to_dataframe.assert_called_once_with(
- bqstorage_client=None, progress_bar_type=magics.context.progress_bar_type
+ bqstorage_client=None,
+ create_bqstorage_client=False,
+ progress_bar_type=magics.context.progress_bar_type,
)
assert isinstance(return_value, pandas.DataFrame)
diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py
index 9483fe8dd..69a6772e5 100644
--- a/tests/unit/test_query.py
+++ b/tests/unit/test_query.py
@@ -13,6 +13,7 @@
# limitations under the License.
import datetime
+import decimal
import unittest
import mock
@@ -430,6 +431,18 @@ def test_positional(self):
self.assertEqual(param.type_, "INT64")
self.assertEqual(param.value, 123)
+ def test_ctor_w_scalar_query_parameter_type(self):
+ from google.cloud.bigquery import enums
+
+ param = self._make_one(
+ name="foo",
+ type_=enums.SqlParameterScalarTypes.BIGNUMERIC,
+ value=decimal.Decimal("123.456"),
+ )
+ self.assertEqual(param.name, "foo")
+ self.assertEqual(param.type_, "BIGNUMERIC")
+ self.assertEqual(param.value, decimal.Decimal("123.456"))
+
def test_from_api_repr_w_name(self):
RESOURCE = {
"name": "foo",
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
index 6fb7f93fd..e0a992f78 100644
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -55,6 +55,18 @@ def test_w_unstructured_requests_chunked_encoding_error(self):
exc = requests.exceptions.ChunkedEncodingError()
self.assertTrue(self._call_fut(exc))
+ def test_w_unstructured_requests_connecttimeout(self):
+ exc = requests.exceptions.ConnectTimeout()
+ self.assertTrue(self._call_fut(exc))
+
+ def test_w_unstructured_requests_readtimeout(self):
+ exc = requests.exceptions.ReadTimeout()
+ self.assertTrue(self._call_fut(exc))
+
+ def test_w_unstructured_requests_timeout(self):
+ exc = requests.exceptions.Timeout()
+ self.assertTrue(self._call_fut(exc))
+
def test_w_auth_transporterror(self):
from google.auth.exceptions import TransportError
@@ -86,3 +98,27 @@ def test_w_unstructured_bad_gateway(self):
exc = BadGateway("testing")
self.assertTrue(self._call_fut(exc))
+
+
+def test_DEFAULT_JOB_RETRY_predicate():
+ from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY
+ from google.api_core.exceptions import ClientError
+
+ assert not DEFAULT_JOB_RETRY._predicate(TypeError())
+ assert not DEFAULT_JOB_RETRY._predicate(ClientError("fail"))
+ assert not DEFAULT_JOB_RETRY._predicate(
+ ClientError("fail", errors=[dict(reason="idk")])
+ )
+
+ assert DEFAULT_JOB_RETRY._predicate(
+ ClientError("fail", errors=[dict(reason="rateLimitExceeded")])
+ )
+ assert DEFAULT_JOB_RETRY._predicate(
+ ClientError("fail", errors=[dict(reason="backendError")])
+ )
+
+
+def test_DEFAULT_JOB_RETRY_deadline():
+ from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY
+
+ assert DEFAULT_JOB_RETRY._deadline == 600
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
index 4b1fd833b..1ce930ee4 100644
--- a/tests/unit/test_table.py
+++ b/tests/unit/test_table.py
@@ -14,15 +14,14 @@
import datetime
import logging
+import re
import time
import types
import unittest
import warnings
import mock
-import pkg_resources
import pytest
-import pytz
import google.api_core.exceptions
from test_utils.imports import maybe_fail_import
@@ -41,14 +40,16 @@
except (ImportError, AttributeError): # pragma: NO COVER
pandas = None
+try:
+ import geopandas
+except (ImportError, AttributeError): # pragma: NO COVER
+ geopandas = None
+
try:
import pyarrow
import pyarrow.types
-
- PYARROW_VERSION = pkg_resources.parse_version(pyarrow.__version__)
except ImportError: # pragma: NO COVER
pyarrow = None
- PYARROW_VERSION = pkg_resources.parse_version("0.0.1")
try:
from tqdm import tqdm
@@ -58,9 +59,6 @@
from google.cloud.bigquery.dataset import DatasetReference
-PYARROW_TIMESTAMP_VERSION = pkg_resources.parse_version("2.0.0")
-
-
def _mock_client():
from google.cloud.bigquery import client
@@ -115,8 +113,6 @@ def _make_one(self, *args, **kw):
return self._get_target_class()(*args, **kw)
def test_ctor_defaults(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset_ref = DatasetReference("project_1", "dataset_1")
table_ref = self._make_one(dataset_ref, "table_1")
@@ -124,8 +120,6 @@ def test_ctor_defaults(self):
self.assertEqual(table_ref.table_id, "table_1")
def test_to_api_repr(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset_ref = DatasetReference("project_1", "dataset_1")
table_ref = self._make_one(dataset_ref, "table_1")
@@ -137,7 +131,6 @@ def test_to_api_repr(self):
)
def test_from_api_repr(self):
- from google.cloud.bigquery.dataset import DatasetReference
from google.cloud.bigquery.table import TableReference
dataset_ref = DatasetReference("project_1", "dataset_1")
@@ -204,8 +197,6 @@ def test_from_string_ignores_default_project(self):
self.assertEqual(got.table_id, "string_table")
def test___eq___wrong_type(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset_ref = DatasetReference("project_1", "dataset_1")
table = self._make_one(dataset_ref, "table_1")
other = object()
@@ -213,8 +204,6 @@ def test___eq___wrong_type(self):
self.assertEqual(table, mock.ANY)
def test___eq___project_mismatch(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
other_dataset = DatasetReference("project_2", "dataset_1")
table = self._make_one(dataset, "table_1")
@@ -222,8 +211,6 @@ def test___eq___project_mismatch(self):
self.assertNotEqual(table, other)
def test___eq___dataset_mismatch(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
other_dataset = DatasetReference("project_1", "dataset_2")
table = self._make_one(dataset, "table_1")
@@ -231,24 +218,18 @@ def test___eq___dataset_mismatch(self):
self.assertNotEqual(table, other)
def test___eq___table_mismatch(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
table = self._make_one(dataset, "table_1")
other = self._make_one(dataset, "table_2")
self.assertNotEqual(table, other)
def test___eq___equality(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
table = self._make_one(dataset, "table_1")
other = self._make_one(dataset, "table_1")
self.assertEqual(table, other)
def test___hash__set_equality(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
table1 = self._make_one(dataset, "table1")
table2 = self._make_one(dataset, "table2")
@@ -257,8 +238,6 @@ def test___hash__set_equality(self):
self.assertEqual(set_one, set_two)
def test___hash__not_equals(self):
- from google.cloud.bigquery.dataset import DatasetReference
-
dataset = DatasetReference("project_1", "dataset_1")
table1 = self._make_one(dataset, "table1")
table2 = self._make_one(dataset, "table2")
@@ -294,8 +273,6 @@ def _get_target_class():
return Table
def _make_one(self, *args, **kw):
- from google.cloud.bigquery.dataset import DatasetReference
-
if len(args) == 0:
dataset = DatasetReference(self.PROJECT, self.DS_ID)
table_ref = dataset.table(self.TABLE_NAME)
@@ -581,6 +558,68 @@ def test_num_rows_getter(self):
with self.assertRaises(ValueError):
getattr(table, "num_rows")
+ def test__eq__wrong_type(self):
+ table = self._make_one("project_foo.dataset_bar.table_baz")
+
+ class TableWannabe:
+ pass
+
+ not_a_table = TableWannabe()
+ not_a_table._properties = table._properties
+
+ assert table != not_a_table # Can't fake it.
+
+ def test__eq__same_table_basic(self):
+ table_1 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_2 = self._make_one("project_foo.dataset_bar.table_baz")
+ assert table_1 == table_2
+
+ def test__eq__same_table_multiple_properties(self):
+ from google.cloud.bigquery import SchemaField
+
+ table_1 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_1.require_partition_filter = True
+ table_1.labels = {"first": "one", "second": "two"}
+
+ table_1.schema = [
+ SchemaField("name", "STRING", "REQUIRED"),
+ SchemaField("age", "INTEGER", "NULLABLE"),
+ ]
+
+ table_2 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_2.require_partition_filter = True
+ table_2.labels = {"first": "one", "second": "two"}
+ table_2.schema = [
+ SchemaField("name", "STRING", "REQUIRED"),
+ SchemaField("age", "INTEGER", "NULLABLE"),
+ ]
+
+ assert table_1 == table_2
+
+ def test__eq__same_table_property_different(self):
+ table_1 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_1.description = "This is table baz"
+
+ table_2 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_2.description = "This is also table baz"
+
+ assert table_1 == table_2 # Still equal, only table reference is important.
+
+ def test__eq__different_table(self):
+ table_1 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_2 = self._make_one("project_foo.dataset_bar.table_baz_2")
+
+ assert table_1 != table_2
+
+ def test_hashable(self):
+ table_1 = self._make_one("project_foo.dataset_bar.table_baz")
+ table_1.description = "This is a table"
+
+ table_1b = self._make_one("project_foo.dataset_bar.table_baz")
+ table_1b.description = "Metadata is irrelevant for hashes"
+
+ assert hash(table_1) == hash(table_1b)
+
def test_schema_setter_non_sequence(self):
dataset = DatasetReference(self.PROJECT, self.DS_ID)
table_ref = dataset.table(self.TABLE_NAME)
@@ -873,7 +912,9 @@ def test_mview_last_refresh_time(self):
}
self.assertEqual(
table.mview_last_refresh_time,
- datetime.datetime(2020, 11, 30, 15, 57, 22, 496000, tzinfo=pytz.utc),
+ datetime.datetime(
+ 2020, 11, 30, 15, 57, 22, 496000, tzinfo=datetime.timezone.utc
+ ),
)
def test_mview_enable_refresh(self):
@@ -1543,6 +1584,148 @@ def test_to_api_repr(self):
table = self._make_one(resource)
self.assertEqual(table.to_api_repr(), resource)
+ def test__eq__wrong_type(self):
+ resource = {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ table = self._make_one(resource)
+
+ class FakeTableListItem:
+ project = "project_foo"
+ dataset_id = "dataset_bar"
+ table_id = "table_baz"
+
+ not_a_table = FakeTableListItem()
+
+ assert table != not_a_table # Can't fake it.
+
+ def test__eq__same_table(self):
+ resource = {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ table_1 = self._make_one(resource)
+ table_2 = self._make_one(resource)
+
+ assert table_1 == table_2
+
+ def test__eq__same_table_property_different(self):
+ table_ref_resource = {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+
+ resource_1 = {"tableReference": table_ref_resource, "friendlyName": "Table One"}
+ table_1 = self._make_one(resource_1)
+
+ resource_2 = {"tableReference": table_ref_resource, "friendlyName": "Table Two"}
+ table_2 = self._make_one(resource_2)
+
+ assert table_1 == table_2 # Still equal, only table reference is important.
+
+ def test__eq__different_table(self):
+ resource_1 = {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ table_1 = self._make_one(resource_1)
+
+ resource_2 = {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_quux",
+ }
+ }
+ table_2 = self._make_one(resource_2)
+
+ assert table_1 != table_2
+
+ def test_hashable(self):
+ resource = {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ table_item = self._make_one(resource)
+ table_item_2 = self._make_one(resource)
+
+ assert hash(table_item) == hash(table_item_2)
+
+
+class TestTableClassesInterchangeability:
+ @staticmethod
+ def _make_table(*args, **kwargs):
+ from google.cloud.bigquery.table import Table
+
+ return Table(*args, **kwargs)
+
+ @staticmethod
+ def _make_table_ref(*args, **kwargs):
+ from google.cloud.bigquery.table import TableReference
+
+ return TableReference(*args, **kwargs)
+
+ @staticmethod
+ def _make_table_list_item(*args, **kwargs):
+ from google.cloud.bigquery.table import TableListItem
+
+ return TableListItem(*args, **kwargs)
+
+ def test_table_eq_table_ref(self):
+
+ table = self._make_table("project_foo.dataset_bar.table_baz")
+ dataset_ref = DatasetReference("project_foo", "dataset_bar")
+ table_ref = self._make_table_ref(dataset_ref, "table_baz")
+
+ assert table == table_ref
+ assert table_ref == table
+
+ def test_table_eq_table_list_item(self):
+ table = self._make_table("project_foo.dataset_bar.table_baz")
+ table_list_item = self._make_table_list_item(
+ {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ )
+
+ assert table == table_list_item
+ assert table_list_item == table
+
+ def test_table_ref_eq_table_list_item(self):
+
+ dataset_ref = DatasetReference("project_foo", "dataset_bar")
+ table_ref = self._make_table_ref(dataset_ref, "table_baz")
+ table_list_item = self._make_table_list_item(
+ {
+ "tableReference": {
+ "projectId": "project_foo",
+ "datasetId": "dataset_bar",
+ "tableId": "table_baz",
+ }
+ }
+ )
+
+ assert table_ref == table_list_item
+ assert table_list_item == table_ref
+
class TestSnapshotDefinition:
@staticmethod
@@ -1665,6 +1848,27 @@ def test_to_dataframe_iterable(self):
self.assertEqual(len(df), 0) # Verify the number of rows.
self.assertEqual(len(df.columns), 0)
+ @mock.patch("google.cloud.bigquery.table.geopandas", new=None)
+ def test_to_geodataframe_if_geopandas_is_none(self):
+ row_iterator = self._make_one()
+ with self.assertRaisesRegex(
+ ValueError,
+ re.escape(
+ "The geopandas library is not installed, please install "
+ "geopandas to use the to_geodataframe() function."
+ ),
+ ):
+ row_iterator.to_geodataframe(create_bqstorage_client=False)
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe(self):
+ row_iterator = self._make_one()
+ df = row_iterator.to_geodataframe(create_bqstorage_client=False)
+ self.assertIsInstance(df, geopandas.GeoDataFrame)
+ self.assertEqual(len(df), 0) # verify the number of rows
+ self.assertEqual(df.crs.srs, "EPSG:4326")
+ self.assertEqual(df.crs.name, "WGS 84")
+
class TestRowIterator(unittest.TestCase):
def _class_under_test(self):
@@ -1702,6 +1906,16 @@ def _make_one(
client, api_request, path, schema, table=table, **kwargs
)
+ def _make_one_from_data(self, schema=(), rows=()):
+ from google.cloud.bigquery.schema import SchemaField
+
+ schema = [SchemaField(*a) for a in schema]
+ rows = [{"f": [{"v": v} for v in row]} for row in rows]
+
+ path = "/foo"
+ api_request = mock.Mock(return_value={"rows": 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
@@ -2695,10 +2909,7 @@ def test_to_dataframe_timestamp_out_of_pyarrow_bounds(self):
df = row_iterator.to_dataframe(create_bqstorage_client=False)
- tzinfo = None
- if PYARROW_VERSION >= PYARROW_TIMESTAMP_VERSION:
- tzinfo = datetime.timezone.utc
-
+ tzinfo = datetime.timezone.utc
self.assertIsInstance(df, pandas.DataFrame)
self.assertEqual(len(df), 2) # verify the number of rows
self.assertEqual(list(df.columns), ["some_timestamp"])
@@ -2996,6 +3207,18 @@ def test_to_dataframe_error_if_pandas_is_none(self):
with self.assertRaises(ValueError):
row_iterator.to_dataframe()
+ @unittest.skipIf(pandas is None, "Requires `pandas`")
+ @mock.patch("google.cloud.bigquery.table.shapely", new=None)
+ def test_to_dataframe_error_if_shapely_is_none(self):
+ with self.assertRaisesRegex(
+ ValueError,
+ re.escape(
+ "The shapely library is not installed, please install "
+ "shapely to use the geography_as_object option."
+ ),
+ ):
+ self._make_one_from_data().to_dataframe(geography_as_object=True)
+
@unittest.skipIf(pandas is None, "Requires `pandas`")
def test_to_dataframe_max_results_w_bqstorage_warning(self):
from google.cloud.bigquery.schema import SchemaField
@@ -3753,6 +3976,199 @@ def test_to_dataframe_concat_categorical_dtype_w_pyarrow(self):
# Don't close the client if it was passed in.
bqstorage_client._transport.grpc_channel.close.assert_not_called()
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_dataframe_geography_as_object(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "GEOGRAPHY")),
+ (
+ ("foo", "Point(0 0)"),
+ ("bar", None),
+ ("baz", "Polygon((0 0, 0 1, 1 0, 0 0))"),
+ ),
+ )
+ df = row_iterator.to_dataframe(
+ create_bqstorage_client=False, geography_as_object=True,
+ )
+ self.assertIsInstance(df, pandas.DataFrame)
+ self.assertEqual(len(df), 3) # verify the number of rows
+ self.assertEqual(list(df), ["name", "geog"]) # verify the column names
+ self.assertEqual(df.name.dtype.name, "object")
+ self.assertEqual(df.geog.dtype.name, "object")
+ self.assertIsInstance(df.geog, pandas.Series)
+ self.assertEqual(
+ [v.__class__.__name__ for v in df.geog], ["Point", "float", "Polygon"]
+ )
+
+ @mock.patch("google.cloud.bigquery.table.geopandas", new=None)
+ def test_to_geodataframe_error_if_geopandas_is_none(self):
+ with self.assertRaisesRegex(
+ ValueError,
+ re.escape(
+ "The geopandas library is not installed, please install "
+ "geopandas to use the to_geodataframe() function."
+ ),
+ ):
+ self._make_one_from_data().to_geodataframe()
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "GEOGRAPHY")),
+ (
+ ("foo", "Point(0 0)"),
+ ("bar", None),
+ ("baz", "Polygon((0 0, 0 1, 1 0, 0 0))"),
+ ),
+ )
+ df = row_iterator.to_geodataframe(create_bqstorage_client=False)
+ self.assertIsInstance(df, geopandas.GeoDataFrame)
+ self.assertEqual(len(df), 3) # verify the number of rows
+ self.assertEqual(list(df), ["name", "geog"]) # verify the column names
+ self.assertEqual(df.name.dtype.name, "object")
+ self.assertEqual(df.geog.dtype.name, "geometry")
+ self.assertIsInstance(df.geog, geopandas.GeoSeries)
+ self.assertEqual(list(map(str, df.area)), ["0.0", "nan", "0.5"])
+ self.assertEqual(list(map(str, df.geog.area)), ["0.0", "nan", "0.5"])
+ self.assertEqual(df.crs.srs, "EPSG:4326")
+ self.assertEqual(df.crs.name, "WGS 84")
+ self.assertEqual(df.geog.crs.srs, "EPSG:4326")
+ self.assertEqual(df.geog.crs.name, "WGS 84")
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe_ambiguous_geog(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "GEOGRAPHY"), ("geog2", "GEOGRAPHY")), ()
+ )
+ with self.assertRaisesRegex(
+ ValueError,
+ re.escape(
+ "There is more than one GEOGRAPHY column in the result. "
+ "The geography_column argument must be used to specify which "
+ "one to use to create a GeoDataFrame"
+ ),
+ ):
+ row_iterator.to_geodataframe(create_bqstorage_client=False)
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe_bad_geography_column(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "GEOGRAPHY"), ("geog2", "GEOGRAPHY")), ()
+ )
+ with self.assertRaisesRegex(
+ ValueError,
+ re.escape(
+ "The given geography column, xxx, doesn't name"
+ " a GEOGRAPHY column in the result."
+ ),
+ ):
+ row_iterator.to_geodataframe(
+ create_bqstorage_client=False, geography_column="xxx"
+ )
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe_no_geog(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "STRING")), ()
+ )
+ with self.assertRaisesRegex(
+ TypeError,
+ re.escape(
+ "There must be at least one GEOGRAPHY column"
+ " to create a GeoDataFrame"
+ ),
+ ):
+ row_iterator.to_geodataframe(create_bqstorage_client=False)
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ def test_to_geodataframe_w_geography_column(self):
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("geog", "GEOGRAPHY"), ("geog2", "GEOGRAPHY")),
+ (
+ ("foo", "Point(0 0)", "Point(1 1)"),
+ ("bar", None, "Point(2 2)"),
+ ("baz", "Polygon((0 0, 0 1, 1 0, 0 0))", "Point(3 3)"),
+ ),
+ )
+ df = row_iterator.to_geodataframe(
+ create_bqstorage_client=False, geography_column="geog"
+ )
+ self.assertIsInstance(df, geopandas.GeoDataFrame)
+ self.assertEqual(len(df), 3) # verify the number of rows
+ self.assertEqual(list(df), ["name", "geog", "geog2"]) # verify the column names
+ self.assertEqual(df.name.dtype.name, "object")
+ self.assertEqual(df.geog.dtype.name, "geometry")
+ self.assertEqual(df.geog2.dtype.name, "object")
+ self.assertIsInstance(df.geog, geopandas.GeoSeries)
+ self.assertEqual(list(map(str, df.area)), ["0.0", "nan", "0.5"])
+ self.assertEqual(list(map(str, df.geog.area)), ["0.0", "nan", "0.5"])
+ self.assertEqual(
+ [v.__class__.__name__ for v in df.geog], ["Point", "NoneType", "Polygon"]
+ )
+
+ # Geog2 isn't a GeoSeries, but it contains geomentries:
+ self.assertIsInstance(df.geog2, pandas.Series)
+ self.assertEqual(
+ [v.__class__.__name__ for v in df.geog2], ["Point", "Point", "Point"]
+ )
+ # and can easily be converted to a GeoSeries
+ self.assertEqual(
+ list(map(str, geopandas.GeoSeries(df.geog2).area)), ["0.0", "0.0", "0.0"]
+ )
+
+ @unittest.skipIf(geopandas is None, "Requires `geopandas`")
+ @mock.patch("google.cloud.bigquery.table.RowIterator.to_dataframe")
+ def test_rowiterator_to_geodataframe_delegation(self, to_dataframe):
+ """
+ RowIterator.to_geodataframe just delegates to RowIterator.to_dataframe.
+
+ This test just demonstrates that. We don't need to test all the
+ variations, which are tested for to_dataframe.
+ """
+ import numpy
+ from shapely import wkt
+
+ row_iterator = self._make_one_from_data(
+ (("name", "STRING"), ("g", "GEOGRAPHY"))
+ )
+ bqstorage_client = object()
+ dtypes = dict(xxx=numpy.dtype("int64"))
+ progress_bar_type = "normal"
+ create_bqstorage_client = False
+ date_as_object = False
+ geography_column = "g"
+
+ to_dataframe.return_value = pandas.DataFrame(
+ dict(name=["foo"], g=[wkt.loads("point(0 0)")],)
+ )
+
+ df = row_iterator.to_geodataframe(
+ bqstorage_client=bqstorage_client,
+ dtypes=dtypes,
+ progress_bar_type=progress_bar_type,
+ create_bqstorage_client=create_bqstorage_client,
+ date_as_object=date_as_object,
+ geography_column=geography_column,
+ )
+
+ to_dataframe.assert_called_once_with(
+ bqstorage_client,
+ dtypes,
+ progress_bar_type,
+ create_bqstorage_client,
+ date_as_object,
+ geography_as_object=True,
+ )
+
+ self.assertIsInstance(df, geopandas.GeoDataFrame)
+ self.assertEqual(len(df), 1) # verify the number of rows
+ self.assertEqual(list(df), ["name", "g"]) # verify the column names
+ self.assertEqual(df.name.dtype.name, "object")
+ self.assertEqual(df.g.dtype.name, "geometry")
+ self.assertIsInstance(df.g, geopandas.GeoSeries)
+ self.assertEqual(list(map(str, df.area)), ["0.0"])
+ self.assertEqual(list(map(str, df.g.area)), ["0.0"])
+ self.assertEqual([v.__class__.__name__ for v in df.g], ["Point"])
+
class TestPartitionRange(unittest.TestCase):
def _get_target_class(self):