diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 7fe412478bfa..fd69502f6ea1 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -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 @@ -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): @@ -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() @@ -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, \ ]): @@ -446,11 +448,7 @@ def get_table(self, table_ref, retry=DEFAULT_RETRY): google.cloud.bigquery.table.Table: A ``Table`` instance. """ - 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) @@ -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) @@ -925,6 +921,7 @@ def load_table_from_uri( URIs of data files to be loaded; in format ``gs:///``. destination (Union[ \ + :class:`~google.cloud.bigquery.table.Table`, \ :class:`~google.cloud.bigquery.table.TableReference`, \ str, \ ]): @@ -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) @@ -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, \ ]): @@ -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() @@ -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, \ ]): @@ -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 ) @@ -1364,6 +1367,7 @@ def extract_table( Arguments: source (Union[ \ + :class:`google.cloud.bigquery.table.Table`, \ :class:`google.cloud.bigquery.table.TableReference`, \ src, \ ]): @@ -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] @@ -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): + 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] @@ -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} @@ -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), @@ -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: diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 4df7a92ba63c..f4c919a2d8fc 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -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 @@ -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 @@ -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 """ @@ -2009,9 +2021,17 @@ 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 @@ -2019,6 +2039,15 @@ 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 """ @@ -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 diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index 35f950e7d94a..46419e70a83e 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -1661,3 +1661,28 @@ def _rows_page_start(iterator, page, response): # pylint: enable=unused-argument + + +def _table_arg_to_table_ref(value, default_project=None): + """Helper to convert a string or Table to TableReference. + + This function keeps TableReference and other kinds of objects unchanged. + """ + if isinstance(value, six.string_types): + value = TableReference.from_string(value, default_project=default_project) + if isinstance(value, (Table, TableListItem)): + value = value.reference + return value + + +def _table_arg_to_table(value, default_project=None): + """Helper to convert a string or TableReference to a Table. + + This function keeps Table and other kinds of objects unchanged. + """ + if isinstance(value, six.string_types): + value = TableReference.from_string(value, default_project=default_project) + if isinstance(value, TableReference): + value = Table(value) + + return value diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index 794b76a0a9f4..780f13bcbe3e 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -2820,6 +2820,33 @@ def test_copy_table_w_client_location(self): method="POST", path="/projects/other-project/jobs", data=resource ) + def test_copy_table_w_source_strings(self): + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + client._connection = _make_connection({}) + sources = [ + "dataset_wo_proj.some_table", + "other_project.other_dataset.other_table", + client.dataset("dataset_from_ref").table("table_from_ref"), + ] + destination = "some_project.some_dataset.destination_table" + + job = client.copy_table(sources, destination) + + expected_sources = [ + client.dataset("dataset_wo_proj").table("some_table"), + client.dataset("other_dataset", project="other_project").table( + "other_table" + ), + client.dataset("dataset_from_ref").table("table_from_ref"), + ] + self.assertEqual(list(job.sources), expected_sources) + expected_destination = client.dataset( + "some_dataset", project="some_project" + ).table("destination_table") + self.assertEqual(job.destination, expected_destination) + def test_extract_table(self): from google.cloud.bigquery.job import ExtractJob @@ -3400,7 +3427,7 @@ def test_query_w_query_parameters(self): ) def test_insert_rows_wo_schema(self): - from google.cloud.bigquery.table import Table, _TABLE_HAS_NO_SCHEMA + from google.cloud.bigquery.table import Table creds = _make_credentials() http = object() @@ -3416,7 +3443,7 @@ def test_insert_rows_wo_schema(self): with self.assertRaises(ValueError) as exc: client.insert_rows(table, ROWS) - self.assertEqual(exc.exception.args, (_TABLE_HAS_NO_SCHEMA,)) + self.assertIn("Could not determine schema for table", exc.exception.args[0]) def test_insert_rows_w_schema(self): import datetime diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 5023c41e6374..833081ce066d 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -3163,6 +3163,49 @@ def test_ctor_w_properties(self): self.assertFalse(config.use_query_cache) self.assertTrue(config.use_legacy_sql) + def test_ctor_w_string_default_dataset(self): + from google.cloud.bigquery import dataset + + default_dataset = "default-proj.default_dset" + config = self._get_target_class()(default_dataset=default_dataset) + expected = dataset.DatasetReference.from_string(default_dataset) + self.assertEqual(config.default_dataset, expected) + + def test_ctor_w_string_destinaton(self): + from google.cloud.bigquery import table + + destination = "dest-proj.dest_dset.dest_tbl" + config = self._get_target_class()(destination=destination) + expected = table.TableReference.from_string(destination) + self.assertEqual(config.destination, expected) + + def test_default_dataset_w_string(self): + from google.cloud.bigquery import dataset + + default_dataset = "default-proj.default_dset" + config = self._make_one() + config.default_dataset = default_dataset + expected = dataset.DatasetReference.from_string(default_dataset) + self.assertEqual(config.default_dataset, expected) + + def test_default_dataset_w_dataset(self): + from google.cloud.bigquery import dataset + + default_dataset = "default-proj.default_dset" + expected = dataset.DatasetReference.from_string(default_dataset) + config = self._make_one() + config.default_dataset = dataset.Dataset(expected) + self.assertEqual(config.default_dataset, expected) + + def test_destinaton_w_string(self): + from google.cloud.bigquery import table + + destination = "dest-proj.dest_dset.dest_tbl" + config = self._make_one() + config.destination = destination + expected = table.TableReference.from_string(destination) + self.assertEqual(config.destination, expected) + def test_time_partitioning(self): from google.cloud.bigquery import table