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)