From c098cd01c755633bfaba7193dd5c044a489a5b61 Mon Sep 17 00:00:00 2001 From: Steffany Brown <30247553+steffnay@users.noreply.github.com> Date: Tue, 14 Dec 2021 08:11:21 -0800 Subject: [PATCH] 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 --- google/cloud/bigquery/dataset.py | 73 ++++++++++---------- google/cloud/bigquery/enums.py | 13 ++++ samples/snippets/authorized_view_tutorial.py | 5 +- samples/snippets/update_dataset_access.py | 4 +- tests/unit/test_dataset.py | 40 ++++++++++- 5 files changed, 92 insertions(+), 43 deletions(-) diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index ff015d605..499072de2 100644 --- a/google/cloud/bigquery/dataset.py +++ b/google/cloud/bigquery/dataset.py @@ -77,10 +77,10 @@ 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. @@ -88,17 +88,18 @@ class AccessEntry(object): 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): diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py index 0eaaffd2e..7fc0a5fd6 100644 --- a/google/cloud/bigquery/enums.py +++ b/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): diff --git a/samples/snippets/authorized_view_tutorial.py b/samples/snippets/authorized_view_tutorial.py index b6a20c6ec..66810c036 100644 --- a/samples/snippets/authorized_view_tutorial.py +++ b/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( diff --git a/samples/snippets/update_dataset_access.py b/samples/snippets/update_dataset_access.py index fb3bfa14f..1448213a6 100644 --- a/samples/snippets/update_dataset_access.py +++ b/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: diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py index b3a53a08d..c554782bf 100644 --- a/tests/unit/test_dataset.py +++ b/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 = {