bigquery: Dataset constructor takes a DatasetReference (#4036)

Author: jba

bigquery/google/cloud/bigquery/client.py

@@ -185,8 +185,6 @@ def create_dataset(self, dataset):
         :rtype: ":class:`~google.cloud.bigquery.dataset.Dataset`"
         :returns: a new ``Dataset`` returned from the service.
         """
-        if dataset.project is None:
-            dataset._project = self.project
         path = '/projects/%s/datasets' % (dataset.project,)
         api_response = self._connection.api_request(
             method='POST', path=path, data=dataset._build_resource())
@@ -244,8 +242,6 @@ def update_dataset(self, dataset, fields):
         :rtype: :class:`google.cloud.bigquery.dataset.Dataset`
         :returns: the modified ``Dataset`` instance
         """
-        if dataset.project is None:
-            dataset._project = self.project
         path = '/projects/%s/datasets/%s' % (dataset.project,
                                              dataset.dataset_id)
         partial = {}

bigquery/google/cloud/bigquery/dataset.py

@@ -106,6 +106,10 @@ class DatasetReference(object):
     """
 
     def __init__(self, project, dataset_id):
+        if not isinstance(project, six.string_types):
+            raise ValueError("Pass a string for project")
+        if not isinstance(dataset_id, six.string_types):
+            raise ValueError("Pass a string for dataset_id")
         self._project = project
         self._dataset_id = dataset_id
 
@@ -154,27 +158,15 @@ class Dataset(object):
     See
     https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets
 
-    :type dataset_id: str
-    :param dataset_id: the ID of the dataset
-
-    :type access_entries: list of :class:`AccessEntry`
-    :param access_entries: roles granted to entities for this dataset
-
-    :type project: str
-    :param project: (Optional) project ID for the dataset.
+    :type dataset_ref: :class:`~google.cloud.bigquery.dataset.DatasetReference`
+    :param dataset_ref: a pointer to a dataset
     """
 
-    _access_entries = None
-
-    def __init__(self,
-                 dataset_id,
-                 access_entries=(),
-                 project=None):
-        self._dataset_id = dataset_id
+    def __init__(self, dataset_ref):
+        self._project = dataset_ref.project
+        self._dataset_id = dataset_ref.dataset_id
         self._properties = {'labels': {}}
-        # Let the @property do validation.
-        self.access_entries = access_entries
-        self._project = project
+        self._access_entries = ()
 
     @property
     def project(self):
@@ -406,7 +398,7 @@ def from_api_repr(cls, resource):
             raise KeyError('Resource lacks required identity information:'
                            '["datasetReference"]["datasetId"]')
         dataset_id = dsr['datasetId']
-        dataset = cls(dataset_id, project=dsr['projectId'])
+        dataset = cls(DatasetReference(dsr['projectId'], dataset_id))
         dataset._set_properties(resource)
         return dataset
 

bigquery/google/cloud/bigquery/job.py

@@ -842,8 +842,9 @@ def from_api_repr(cls, resource, client):
         """
         job_id, config = cls._get_resource_config(resource)
         dest_config = config['destinationTable']
-        dataset = Dataset(dest_config['datasetId'],
-                          project=dest_config['projectId'])
+        ds_ref = DatasetReference(dest_config['projectId'],
+                                  dest_config['datasetId'],)
+        dataset = Dataset(ds_ref)
         table_ref = TableReference(dataset, dest_config['tableId'])
         destination = Table(table_ref, client=client)
         source_urls = config.get('sourceUris', ())
@@ -959,8 +960,9 @@ def from_api_repr(cls, resource, client):
         """
         job_id, config = cls._get_resource_config(resource)
         dest_config = config['destinationTable']
-        dataset = Dataset(dest_config['datasetId'],
-                          project=dest_config['projectId'])
+        ds_ref = DatasetReference(dest_config['projectId'],
+                                  dest_config['datasetId'],)
+        dataset = Dataset(ds_ref)
         table_ref = TableReference(dataset, dest_config['tableId'])
         destination = Table(table_ref, client=client)
         sources = []
@@ -972,9 +974,9 @@ def from_api_repr(cls, resource, client):
                     "Resource missing 'sourceTables' / 'sourceTable'")
             source_configs = [single]
         for source_config in source_configs:
-            dataset = Dataset(source_config['datasetId'],
-                              project=source_config['projectId'])
-            table_ref = TableReference(dataset, source_config['tableId'])
+            ds_ref = DatasetReference(source_config['projectId'],
+                                      source_config['datasetId'])
+            table_ref = ds_ref.table(source_config['tableId'])
             sources.append(Table(table_ref, client=client))
         job = cls(job_id, destination, sources, client=client)
         job._set_properties(resource)
@@ -1426,17 +1428,17 @@ def _copy_configuration_properties(self, configuration):
             dest_local = self._destination_table_resource()
             if dest_remote != dest_local:
                 project = dest_remote['projectId']
