From 7e137ec1757ea492fb2474e6a05dfb157f9a22c5 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 11 Oct 2018 15:11:05 -0700 Subject: [PATCH] Allow strings or references in create_dataset and create_table This makes it easier to create an empty dataset or table when all you care about is the ID and not any of the extra metadata. When creating an empty dataset, the default location is included. --- bigquery/google/cloud/bigquery/client.py | 39 ++++- bigquery/tests/unit/test_client.py | 194 +++++++++++++++++++++++ 2 files changed, 227 insertions(+), 6 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 47b655ab2345..4999337dbf3f 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -285,8 +285,14 @@ def create_dataset(self, dataset): https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/insert Args: - dataset (google.cloud.bigquery.dataset.Dataset): - A ``Dataset`` populated with the desired initial state. + dataset (Union[ \ + :class:`~google.cloud.bigquery.dataset.Dataset`, \ + :class:`~google.cloud.bigquery.dataset.DatasetReference`, \ + str, \ + ]): + A :class:`~google.cloud.bigquery.dataset.Dataset` to create. + If ``dataset`` is a reference, an empty dataset is created + with the specified ID and client's default location. Returns: google.cloud.bigquery.dataset.Dataset: @@ -300,6 +306,12 @@ def create_dataset(self, dataset): >>> dataset = client.create_dataset(dataset) """ + if isinstance(dataset, str): + dataset = DatasetReference.from_string( + dataset, default_project=self.project) + if isinstance(dataset, DatasetReference): + dataset = Dataset(dataset) + path = '/projects/%s/datasets' % (dataset.project,) data = dataset.to_api_repr() @@ -317,12 +329,27 @@ def create_table(self, table): See https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/insert - :type table: :class:`~google.cloud.bigquery.table.Table` - :param table: A ``Table`` populated with the desired initial state. + Args: + table (Union[ \ + :class:`~google.cloud.bigquery.table.Table`, \ + :class:`~google.cloud.bigquery.table.TableReference`, \ + str, \ + ]): + A :class:`~google.cloud.bigquery.table.Table` to create. + If ``table`` is a reference, an empty table is created + with the specified ID. The dataset that the table belongs to + must already exist. - :rtype: ":class:`~google.cloud.bigquery.table.Table`" - :returns: a new ``Table`` returned from the service. + Returns: + 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) + path = '/projects/%s/datasets/%s/tables' % ( table.project, table.dataset_id) api_response = self._connection.api_request( diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index 56f4aabeb651..60c81b59c482 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -711,6 +711,109 @@ def test_create_dataset_w_client_location_w_dataset_location(self): 'location': OTHER_LOCATION, }) + def test_create_dataset_w_reference(self): + path = '/projects/%s/datasets' % self.PROJECT + resource = { + 'datasetReference': + {'projectId': self.PROJECT, 'datasetId': self.DS_ID}, + 'etag': "etag", + 'id': "%s:%s" % (self.PROJECT, self.DS_ID), + 'location': self.LOCATION, + } + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION) + conn = client._connection = _make_connection(resource) + + dataset = client.create_dataset(client.dataset(self.DS_ID)) + + self.assertEqual(dataset.dataset_id, self.DS_ID) + self.assertEqual(dataset.project, self.PROJECT) + self.assertEqual(dataset.etag, resource['etag']) + self.assertEqual(dataset.full_dataset_id, resource['id']) + self.assertEqual(dataset.location, self.LOCATION) + + conn.api_request.assert_called_once_with( + method='POST', + path=path, + data={ + 'datasetReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + }, + 'labels': {}, + 'location': self.LOCATION, + }) + + def test_create_dataset_w_fully_qualified_string(self): + path = '/projects/%s/datasets' % self.PROJECT + resource = { + 'datasetReference': + {'projectId': self.PROJECT, 'datasetId': self.DS_ID}, + 'etag': "etag", + 'id': "%s:%s" % (self.PROJECT, self.DS_ID), + 'location': self.LOCATION, + } + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION) + conn = client._connection = _make_connection(resource) + + dataset = client.create_dataset( + '{}.{}'.format(self.PROJECT, self.DS_ID)) + + self.assertEqual(dataset.dataset_id, self.DS_ID) + self.assertEqual(dataset.project, self.PROJECT) + self.assertEqual(dataset.etag, resource['etag']) + self.assertEqual(dataset.full_dataset_id, resource['id']) + self.assertEqual(dataset.location, self.LOCATION) + + conn.api_request.assert_called_once_with( + method='POST', + path=path, + data={ + 'datasetReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + }, + 'labels': {}, + 'location': self.LOCATION, + }) + + def test_create_dataset_w_string(self): + path = '/projects/%s/datasets' % self.PROJECT + resource = { + 'datasetReference': + {'projectId': self.PROJECT, 'datasetId': self.DS_ID}, + 'etag': "etag", + 'id': "%s:%s" % (self.PROJECT, self.DS_ID), + 'location': self.LOCATION, + } + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION) + conn = client._connection = _make_connection(resource) + + dataset = client.create_dataset(self.DS_ID) + + self.assertEqual(dataset.dataset_id, self.DS_ID) + self.assertEqual(dataset.project, self.PROJECT) + self.assertEqual(dataset.etag, resource['etag']) + self.assertEqual(dataset.full_dataset_id, resource['id']) + self.assertEqual(dataset.location, self.LOCATION) + + conn.api_request.assert_called_once_with( + method='POST', + path=path, + data={ + 'datasetReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + }, + 'labels': {}, + 'location': self.LOCATION, + }) + def test_create_table_w_day_partition(self): from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning @@ -991,6 +1094,97 @@ def test_create_table_w_external(self): SourceFormat.CSV) self.assertEqual(got.external_data_configuration.autodetect, True) + def test_create_table_w_reference(self): + path = 'projects/%s/datasets/%s/tables' % ( + self.PROJECT, self.DS_ID) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + resource = { + 'id': '%s:%s:%s' % (self.PROJECT, self.DS_ID, self.TABLE_ID), + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + } + conn = client._connection = _make_connection(resource) + + got = client.create_table(self.TABLE_REF) + + conn.api_request.assert_called_once_with( + method='POST', + path='/%s' % path, + data={ + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + 'labels': {}, + }) + self.assertEqual(got.table_id, self.TABLE_ID) + + def test_create_table_w_fully_qualified_string(self): + path = 'projects/%s/datasets/%s/tables' % ( + self.PROJECT, self.DS_ID) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + resource = { + 'id': '%s:%s:%s' % (self.PROJECT, self.DS_ID, self.TABLE_ID), + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + } + conn = client._connection = _make_connection(resource) + + got = client.create_table( + '{}.{}.{}'.format(self.PROJECT, self.DS_ID, self.TABLE_ID)) + + conn.api_request.assert_called_once_with( + method='POST', + path='/%s' % path, + data={ + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + 'labels': {}, + }) + self.assertEqual(got.table_id, self.TABLE_ID) + + def test_create_table_w_string(self): + path = 'projects/%s/datasets/%s/tables' % ( + self.PROJECT, self.DS_ID) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + resource = { + 'id': '%s:%s:%s' % (self.PROJECT, self.DS_ID, self.TABLE_ID), + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + } + conn = client._connection = _make_connection(resource) + + got = client.create_table('{}.{}'.format(self.DS_ID, self.TABLE_ID)) + + conn.api_request.assert_called_once_with( + method='POST', + path='/%s' % path, + data={ + 'tableReference': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': self.TABLE_ID + }, + 'labels': {}, + }) + self.assertEqual(got.table_id, self.TABLE_ID) + def test_get_table(self): path = 'projects/%s/datasets/%s/tables/%s' % ( self.PROJECT, self.DS_ID, self.TABLE_ID)