Skip to content

Allow string for copy sources, query destination, and default dataset #7560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JustinBeckwith merged 4 commits into from
Mar 27, 2019
Merged

Allow string for copy sources, query destination, and default dataset #7560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8681335
Allow string for copy sources, query destination, and default dataset
tswast Mar 26, 2019
3ab95d4
Refactor str/table/ref to ref helper
tswast Mar 26, 2019
84dbed9
Rename table_ref arg to table. Fix lint errors.
tswast Mar 27, 2019
fd8b06d
Fix typo in docstring.
tswast Mar 27, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 77 additions & 65 deletions bigquery/google/cloud/bigquery/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,11 +48,12 @@
from google.cloud.bigquery import job
from google.cloud.bigquery.query import _QueryResults
from google.cloud.bigquery.retry import DEFAULT_RETRY
from google.cloud.bigquery.table import _table_arg_to_table
from google.cloud.bigquery.table import _table_arg_to_table_ref
from google.cloud.bigquery.table import Table
from google.cloud.bigquery.table import TableListItem
from google.cloud.bigquery.table import TableReference
from google.cloud.bigquery.table import RowIterator
from google.cloud.bigquery.table import _TABLE_HAS_NO_SCHEMA


_DEFAULT_CHUNKSIZE = 1048576 # 1024 * 1024 B = 1 MB
Expand All @@ -68,6 +69,9 @@
_READ_LESS_THAN_SIZE = (
"Size {:d} was specified but the file-like object only had " "{:d} bytes remaining."
)
_NEED_TABLE_ARGUMENT = (
"The table argument should be a table ID string, Table, or TableReference"
)


class Project(object):
Expand Down Expand Up @@ -379,10 +383,7 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY):
google.cloud.bigquery.table.Table:
A new ``Table`` returned from the service.
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)
if isinstance(table, TableReference):
table = Table(table)
table = _table_arg_to_table(table, default_project=self.project)

path = "/projects/%s/datasets/%s/tables" % (table.project, table.dataset_id)
data = table.to_api_repr()
Expand Down Expand Up @@ -427,11 +428,12 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY):
api_response = self._call_api(retry, method="GET", path=dataset_ref.path)
return Dataset.from_api_repr(api_response)

def get_table(self, table_ref, retry=DEFAULT_RETRY):
"""Fetch the table referenced by ``table_ref``.
def get_table(self, table, retry=DEFAULT_RETRY):
"""Fetch the table referenced by ``table``.