-                dataset = Dataset(dest_remote['datasetId'], project=project)
+                dataset = Dataset(DatasetReference(project,
+                                                   dest_remote['datasetId']))
                 self.destination = dataset.table(dest_remote['tableId'])
 
         def_ds = configuration.get('defaultDataset')
         if def_ds is None:
             if self.default_dataset is not None:
                 del self.default_dataset
         else:
-            self.default_dataset = Dataset(def_ds['datasetId'],
-                                           project=def_ds['projectId'])
-
+            self.default_dataset = Dataset(
+                DatasetReference(def_ds['projectId'], def_ds['datasetId']))
         udf_resources = []
         for udf_mapping in configuration.get(self._UDF_KEY, ()):
             key_val, = udf_mapping.items()
@@ -1587,11 +1589,11 @@ def referenced_tables(self):
 
             t_project = table['projectId']
 
-            ds_name = table['datasetId']
-            t_dataset = datasets_by_project_name.get((t_project, ds_name))
+            ds_id = table['datasetId']
+            t_dataset = datasets_by_project_name.get((t_project, ds_id))
             if t_dataset is None:
-                t_dataset = Dataset(ds_name, project=t_project)
-                datasets_by_project_name[(t_project, ds_name)] = t_dataset
+                t_dataset = DatasetReference(t_project, ds_id)
+                datasets_by_project_name[(t_project, ds_id)] = t_dataset
 
             t_name = table['tableId']
             tables.append(t_dataset.table(t_name))

bigquery/tests/system.py

@@ -112,8 +112,7 @@ def _still_in_use(bad_request):
 
     def test_create_dataset(self):
         DATASET_ID = _make_dataset_id('create_dataset')
-        dataset = retry_403(Config.CLIENT.create_dataset)(Dataset(DATASET_ID))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(DATASET_ID)
 
         self.assertTrue(_dataset_exists(dataset))
         self.assertEqual(dataset.dataset_id, DATASET_ID)
@@ -122,7 +121,7 @@ def test_create_dataset(self):
     def test_get_dataset(self):
         DATASET_ID = _make_dataset_id('get_dataset')
         client = Config.CLIENT
-        dataset_arg = Dataset(DATASET_ID, project=client.project)
+        dataset_arg = Dataset(client.dataset(DATASET_ID))
         dataset_arg.friendly_name = 'Friendly'
         dataset_arg.description = 'Description'
         dataset = retry_403(client.create_dataset)(dataset_arg)
@@ -135,10 +134,7 @@ def test_get_dataset(self):
         self.assertEqual(got.description, 'Description')
 
     def test_update_dataset(self):
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('update_dataset')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('update_dataset'))
         self.assertTrue(_dataset_exists(dataset))
         self.assertIsNone(dataset.friendly_name)
         self.assertIsNone(dataset.description)
@@ -163,16 +159,15 @@ def test_update_dataset(self):
         self.assertEqual(ds3.labels, {'color': 'green', 'shape': 'circle'})
 
     # TODO(jba): test that read-modify-write with ETag works.
+
     def test_list_datasets(self):
         datasets_to_create = [
             'new' + unique_resource_id(),
             'newer' + unique_resource_id(),
             'newest' + unique_resource_id(),
         ]
         for dataset_id in datasets_to_create:
-            created_dataset = retry_403(Config.CLIENT.create_dataset)(
-                Dataset(dataset_id))
-            self.to_delete.append(created_dataset)
+            self.temp_dataset(dataset_id)
 
         # Retrieve the datasets.
         iterator = Config.CLIENT.list_datasets()
@@ -184,9 +179,7 @@ def test_list_datasets(self):
         self.assertEqual(len(created), len(datasets_to_create))
 
     def test_create_table(self):
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-                Dataset(_make_dataset_id('create_table')))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(_make_dataset_id('create_table'))
 
         TABLE_NAME = 'test_table'
         full_name = bigquery.SchemaField('full_name', 'STRING',
@@ -217,9 +210,7 @@ def test_get_table_w_public_dataset(self):
 
     def test_list_dataset_tables(self):
         DATASET_ID = _make_dataset_id('list_tables')
-        dataset = retry_403(Config.CLIENT.create_dataset)(Dataset(DATASET_ID))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(DATASET_ID)
         # Retrieve tables before any are created for the dataset.
         iterator = Config.CLIENT.list_dataset_tables(dataset)
         all_tables = list(iterator)
@@ -252,9 +243,7 @@ def test_list_dataset_tables(self):
         self.assertEqual(len(created), len(tables_to_create))
 
     def test_patch_table(self):
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('patch_table')))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(_make_dataset_id('patch_table'))
 
         TABLE_NAME = 'test_table'
         full_name = bigquery.SchemaField('full_name', 'STRING',
@@ -273,9 +262,7 @@ def test_patch_table(self):
         self.assertEqual(table.description, 'Description')
 
     def test_update_table(self):
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('update_table')))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(_make_dataset_id('update_table'))
 
         TABLE_NAME = 'test_table'
         full_name = bigquery.SchemaField('full_name', 'STRING',
@@ -316,10 +303,7 @@ def test_insert_data_then_dump_table(self):
         ]
         ROW_IDS = range(len(ROWS))
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('insert_data_then_dump')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('insert_data_then_dump'))
         TABLE_NAME = 'test_table'
         full_name = bigquery.SchemaField('full_name', 'STRING',
                                          mode='REQUIRED')
