rbramwell
diff --git a/‎healthcare/api-client/datasets/datasets.py
Copy file name to clipboardExpand all lines: healthcare/api-client/datasets/datasets.py
+106-2Lines changed: 106 additions & 2 deletions b/‎healthcare/api-client/datasets/datasets.py
Copy file name to clipboardExpand all lines: healthcare/api-client/datasets/datasets.py
+106-2Lines changed: 106 additions & 2 deletions
diff --git a/‎healthcare/api-client/datasets/datasets_test.py
Copy file name to clipboardExpand all lines: healthcare/api-client/datasets/datasets_test.py
+41-4Lines changed: 41 additions & 4 deletions b/‎healthcare/api-client/datasets/datasets_test.py
Copy file name to clipboardExpand all lines: healthcare/api-client/datasets/datasets_test.py
+41-4Lines changed: 41 additions & 4 deletions
@@ -25,15 +25,15 @@ def get_client(service_account_json, api_key):
     """Returns an authorized API client by discovering the Healthcare API and
     creating a service object using the service account credentials JSON."""
     api_scopes = ['https://www.googleapis.com/auth/cloud-platform']
-    api_version = 'v1alpha2'
+    api_version = 'v1beta1'
     discovery_api = 'https://healthcare.googleapis.com/$discovery/rest'
     service_name = 'healthcare'
 
     credentials = service_account.Credentials.from_service_account_file(
         service_account_json)
     scoped_credentials = credentials.with_scopes(api_scopes)
 