Args:
table_ref (Union[ \
table (Union[ \
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
]):
Expand All @@ -446,11 +448,7 @@ def get_table(self, table_ref, retry=DEFAULT_RETRY):
google.cloud.bigquery.table.Table:
A ``Table`` instance.
"""
Comment thread
alixhami marked this conversation as resolved.
if isinstance(table_ref, str):
table_ref = TableReference.from_string(
table_ref, default_project=self.project
)

table_ref = _table_arg_to_table_ref(table, default_project=self.project)
api_response = self._call_api(retry, method="GET", path=table_ref.path)
return Table.from_api_repr(api_response)

Expand Down Expand Up @@ -653,11 +651,9 @@ def delete_table(self, table, retry=DEFAULT_RETRY, not_found_ok=False):
Defaults to ``False``. If ``True``, ignore "not found" errors
when deleting the table.
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)

if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or a TableReference")
table = _table_arg_to_table_ref(table, default_project=self.project)
if not isinstance(table, TableReference):
raise TypeError("Unable to get TableReference for table '{}'".format(table))

try:
self._call_api(retry, method="DELETE", path=table.path)
Expand Down Expand Up @@ -925,6 +921,7 @@ def load_table_from_uri(
URIs of data files to be loaded; in format
``gs://<bucket_name>/<object_name_or_glob>``.
destination (Union[ \
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
]):
Expand Down Expand Up @@ -966,11 +963,7 @@ def load_table_from_uri(
if isinstance(source_uris, six.string_types):
source_uris = [source_uris]

if isinstance(destination, str):
destination = TableReference.from_string(
destination, default_project=self.project
)

destination = _table_arg_to_table_ref(destination, default_project=self.project)
load_job = job.LoadJob(job_ref, source_uris, destination, self, job_config)
load_job._begin(retry=retry)

Expand All @@ -997,6 +990,7 @@ def load_table_from_file(
Arguments:
file_obj (file): A file handle opened in binary mode for reading.
destination (Union[ \
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
]):
Expand Down Expand Up @@ -1045,11 +1039,7 @@ def load_table_from_file(
if location is None:
location = self.location

if isinstance(destination, str):
destination = TableReference.from_string(
destination, default_project=self.project
)

destination = _table_arg_to_table_ref(destination, default_project=self.project)
job_ref = job._JobReference(job_id, project=project, location=location)
load_job = job.LoadJob(job_ref, None, destination, self, job_config)
job_resource = load_job.to_api_repr()
Expand Down Expand Up @@ -1286,13 +1276,20 @@ def copy_table(

Arguments:
sources (Union[ \
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
Sequence[ \
:class:`~google.cloud.bigquery.table.TableReference`], \
Union[ \
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
] \
], \
]):
Table or tables to be copied.
destination (Union[
:class:`~google.cloud.bigquery.table.Table`, \
:class:`~google.cloud.bigquery.table.TableReference`, \
str, \
]):
Expand Down Expand Up @@ -1328,17 +1325,23 @@ def copy_table(

job_ref = job._JobReference(job_id, project=project, location=location)

if isinstance(sources, str):
sources = TableReference.from_string(sources, default_project=self.project)

if isinstance(destination, str):
destination = TableReference.from_string(
destination, default_project=self.project
)
# sources can be one of many different input types. (string, Table,
# TableReference, or a sequence of any of those.) Convert them all to a
# list of TableReferences.
#
# _table_arg_to_table_ref leaves lists unmodified.
sources = _table_arg_to_table_ref(sources, default_project=self.project)

if not isinstance(sources, collections_abc.Sequence):
sources = [sources]

sources = [
_table_arg_to_table_ref(source, default_project=self.project)
for source in sources
]

destination = _table_arg_to_table_ref(destination, default_project=self.project)

copy_job = job.CopyJob(
job_ref, sources, destination, client=self, job_config=job_config
)
Expand All @@ -1364,6 +1367,7 @@ def extract_table(

Arguments:
source (Union[ \
:class:`google.cloud.bigquery.table.Table`, \
:class:`google.cloud.bigquery.table.TableReference`, \
src, \
]):
Expand Down Expand Up @@ -1405,9 +1409,7 @@ def extract_table(
location = self.location

job_ref = job._JobReference(job_id, project=project, location=location)

if isinstance(source, str):
source = TableReference.from_string(source, default_project=self.project)
source = _table_arg_to_table_ref(source, default_project=self.project)

if isinstance(destination_uris, six.string_types):
destination_uris = [destination_uris]
Expand Down Expand Up @@ -1529,19 +1531,24 @@ def insert_rows(self, table, rows, selected_fields=None, **kwargs):
Raises:
ValueError: if table's schema is not set
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)
table = _table_arg_to_table(table, default_project=self.project)

if not isinstance(table, Table):
Comment thread
alixhami marked this conversation as resolved.
raise TypeError(_NEED_TABLE_ARGUMENT)

schema = table.schema

# selected_fields can override the table schema.
if selected_fields is not None:
schema = selected_fields
elif isinstance(table, TableReference):
raise ValueError("need selected_fields with TableReference")
elif isinstance(table, Table):
if len(table.schema) == 0:
raise ValueError(_TABLE_HAS_NO_SCHEMA)
schema = table.schema
else:
raise TypeError("table should be Table or TableReference")

if len(schema) == 0:
raise ValueError(
(
"Could not determine schema for table '{}'. Call client.get_table() "
"or pass in a list of schema fields to the selected_fields argument."
).format(table)
)

json_rows = [_record_field_to_json(schema, row) for row in rows]

Expand Down Expand Up @@ -1596,9 +1603,10 @@ def insert_rows_json(
identifies the row, and the "errors" key contains a list of
the mappings describing one or more problems with the row.
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)

# Convert table to just a reference because unlike insert_rows,
# insert_rows_json doesn't need the table schema. It's not doing any
# type conversions.
table = _table_arg_to_table_ref(table, default_project=self.project)
rows_info = []
data = {"rows": rows_info}

Expand Down Expand Up @@ -1647,9 +1655,7 @@ def list_partitions(self, table, retry=DEFAULT_RETRY):
List[str]:
A list of the partition ids present in the partitioned table
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)

table = _table_arg_to_table_ref(table, default_project=self.project)
meta_table = self.get_table(
TableReference(
self.dataset(table.dataset_id, project=table.project),
Expand Down Expand Up @@ -1724,19 +1730,25 @@ def list_rows(
(this is distinct from the total number of rows in the
current page: ``iterator.page.num_items``).
"""
if isinstance(table, str):
table = TableReference.from_string(table, default_project=self.project)
table = _table_arg_to_table(table, default_project=self.project)

if not isinstance(table, Table):
raise TypeError(_NEED_TABLE_ARGUMENT)

schema = table.schema

# selected_fields can override the table schema.
if selected_fields is not None:
schema = selected_fields
elif isinstance(table, TableReference):
raise ValueError("need selected_fields with TableReference")
elif isinstance(table, Table):
if len(table.schema) == 0 and table.created is None:
raise ValueError(_TABLE_HAS_NO_SCHEMA)
schema = table.schema
else:
raise TypeError("table should be Table or TableReference")

# Allow listing rows of an empty table by not raising if the table exists.
elif len(schema) == 0 and table.created is None:
raise ValueError(
(
"Could not determine schema for table '{}'. Call client.get_table() "
"or pass in a list of schema fields to the selected_fields argument."
).format(table)
)

params = {}
if selected_fields is not None:
Expand Down
44 changes: 38 additions & 6 deletions bigquery/google/cloud/bigquery/job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,14 @@
import copy
import threading

import six
from six.moves import http_client

import google.api_core.future.polling
from google.cloud import exceptions
from google.cloud.exceptions import NotFound
from google.cloud.bigquery.dataset import Dataset
from google.cloud.bigquery.dataset import DatasetListItem
from google.cloud.bigquery.dataset import DatasetReference
from google.cloud.bigquery.external_config import ExternalConfig
from google.cloud.bigquery.query import _query_param_from_api_repr
Expand All @@ -33,6 +36,7 @@
from google.cloud.bigquery.schema import SchemaField
from google.cloud.bigquery.table import _EmptyRowIterator
from google.cloud.bigquery.table import EncryptionConfiguration
from google.cloud.bigquery.table import _table_arg_to_table_ref
from google.cloud.bigquery.table import TableReference
from google.cloud.bigquery.table import Table
from google.cloud.bigquery.table import TimePartitioning
Expand Down Expand Up @@ -1999,6 +2003,14 @@ def default_dataset(self):
to use for unqualified table names in the query or :data:`None` if not
set.

The ``default_dataset`` setter accepts:

- a :class:`~google.cloud.bigquery.dataset.Dataset`, or
- a :class:`~google.cloud.bigquery.dataset.DatasetReference`, or
- a :class:`str` of the fully-qualified dataset ID in standard SQL
format. The value must included a project ID and dataset ID
separated by ``.``. For example: ``your-project.your_dataset``.

See
https://g.co/cloud/bigquery/docs/reference/v2/jobs#configuration.query.defaultDataset
"""
Expand All @@ -2009,16 +2021,33 @@ def default_dataset(self):

@default_dataset.setter
def default_dataset(self, value):
resource = None
if value is not None:
resource = value.to_api_repr()
if value is None:
self._set_sub_prop("defaultDataset", None)
return

if isinstance(value, six.string_types):
value = DatasetReference.from_string(value)

if isinstance(value, (Dataset, DatasetListItem)):
value = value.reference

resource = value.to_api_repr()
self._set_sub_prop("defaultDataset", resource)

@property
def destination(self):
"""google.cloud.bigquery.table.TableReference: table where results are
written or :data:`None` if not set.

The ``destination`` setter accepts:

- a :class:`~google.cloud.bigquery.table.Table`, or
- a :class:`~google.cloud.bigquery.table.TableReference`, or
- a :class:`str` of the fully-qualified table ID in standard SQL
format. The value must included a project ID, dataset ID, and table
ID, each separated by ``.``. For example:
``your-project.your_dataset.your_table``.

See
https://g.co/cloud/bigquery/docs/reference/rest/v2/jobs#configuration.query.destinationTable
"""
Expand All @@ -2029,9 +2058,12 @@ def destination(self):

@destination.setter
def destination(self, value):
resource = None
if value is not None:
resource = value.to_api_repr()
if value is None:
self._set_sub_prop("destinationTable", None)
return

value = _table_arg_to_table_ref(value)
resource = value.to_api_repr()
self._set_sub_prop("destinationTable", resource)

@property
Expand Down
Loading