@@ -358,10 +342,7 @@ def test_load_table_from_local_file_then_dump_table(self):
         ]
         TABLE_NAME = 'test_table'
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('load_local_then_dump')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('load_local_then_dump'))
         full_name = bigquery.SchemaField('full_name', 'STRING',
                                          mode='REQUIRED')
         age = bigquery.SchemaField('age', 'INTEGER', mode='REQUIRED')
@@ -406,10 +387,7 @@ def test_load_table_from_local_avro_file_then_dump_table(self):
             ("orange", 590),
             ("red", 650)]
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('load_local_then_dump')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('load_local_then_dump'))
         table = Table(dataset.table(TABLE_NAME), client=Config.CLIENT)
         self.to_delete.insert(0, table)
 
@@ -467,9 +445,7 @@ def test_load_table_from_storage_then_dump_table(self):
 
         self.to_delete.insert(0, blob)
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('load_gcs_then_dump')))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(_make_dataset_id('load_gcs_then_dump'))
 
         full_name = bigquery.SchemaField('full_name', 'STRING',
                                          mode='REQUIRED')
@@ -536,10 +512,7 @@ def test_load_table_from_storage_w_autodetect_schema(self):
 
         self.to_delete.insert(0, blob)
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('load_gcs_then_dump')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('load_gcs_then_dump'))
         table_ref = dataset.table(table_name)
 
         job = Config.CLIENT.load_table_from_storage(
@@ -589,9 +562,7 @@ def _load_table_for_extract_table(
                 blob.upload_from_file(csv_read, content_type='text/csv')
         self.to_delete.insert(0, blob)
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(table.dataset_id))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(table.dataset_id)
         table_ref = dataset.table(table.table_id)
         job = Config.CLIENT.load_table_from_storage(
             'bq_extract_storage_test_' + local_id, table_ref, gs_url)
@@ -676,8 +647,7 @@ def test_job_cancel(self):
         TABLE_NAME = 'test_table'
         QUERY = 'SELECT * FROM %s.%s' % (DATASET_ID, TABLE_NAME)
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(Dataset(DATASET_ID))
-        self.to_delete.append(dataset)
+        dataset = self.temp_dataset(DATASET_ID)
 
         full_name = bigquery.SchemaField('full_name', 'STRING',
                                          mode='REQUIRED')
@@ -866,9 +836,7 @@ def test_dbapi_fetchall(self):
     def _load_table_for_dml(self, rows, dataset_id, table_id):
         from google.cloud._testing import _NamedTemporaryFile
 
-        dataset = retry_403(Config.CLIENT.create_dataset)(Dataset(dataset_id))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(dataset_id)
         greeting = bigquery.SchemaField(
             'greeting', 'STRING', mode='NULLABLE')
         table = Table(dataset.table(table_id), schema=[greeting],
@@ -1190,8 +1158,7 @@ def test_dump_table_w_public_data(self):
         DATASET_ID = 'samples'
         TABLE_NAME = 'natality'
 
-        dataset = Dataset(DATASET_ID, project=PUBLIC)
-        table_ref = dataset.table(TABLE_NAME)
+        table_ref = DatasetReference(PUBLIC, DATASET_ID).table(TABLE_NAME)
         table = Config.CLIENT.get_table(table_ref)
         self._fetch_single_page(table)
 
@@ -1242,10 +1209,7 @@ def test_insert_nested_nested(self):
             ('Some value', record)
         ]
         table_name = 'test_table'
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('issue_2951')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(_make_dataset_id('issue_2951'))
         table = Table(dataset.table(table_name), schema=schema,
                       client=Config.CLIENT)
         table.create()
@@ -1260,10 +1224,8 @@ def test_insert_nested_nested(self):
 
     def test_create_table_insert_fetch_nested_schema(self):
         table_name = 'test_table'
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('create_table_nested_schema')))
-        self.to_delete.append(dataset)
-
+        dataset = self.temp_dataset(
+            _make_dataset_id('create_table_nested_schema'))
         schema = _load_json_schema()
         table = Table(dataset.table(table_name), schema=schema,
                       client=Config.CLIENT)
@@ -1321,6 +1283,12 @@ def test_create_table_insert_fetch_nested_schema(self):
             e_favtime = datetime.datetime(*parts[0:6])
             self.assertEqual(found[7], e_favtime)              # FavoriteTime
 
+    def temp_dataset(self, dataset_id):
+        dataset = retry_403(Config.CLIENT.create_dataset)(
+            Dataset(Config.CLIENT.dataset(dataset_id)))
+        self.to_delete.append(dataset)
+        return dataset
+
 
 def _job_done(instance):
     return instance.state.lower() == 'done'