Skip to content

Commit c8b5581

Browse files
jimfultonjimfulton
authored
feat: accept DatasetListItem where DatasetReference is accepted (googleapis#597)
* split out and pytestify list_tables tests. Also, exercise dataset polymorphism in some of the tests. * list_tables now accepts DatasetListItem objects * Get coverage to 100% But why do we run coverage on test code? * lint * Update exception text for DatasetListItem * Bypass opentelemetry tracing in unit tests. * Got rid of opentelemetry tracing checks. They aren't needed. * abstracted dataset-argument handling And applied it to `list_tables` and `list_models`. * Converted list_model tests to pytest and included check for dataset polymorphism * removed unneeded blanl lines. * Made list_routines accept DatasetListItem and conveted list_routines tests to pytest. * create_dataset accepts DatasetListItem Also converted create_dataset tests to pytest. (And fixed some long lines.) * Converted list_routine tests to pytest * include string dataset representation in dataset polymorphism. * removed some unused imports * Updated delete_dataset tests - Polymorphoc on dataset - pytest * black * lint * We don't actually need to avoid opentelemetry And a 3.6 test dependened on it. * fixed docstrings to include DatasetListItem in dataset polymorphic APIs.
1 parent 8f4c0b8 commit c8b5581

File tree

9 files changed

+2005
-2099
lines changed

9 files changed

+2005
-2099
lines changed

‎google/cloud/bigquery/client.py‎

Lines changed: 26 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,22 @@ def _create_bqstorage_client(self):
449449

450450
return bigquery_storage.BigQueryReadClient(credentials=self._credentials)
451451

452+
def _dataset_from_arg(self, dataset):
453+
if isinstance(dataset, str):
454+
dataset = DatasetReference.from_string(
455+
dataset, default_project=self.project
456+
)
457+
458+
if not isinstance(dataset, (Dataset, DatasetReference)):
459+
if isinstance(dataset, DatasetListItem):
460+
dataset = dataset.reference
461+
else:
462+
raise TypeError(
463+
"dataset must be a Dataset, DatasetReference, DatasetListItem,"
464+
" or string"
465+
)
466+
return dataset
467+
452468
def create_dataset(
453469
self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None
454470
):
@@ -461,6 +477,7 @@ def create_dataset(
461477
dataset (Union[ \
462478
google.cloud.bigquery.dataset.Dataset, \
463479
google.cloud.bigquery.dataset.DatasetReference, \
480+
google.cloud.bigquery.dataset.DatasetListItem, \
464481
str, \
465482
]):
466483
A :class:`~google.cloud.bigquery.dataset.Dataset` to create.
@@ -491,10 +508,7 @@ def create_dataset(
491508
>>> dataset = client.create_dataset(dataset)
492509
493510
"""
494-
if isinstance(dataset, str):
495-
dataset = DatasetReference.from_string(
496-
dataset, default_project=self.project
497-
)
511+
dataset = self._dataset_from_arg(dataset)
498512
if isinstance(dataset, DatasetReference):
499513
dataset = Dataset(dataset)
500514

@@ -1133,6 +1147,7 @@ def list_models(
11331147
dataset (Union[ \
11341148
google.cloud.bigquery.dataset.Dataset, \
11351149
google.cloud.bigquery.dataset.DatasetReference, \
1150+
google.cloud.bigquery.dataset.DatasetListItem, \
11361151
str, \
11371152
]):
11381153
A reference to the dataset whose models to list from the
@@ -1160,13 +1175,7 @@ def list_models(
11601175
:class:`~google.cloud.bigquery.model.Model` contained
11611176
within the requested dataset.
11621177
"""
1163-
if isinstance(dataset, str):
1164-
dataset = DatasetReference.from_string(
1165-
dataset, default_project=self.project
1166-
)
1167-
1168-
if not isinstance(dataset, (Dataset, DatasetReference)):
1169-
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
1178+
dataset = self._dataset_from_arg(dataset)
11701179

11711180
path = "%s/models" % dataset.path
11721181
span_attributes = {"path": path}
@@ -1210,6 +1219,7 @@ def list_routines(
12101219
dataset (Union[ \
12111220
google.cloud.bigquery.dataset.Dataset, \
12121221
google.cloud.bigquery.dataset.DatasetReference, \
1222+
google.cloud.bigquery.dataset.DatasetListItem, \
12131223
str, \
12141224
]):
12151225
A reference to the dataset whose routines to list from the
@@ -1237,14 +1247,7 @@ def list_routines(
12371247
:class:`~google.cloud.bigquery.routine.Routine`s contained
12381248
within the requested dataset, limited by ``max_results``.
12391249
"""
1240-
if isinstance(dataset, str):
1241-
dataset = DatasetReference.from_string(
1242-
dataset, default_project=self.project
1243-
)
1244-
1245-
if not isinstance(dataset, (Dataset, DatasetReference)):
1246-
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
1247-
1250+
dataset = self._dataset_from_arg(dataset)
12481251
path = "{}/routines".format(dataset.path)
12491252

12501253
span_attributes = {"path": path}
@@ -1288,6 +1291,7 @@ def list_tables(
12881291
dataset (Union[ \
12891292
google.cloud.bigquery.dataset.Dataset, \
12901293
google.cloud.bigquery.dataset.DatasetReference, \
1294+
google.cloud.bigquery.dataset.DatasetListItem, \
12911295
str, \
12921296
]):
12931297
A reference to the dataset whose tables to list from the
@@ -1315,14 +1319,7 @@ def list_tables(
13151319
:class:`~google.cloud.bigquery.table.TableListItem` contained
13161320
within the requested dataset.
13171321
"""
1318-
if isinstance(dataset, str):
1319-
dataset = DatasetReference.from_string(
1320-
dataset, default_project=self.project
1321-
)
1322-
1323-
if not isinstance(dataset, (Dataset, DatasetReference)):
1324-
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
1325-
1322+
dataset = self._dataset_from_arg(dataset)
13261323
path = "%s/tables" % dataset.path
13271324
span_attributes = {"path": path}
13281325

@@ -1365,6 +1362,7 @@ def delete_dataset(
13651362
dataset (Union[ \
13661363
google.cloud.bigquery.dataset.Dataset, \
13671364
google.cloud.bigquery.dataset.DatasetReference, \
1365+
google.cloud.bigquery.dataset.DatasetListItem, \
13681366
str, \
13691367
]):
13701368
A reference to the dataset to delete. If a string is passed
@@ -1384,14 +1382,7 @@ def delete_dataset(
13841382
Defaults to ``False``. If ``True``, ignore "not found" errors
13851383
when deleting the dataset.
13861384
"""
1387-
if isinstance(dataset, str):
1388-
dataset = DatasetReference.from_string(
1389-
dataset, default_project=self.project
1390-
)
1391-
1392-
if not isinstance(dataset, (Dataset, DatasetReference)):
1393-
raise TypeError("dataset must be a Dataset or a DatasetReference")
1394-
1385+
dataset = self._dataset_from_arg(dataset)
13951386
params = {}
13961387
path = dataset.path
13971388
if delete_contents:

‎tests/unit/conftest.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
import pytest
2+
3+
from .helpers import make_client
4+
5+
6+
@pytest.fixture
7+
def client():
8+
yield make_client()
9+
10+
11+
@pytest.fixture
12+
def PROJECT():
13+
yield "PROJECT"
14+
15+
16+
@pytest.fixture
17+
def DS_ID():
18+
yield "DATASET_ID"
19+
20+
21+
@pytest.fixture
22+
def LOCATION():
23+
yield "us-central"

‎tests/unit/helpers.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import google.cloud.bigquery.client
16+
import google.cloud.bigquery.dataset
17+
import mock
18+
import pytest
19+
1520

1621
def make_connection(*responses):
1722
import google.cloud.bigquery._http
@@ -31,3 +36,47 @@ def _to_pyarrow(value):
3136
import pyarrow
3237

3338
return pyarrow.array([value])[0]
39+
40+
41+
def make_client(project="PROJECT", **kw):
42+
credentials = mock.Mock(spec=google.auth.credentials.Credentials)
43+
return google.cloud.bigquery.client.Client(project, credentials, **kw)
44+
45+
46+
def make_dataset_reference_string(project, ds_id):
47+
return f"{project}.{ds_id}"
48+
49+
50+
def make_dataset(project, ds_id):
51+
return google.cloud.bigquery.dataset.Dataset(
52+
google.cloud.bigquery.dataset.DatasetReference(project, ds_id)
53+
)
54+
55+
56+
def make_dataset_list_item(project, ds_id):
57+
return google.cloud.bigquery.dataset.DatasetListItem(
58+
dict(datasetReference=dict(projectId=project, datasetId=ds_id))
59+
)
60+
61+
62+
def identity(x):
63+
return x
64+
65+
66+
def get_reference(x):
67+
return x.reference
68+
69+
70+
dataset_like = [
71+
(google.cloud.bigquery.dataset.DatasetReference, identity),
72+
(make_dataset, identity),
73+
(make_dataset_list_item, get_reference),
74+
(
75+
make_dataset_reference_string,
76+
google.cloud.bigquery.dataset.DatasetReference.from_string,
77+
),
78+
]
79+
80+
dataset_polymorphic = pytest.mark.parametrize(
81+
"make_dataset,get_reference", dataset_like
82+
)

0 commit comments

Comments
 (0)