diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 96f1310c3f99..95d49e12968a 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -32,6 +32,7 @@ from google.resumable_media.requests import MultipartUpload from google.resumable_media.requests import ResumableUpload +import google.api_core.exceptions from google.api_core import page_iterator import google.cloud._helpers from google.cloud import exceptions @@ -297,7 +298,7 @@ def dataset(self, dataset_id, project=None): return DatasetReference(project, dataset_id) - def create_dataset(self, dataset): + def create_dataset(self, dataset, exists_ok=False, retry=DEFAULT_RETRY): """API call: create the dataset via a POST request. See @@ -312,6 +313,11 @@ def create_dataset(self, dataset): 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. + exists_ok (bool): + Defaults to ``False``. If ``True``, ignore "already exists" + errors when creating the dataset. + retry (google.api_core.retry.Retry): + Optional. How to retry the RPC. Returns: google.cloud.bigquery.dataset.Dataset: @@ -338,11 +344,15 @@ def create_dataset(self, dataset): if data.get("location") is None and self.location is not None: data["location"] = self.location - api_response = self._connection.api_request(method="POST", path=path, data=data) - - return Dataset.from_api_repr(api_response) - - def create_table(self, table): + try: + api_response = self._call_api(retry, method="POST", path=path, data=data) + return Dataset.from_api_repr(api_response) + except google.api_core.exceptions.Conflict: + if not exists_ok: + raise + return self.get_dataset(dataset.reference, retry=retry) + + def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY): """API call: create a table via a PUT request See @@ -358,6 +368,11 @@ def create_table(self, table): If ``table`` is a reference, an empty table is created with the specified ID. The dataset that the table belongs to must already exist. + exists_ok (bool): + Defaults to ``False``. If ``True``, ignore "already exists" + errors when creating the table. + retry (google.api_core.retry.Retry): + Optional. How to retry the RPC. Returns: google.cloud.bigquery.table.Table: @@ -369,10 +384,14 @@ def create_table(self, table): table = Table(table) path = "/projects/%s/datasets/%s/tables" % (table.project, table.dataset_id) - api_response = self._connection.api_request( - method="POST", path=path, data=table.to_api_repr() - ) - return Table.from_api_repr(api_response) + data = table.to_api_repr() + try: + api_response = self._call_api(retry, method="POST", path=path, data=data) + return Table.from_api_repr(api_response) + except google.api_core.exceptions.Conflict: + if not exists_ok: + raise + return self.get_table(table.reference, retry=retry) def _call_api(self, retry, **kwargs): call = functools.partial(self._connection.api_request, **kwargs) @@ -563,7 +582,9 @@ def list_tables( result.dataset = dataset return result - def delete_dataset(self, dataset, delete_contents=False, retry=DEFAULT_RETRY): + def delete_dataset( + self, dataset, delete_contents=False, retry=DEFAULT_RETRY, not_found_ok=False + ): """Delete a dataset. See @@ -579,12 +600,15 @@ def delete_dataset(self, dataset, delete_contents=False, retry=DEFAULT_RETRY): in, this method attempts to create a dataset reference from a string using :func:`google.cloud.bigquery.dataset.DatasetReference.from_string`. - retry (:class:`google.api_core.retry.Retry`): - (Optional) How to retry the RPC. delete_contents (boolean): (Optional) If True, delete all the tables in the dataset. If False and the dataset contains tables, the request will fail. Default is False. + retry (:class:`google.api_core.retry.Retry`): + (Optional) How to retry the RPC. + not_found_ok (bool): + Defaults to ``False``. If ``True``, ignore "not found" errors + when deleting the dataset. """ if isinstance(dataset, str): dataset = DatasetReference.from_string( @@ -598,9 +622,15 @@ def delete_dataset(self, dataset, delete_contents=False, retry=DEFAULT_RETRY): if delete_contents: params["deleteContents"] = "true" - self._call_api(retry, method="DELETE", path=dataset.path, query_params=params) + try: + self._call_api( + retry, method="DELETE", path=dataset.path, query_params=params + ) + except google.api_core.exceptions.NotFound: + if not not_found_ok: + raise - def delete_table(self, table, retry=DEFAULT_RETRY): + def delete_table(self, table, retry=DEFAULT_RETRY, not_found_ok=False): """Delete a table See @@ -618,13 +648,21 @@ def delete_table(self, table, retry=DEFAULT_RETRY): :func:`google.cloud.bigquery.table.TableReference.from_string`. retry (:class:`google.api_core.retry.Retry`): (Optional) How to retry the RPC. + not_found_ok (bool): + 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") - self._call_api(retry, method="DELETE", path=table.path) + + try: + self._call_api(retry, method="DELETE", path=table.path) + except google.api_core.exceptions.NotFound: + if not not_found_ok: + raise def _get_query_results( self, job_id, retry, project=None, timeout_ms=None, location=None diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index c3d90ed640fb..a98ee79aa116 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -36,6 +36,7 @@ except (ImportError, AttributeError): # pragma: NO COVER pyarrow = None +import google.api_core.exceptions from google.cloud.bigquery.dataset import DatasetReference @@ -804,6 +805,61 @@ def test_create_dataset_w_string(self): }, ) + def test_create_dataset_alreadyexists_w_exists_ok_false(self): + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION + ) + client._connection = _make_connection( + google.api_core.exceptions.AlreadyExists("dataset already exists") + ) + + with pytest.raises(google.api_core.exceptions.AlreadyExists): + client.create_dataset(self.DS_ID) + + def test_create_dataset_alreadyexists_w_exists_ok_true(self): + post_path = "/projects/{}/datasets".format(self.PROJECT) + get_path = "/projects/{}/datasets/{}".format(self.PROJECT, self.DS_ID) + resource = { + "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID}, + "etag": "etag", + "id": "{}:{}".format(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( + google.api_core.exceptions.AlreadyExists("dataset already exists"), resource + ) + + dataset = client.create_dataset(self.DS_ID, exists_ok=True) + + 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_has_calls( + [ + mock.call( + method="POST", + path=post_path, + data={ + "datasetReference": { + "projectId": self.PROJECT, + "datasetId": self.DS_ID, + }, + "labels": {}, + "location": self.LOCATION, + }, + ), + mock.call(method="GET", path=get_path), + ] + ) + def test_create_table_w_day_partition(self): from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning @@ -1177,6 +1233,79 @@ def test_create_table_w_string(self): ) self.assertEqual(got.table_id, self.TABLE_ID) + def test_create_table_alreadyexists_w_exists_ok_false(self): + post_path = "/projects/{}/datasets/{}/tables".format(self.PROJECT, self.DS_ID) + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION + ) + conn = client._connection = _make_connection( + google.api_core.exceptions.AlreadyExists("table already exists") + ) + + with pytest.raises(google.api_core.exceptions.AlreadyExists): + client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + + conn.api_request.assert_called_once_with( + method="POST", + path=post_path, + data={ + "tableReference": { + "projectId": self.PROJECT, + "datasetId": self.DS_ID, + "tableId": self.TABLE_ID, + }, + "labels": {}, + }, + ) + + def test_create_table_alreadyexists_w_exists_ok_true(self): + post_path = "/projects/{}/datasets/{}/tables".format(self.PROJECT, self.DS_ID) + get_path = "/projects/{}/datasets/{}/tables/{}".format( + self.PROJECT, self.DS_ID, self.TABLE_ID + ) + 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, + }, + } + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION + ) + conn = client._connection = _make_connection( + google.api_core.exceptions.AlreadyExists("table already exists"), resource + ) + + got = client.create_table( + "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True + ) + + self.assertEqual(got.project, self.PROJECT) + self.assertEqual(got.dataset_id, self.DS_ID) + self.assertEqual(got.table_id, self.TABLE_ID) + + conn.api_request.assert_has_calls( + [ + mock.call( + method="POST", + path=post_path, + data={ + "tableReference": { + "projectId": self.PROJECT, + "datasetId": self.DS_ID, + "tableId": self.TABLE_ID, + }, + "labels": {}, + }, + ), + mock.call(method="GET", path=get_path), + ] + ) + def test_get_table(self): path = "projects/%s/datasets/%s/tables/%s" % ( self.PROJECT, @@ -1804,6 +1933,33 @@ def test_delete_dataset_wrong_type(self): with self.assertRaises(TypeError): client.delete_dataset(client.dataset(self.DS_ID).table("foo")) + def test_delete_dataset_w_not_found_ok_false(self): + path = "/projects/{}/datasets/{}".format(self.PROJECT, self.DS_ID) + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + conn = client._connection = _make_connection( + google.api_core.exceptions.NotFound("dataset not found") + ) + + with self.assertRaises(google.api_core.exceptions.NotFound): + client.delete_dataset(self.DS_ID) + + conn.api_request.assert_called_with(method="DELETE", path=path, query_params={}) + + def test_delete_dataset_w_not_found_ok_true(self): + path = "/projects/{}/datasets/{}".format(self.PROJECT, self.DS_ID) + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + conn = client._connection = _make_connection( + google.api_core.exceptions.NotFound("dataset not found") + ) + + client.delete_dataset(self.DS_ID, not_found_ok=True) + + conn.api_request.assert_called_with(method="DELETE", path=path, query_params={}) + def test_delete_table(self): from google.cloud.bigquery.table import Table @@ -1836,6 +1992,39 @@ def test_delete_table_w_wrong_type(self): with self.assertRaises(TypeError): client.delete_table(client.dataset(self.DS_ID)) + def test_delete_table_w_not_found_ok_false(self): + path = "/projects/{}/datasets/{}/tables/{}".format( + self.PROJECT, self.DS_ID, self.TABLE_ID + ) + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + conn = client._connection = _make_connection( + google.api_core.exceptions.NotFound("table not found") + ) + + with self.assertRaises(google.api_core.exceptions.NotFound): + client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + + conn.api_request.assert_called_with(method="DELETE", path=path) + + def test_delete_table_w_not_found_ok_true(self): + path = "/projects/{}/datasets/{}/tables/{}".format( + self.PROJECT, self.DS_ID, self.TABLE_ID + ) + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + conn = client._connection = _make_connection( + google.api_core.exceptions.NotFound("table not found") + ) + + client.delete_table( + "{}.{}".format(self.DS_ID, self.TABLE_ID), not_found_ok=True + ) + + conn.api_request.assert_called_with(method="DELETE", path=path) + def test_job_from_resource_unknown_type(self): from google.cloud.bigquery.job import UnknownJob