feat: support authorized dataset entity (#1075)

* feat: support authorized dataset entity

* cleanup

* add test and cache the resource from from_api_repr in a _properties value

* lint

* update samples to use enums

* update to_api_repr and add tests

* refactor

Author: steffnay

packages/google-cloud-bigquery/google/cloud/bigquery/dataset.py

@@ -77,28 +77,29 @@ def _get_routine_reference(self, routine_id):
 class AccessEntry(object):
     """Represents grant of an access role to an entity.
 
-    An entry must have exactly one of the allowed :attr:`ENTITY_TYPES`. If
-    anything but ``view`` or ``routine`` are set, a ``role`` is also required.
-    ``role`` is omitted for ``view`` and ``routine``, because they are always
-    read-only.
+    An entry must have exactly one of the allowed
+    :class:`google.cloud.bigquery.enums.EntityTypes`. If anything but ``view``, ``routine``,
+    or ``dataset`` are set, a ``role`` is also required. ``role`` is omitted for ``view``,
+    ``routine``, ``dataset``, because they are always read-only.
 
     See https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets.
 
     Args:
         role (str):
             Role granted to the entity. The following string values are
             supported: `'READER'`, `'WRITER'`, `'OWNER'`. It may also be
-            :data:`None` if the ``entity_type`` is ``view`` or ``routine``.
+            :data:`None` if the ``entity_type`` is ``view``, ``routine``, or ``dataset``.
 
         entity_type (str):
-            Type of entity being granted the role. One of :attr:`ENTITY_TYPES`.
+            Type of entity being granted the role. See
+            :class:`google.cloud.bigquery.enums.EntityTypes` for supported types.
 
         entity_id (Union[str, Dict[str, str]]):
-            If the ``entity_type`` is not 'view' or 'routine', the ``entity_id``
-            is the ``str`` ID of the entity being granted the role. If the
-            ``entity_type`` is 'view' or 'routine', the ``entity_id`` is a ``dict``
-            representing the view  or routine from a different dataset to grant
-            access to in the following format for views::
+            If the ``entity_type`` is not 'view', 'routine', or 'dataset', the
+            ``entity_id`` is the ``str`` ID of the entity being granted the role. If
+            the ``entity_type`` is 'view' or 'routine', the ``entity_id`` is a ``dict``
+            representing the view or routine from a different dataset to grant access
+            to in the following format for views::
 
                 {
                     'projectId': string,
@@ -114,11 +115,22 @@ class AccessEntry(object):
                     'routineId': string
                 }
 
+            If the ``entity_type`` is 'dataset', the ``entity_id`` is a ``dict`` that includes
+            a 'dataset' field with a ``dict`` representing the dataset and a 'target_types'
+            field with a ``str`` value of the dataset's resource type::
+
+                {
+                    'dataset': {
+                        'projectId': string,
+                        'datasetId': string,
+                    },
+                    'target_types: 'VIEWS'
+                }
+
     Raises:
         ValueError:
-            If the ``entity_type`` is not among :attr:`ENTITY_TYPES`, or if a
-            ``view`` or a ``routine`` has ``role`` set, or a non ``view`` and
-            non ``routine`` **does not** have a ``role`` set.
+            If a ``view``, ``routine``, or ``dataset`` has ``role`` set, or a non ``view``,
+            non ``routine``, and non ``dataset`` **does not** have a ``role`` set.
 
     Examples:
         >>> entry = AccessEntry('OWNER', 'userByEmail', 'user@example.com')
@@ -131,27 +143,9 @@ class AccessEntry(object):
         >>> entry = AccessEntry(None, 'view', view)
     """
 
-    ENTITY_TYPES = frozenset(
-        [
-            "userByEmail",
-            "groupByEmail",
-            "domain",
-            "specialGroup",
-            "view",
-            "iamMember",
-            "routine",
-        ]
-    )
-    """Allowed entity types."""
-
-    def __init__(self, role, entity_type, entity_id):
-        if entity_type not in self.ENTITY_TYPES:
-            message = "Entity type %r not among: %s" % (
-                entity_type,
-                ", ".join(self.ENTITY_TYPES),
-            )
-            raise ValueError(message)
-        if entity_type in ("view", "routine"):
+    def __init__(self, role=None, entity_type=None, entity_id=None):
+        self._properties = {}
+        if entity_type in ("view", "routine", "dataset"):
             if role is not None:
                 raise ValueError(
                     "Role must be None for a %r. Received "
@@ -162,7 +156,6 @@ def __init__(self, role, entity_type, entity_id):
                 raise ValueError(
                     "Role must be set for entity " "type %r" % (entity_type,)
                 )
-
         self._role = role
         self._entity_type = entity_type
         self._entity_id = entity_id
@@ -214,7 +207,8 @@ def to_api_repr(self):
         Returns:
             Dict[str, object]: Access entry represented as an API resource
         """
-        resource = {self._entity_type: self._entity_id}
+        resource = copy.deepcopy(self._properties)
+        resource[self._entity_type] = self._entity_id
         if self._role is not None:
             resource["role"] = self._role
         return resource
@@ -241,7 +235,10 @@ def from_api_repr(cls, resource: dict) -> "AccessEntry":
         entity_type, entity_id = entry.popitem()
         if len(entry) != 0:
             raise ValueError("Entry has unexpected keys remaining.", entry)
-        return cls(role, entity_type, entity_id)
+
+        config = cls(role, entity_type, entity_id)
+        config._properties = copy.deepcopy(resource)
+        return config
 
 
 class DatasetReference(object):

packages/google-cloud-bigquery/google/cloud/bigquery/enums.py

@@ -232,6 +232,19 @@ def _make_sql_scalars_enum():
 StandardSqlDataTypes = _make_sql_scalars_enum()
 
 
+class EntityTypes(str, enum.Enum):
+    """Enum of allowed entity type names in AccessEntry"""
+
+    USER_BY_EMAIL = "userByEmail"
+    GROUP_BY_EMAIL = "groupByEmail"
+    DOMAIN = "domain"
+    DATASET = "dataset"
+    SPECIAL_GROUP = "specialGroup"
+    VIEW = "view"
+    IAM_MEMBER = "iamMember"
+    ROUTINE = "routine"
+
+
 # See also: https://cloud.google.com/bigquery/data-types#legacy_sql_data_types
 # and https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types
 class SqlTypeNames(str, enum.Enum):

packages/google-cloud-bigquery/samples/snippets/authorized_view_tutorial.py

@@ -24,6 +24,7 @@ def run_authorized_view_tutorial(override_values={}):
     # Create a source dataset
     # [START bigquery_avt_create_source_dataset]
     from google.cloud import bigquery
+    from google.cloud.bigquery.enums import EntityTypes
 
     client = bigquery.Client()
     source_dataset_id = "github_source_data"
@@ -106,7 +107,7 @@ def run_authorized_view_tutorial(override_values={}):
     # analyst_group_email = 'data_analysts@example.com'
     access_entries = shared_dataset.access_entries
     access_entries.append(
-        bigquery.AccessEntry("READER", "groupByEmail", analyst_group_email)
+        bigquery.AccessEntry("READER", EntityTypes.GROUP_BY_EMAIL, analyst_group_email)
     )
     shared_dataset.access_entries = access_entries
     shared_dataset = client.update_dataset(
@@ -118,7 +119,7 @@ def run_authorized_view_tutorial(override_values={}):
     # [START bigquery_avt_source_dataset_access]
     access_entries = source_dataset.access_entries
     access_entries.append(
-        bigquery.AccessEntry(None, "view", view.reference.to_api_repr())
+        bigquery.AccessEntry(None, EntityTypes.VIEW, view.reference.to_api_repr())
     )
     source_dataset.access_entries = access_entries
     source_dataset = client.update_dataset(

packages/google-cloud-bigquery/samples/snippets/update_dataset_access.py

@@ -27,6 +27,8 @@ def update_dataset_access(dataset_id: str, entity_id: str):
     # of the entity, such as a view's table reference.
     entity_id = "user-or-group-to-add@example.com"
 
+    from google.cloud.bigquery.enums import EntityTypes
+
     # TODO(developer): Set entity_type to the type of entity you are granting access to.
     # Common types include:
     #
@@ -37,7 +39,7 @@ def update_dataset_access(dataset_id: str, entity_id: str):
     #
     # For a complete reference, see the REST API reference documentation:
     # https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#Dataset.FIELDS.access
-    entity_type = "groupByEmail"
+    entity_type = EntityTypes.GROUP_BY_EMAIL
 
     # TODO(developer): Set role to a one of the "Basic roles for datasets"
     # described here:

packages/google-cloud-bigquery/tests/unit/test_dataset.py

@@ -141,6 +141,28 @@ def test_to_api_repr_routine(self):
         exp_resource = {"routine": routine}
         self.assertEqual(resource, exp_resource)
 
+    def test_to_api_repr_dataset(self):
+        dataset = {
+            "dataset": {"projectId": "my-project", "datasetId": "my_dataset"},
+            "target_types": "VIEWS",
+        }
+        entry = self._make_one(None, "dataset", dataset)
+        resource = entry.to_api_repr()
+        exp_resource = {"dataset": dataset}
+        self.assertEqual(resource, exp_resource)
+
+    def test_to_api_w_incorrect_role(self):
+        dataset = {
+            "dataset": {
+                "projectId": "my-project",
+                "datasetId": "my_dataset",
+                "tableId": "my_table",
+            },
+            "target_type": "VIEW",
+        }
+        with self.assertRaises(ValueError):
+            self._make_one("READER", "dataset", dataset)
+
     def test_from_api_repr(self):
         resource = {"role": "OWNER", "userByEmail": "salmon@example.com"}
         entry = self._get_target_class().from_api_repr(resource)
@@ -150,8 +172,22 @@ def test_from_api_repr(self):
 
     def test_from_api_repr_w_unknown_entity_type(self):
         resource = {"role": "READER", "unknown": "UNKNOWN"}
-        with self.assertRaises(ValueError):
-            self._get_target_class().from_api_repr(resource)
+        entry = self._get_target_class().from_api_repr(resource)
+        self.assertEqual(entry.role, "READER")
+        self.assertEqual(entry.entity_type, "unknown")
+        self.assertEqual(entry.entity_id, "UNKNOWN")
+        exp_resource = entry.to_api_repr()
+        self.assertEqual(resource, exp_resource)
+
+    def test_to_api_repr_w_extra_properties(self):
+        resource = {
+            "role": "READER",
+            "userByEmail": "salmon@example.com",
+        }
+        entry = self._get_target_class().from_api_repr(resource)
+        entry._properties["specialGroup"] = resource["specialGroup"] = "projectReaders"
+        exp_resource = entry.to_api_repr()
+        self.assertEqual(resource, exp_resource)
 
     def test_from_api_repr_entries_w_extra_keys(self):
         resource = {