-    discovery_url = '{}?labels=CHC_ALPHA&version={}&key={}'.format(
+    discovery_url = '{}?labels=CHC_BETA&version={}&key={}'.format(
         discovery_api, api_version, api_key)
 
     return discovery.build(
@@ -237,6 +237,80 @@ def deidentify_dataset(
 # [END healthcare_deidentify_dataset]
 
 
+# [START healthcare_dataset_get_iam_policy]
+def get_dataset_iam_policy(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id):
+    """Gets the IAM policy for the specified dataset."""
+    client = get_client(service_account_json, api_key)
+    dataset_name = 'projects/{}/locations/{}/datasets/{}'.format(
+        project_id, cloud_region, dataset_id)
+
+    request = client.projects().locations().datasets().getIamPolicy(
+        resource=dataset_name)
+    response = request.execute()
+
+    print('etag: {}'.format(response.get('name')))
+    return response
+# [END healthcare_dataset_get_iam_policy]
+
+
+# [START healthcare_dataset_set_iam_policy]
+def set_dataset_iam_policy(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id,
+        member,
+        role,
+        etag=None):
+    """Sets the IAM policy for the specified dataset.
+
+        A single member will be assigned a single role. A member can be any of:
+
+        - allUsers, that is, anyone
+        - allAuthenticatedUsers, anyone authenticated with a Google account
+        - user:email, as in 'user:somebody@example.com'
+        - group:email, as in 'group:admins@example.com'
+        - domain:domainname, as in 'domain:example.com'
+        - serviceAccount:email,
+            as in 'serviceAccount:my-other-app@appspot.gserviceaccount.com'
+
+        A role can be any IAM role, such as 'roles/viewer', 'roles/owner',
+        or 'roles/editor'
+    """
+    client = get_client(service_account_json, api_key)
+    dataset_name = 'projects/{}/locations/{}/datasets/{}'.format(
+        project_id, cloud_region, dataset_id)
+
+    policy = {
+        "bindings": [
+            {
+              "role": role,
+              "members": [
+                member
+              ]
+            }
+        ]
+    }
+
+    if etag is not None:
+        policy['etag'] = etag
+
+    request = client.projects().locations().datasets().setIamPolicy(
+        resource=dataset_name, body={'policy': policy})
+    response = request.execute()
+
+    print('etag: {}'.format(response.get('name')))
+    print('bindings: {}'.format(response.get('bindings')))
+    return response
+# [END healthcare_dataset_set_iam_policy]
+
+
 def parse_command_line_args():
     """Parses command line arguments."""
 
@@ -286,13 +360,25 @@ def parse_command_line_args():
         help='The data to keeplist, for example "PatientID" '
         'or "StudyInstanceUID"')
 
+    parser.add_argument(
+        '--member',
+        default=None,
+        help='Member to add to IAM policy (e.g. "domain:example.com")')
+
+    parser.add_argument(
+        '--role',
+        default=None,
+        help='IAM Role to give to member (e.g. "roles/viewer")')
+
     command = parser.add_subparsers(dest='command')
 
     command.add_parser('create-dataset', help=create_dataset.__doc__)
     command.add_parser('delete-dataset', help=delete_dataset.__doc__)
     command.add_parser('get-dataset', help=get_dataset.__doc__)
     command.add_parser('list-datasets', help=list_datasets.__doc__)
     command.add_parser('patch-dataset', help=patch_dataset.__doc__)
+    command.add_parser('get_iam_policy', help=get_dataset_iam_policy.__doc__)
+    command.add_parser('set_iam_policy', help=set_dataset_iam_policy.__doc__)
 
     command.add_parser('deidentify-dataset', help=deidentify_dataset.__doc__)
 
@@ -356,6 +442,24 @@ def run_command(args):
             args.destination_dataset_id,
             args.keeplist_tags)
 
+    elif args.command == 'get_iam_policy':
+        get_dataset_iam_policy(
+            args.service_account_json,
+            args.api_key,
+            args.project_id,
+            args.cloud_region,
+            args.dataset_id)
+
+    elif args.command == 'set_iam_policy':
+        set_dataset_iam_policy(
+            args.service_account_json,
+            args.api_key,
+            args.project_id,
+            args.cloud_region,
+            args.dataset_id,
+            args.member,
+            args.role)
+
 
 def main():
     args = parse_command_line_args()
 
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import os
-import pytest
 import time
 
 import datasets
@@ -29,7 +28,6 @@
 time_zone = 'UTC'
 
 
-@pytest.mark.skip(reason='disable until API whitelisted / enabled')
 def test_CRUD_dataset(capsys):
     datasets.create_dataset(
         service_account_json,
@@ -57,7 +55,6 @@ def test_CRUD_dataset(capsys):
     assert 'Deleted dataset' in out
 
 
-@pytest.mark.skip(reason='disable until API whitelisted / enabled')
 def test_patch_dataset(capsys):
     datasets.create_dataset(
         service_account_json,
@@ -84,7 +81,6 @@ def test_patch_dataset(capsys):
     assert 'UTC' in out
 
 
-@pytest.mark.skip(reason='disable until API whitelisted / enabled')
 def test_deidentify_dataset(capsys):
     datasets.create_dataset(
         service_account_json,
@@ -116,3 +112,44 @@ def test_deidentify_dataset(capsys):
 
     # Check that de-identify worked
     assert 'De-identified data written to' in out
+
+
+def test_get_set_dataset_iam_policy(capsys):
+    datasets.create_dataset(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id)
+
+    get_response = datasets.get_dataset_iam_policy(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id)
+
+    set_response = datasets.set_dataset_iam_policy(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id,
+        'serviceAccount:python-docs-samples-tests@appspot.gserviceaccount.com',
+        'roles/viewer')
+
+    # Clean up
+    datasets.delete_dataset(
+        service_account_json,
+        api_key,
+        project_id,
+        cloud_region,
+        dataset_id)
+
+    out, _ = capsys.readouterr()
+
+    assert 'etag' in get_response
+    assert 'bindings' in set_response
+    assert len(set_response['bindings']) == 1
+    assert 'python-docs-samples-tests' in str(set_response['bindings'])
+    assert 'roles/viewer' in str(set_response['bindings'])