test delegated access

Author: Matt Sokoloff

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+# Version 3.4.0 (2021-09-10)
+## Added
+* New `IAMIntegration` entity
+* `iam_integration` parameter to `Client.create_dataset()` to specify the iam integration
+    - `create_dataset` will use the default integration if this param is not set
+* `Organization.get_iam_integrations()` to list all integrations available to your org
+* `Organization.get_default_iam_integration()` to only get the defaault iam integration
+
 # Version 3.3.0 (2021-09-02)
 ## Added
 * `Dataset.create_data_rows_sync()` for synchronous bulk uploads of data rows

labelbox/__init__.py

@@ -1,5 +1,5 @@
 name = "labelbox"
-__version__ = "3.3.0"
+__version__ = "3.4.0"
 
 from labelbox.schema.project import Project
 from labelbox.client import Client

labelbox/client.py

@@ -1,7 +1,7 @@
 # type: ignore
 from datetime import datetime, timezone
 import json
-from labelbox.schema.iam_integration import IAMIntegration
+
 import logging
 import mimetypes
 import os
@@ -10,8 +10,9 @@
 import requests
 import requests.exceptions
 
-from labelbox import utils
 import labelbox.exceptions
+from labelbox import utils
+from labelbox import __version__ as SDK_VERSION
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
 from labelbox.pagination import PaginatedCollection
@@ -23,8 +24,8 @@
 from labelbox.schema.organization import Organization
 from labelbox.schema.data_row_metadata import DataRowMetadataOntology
 from labelbox.schema.labeling_frontend import LabelingFrontend
+from labelbox.schema.iam_integration import IAMIntegration
 from labelbox.schema import role
-from labelbox import __version__ as SDK_VERSION
 
 logger = logging.getLogger(__name__)
 
@@ -504,15 +505,18 @@ def _create(self, db_object_type, data):
         res = res["create%s" % db_object_type.type_name()]
         return db_object_type(self, res)
 
-    def create_dataset(self, **kwargs):
+    def create_dataset(self, iam_integration=IAMIntegration._DEFAULT, **kwargs):
         """ Creates a Dataset object on the server.
+        This will attempt to connect the organization's default IAM integration if it exists
 
         Attribute values are passed as keyword arguments.
 
         >>> project = client.get_project("<project_uid>")
         >>> dataset = client.create_dataset(name="<dataset_name>", projects=project)
 
         Args:
+            iam_integration (IAMIntegration) : Uses the default integration.
+                Optionally specify another integration or set as None to not use delegated access
             **kwargs: Keyword arguments with Dataset attribute values.
         Returns:
             A new Dataset object.
@@ -521,31 +525,41 @@ def create_dataset(self, **kwargs):
                 any of the attribute names given in kwargs.
         """
         dataset = self._create(Dataset, kwargs)
-        iam_integration = kwargs.get('iam_integration') or self.get_organization().get_default_iam_integration()
-        if iam_integration is not None:
-            if not isinstance(iam_integration, IAMIntegration):
-                raise TypeError(f"iam integration must be a reference an `IAMIntegration` object. Found {type(iam_integration)}")
-
-            if not iam_integration.valid:
-                raise ValueError("Invalid integration is invalid. Please select another integration or remove default.")
-            try:
-                self.execute("""
-                    mutation setSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) {
-                        setSignerForDataset(data: { signerId: $signerId}, where: {id: $datasetId}){id}}
-                    """, {'signerId' : iam_integration.uid, 'datasetId' : dataset.uid})
-                validation_result = self.execute("""
-                    mutation validateDatasetPyApi($id: ID!){validateDataset(where: {id : $id}){
-                        valid checks{name, success}}}
-                    """, {'id' : dataset.uid})
-                if not validation_result['validateDataset']['checks'][0]['success']:
-                    raise labelbox.exceptions.LabelboxError(
-                        f"IAMIntegration {validation_result['validateDataset']['checks']['name']} was not successfully added added to the project."
-                    )
-            except Exception as e:
-                dataset.delete()
-                raise e
-        return dataset
 
+        if iam_integration == IAMIntegration._DEFAULT:
+            iam_integration = self.get_organization(
+            ).get_default_iam_integration()
+
+        if iam_integration is None:
+            return dataset
+
+        if not isinstance(iam_integration, IAMIntegration):
+            raise TypeError(
+                f"iam integration must be a reference an `IAMIntegration` object. Found {type(iam_integration)}"
+            )
+
+        if not iam_integration.valid:
+            raise ValueError("Integration is not valid. Please select another.")
+        try:
+            self.execute(
+                """mutation setSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) {
+                    setSignerForDataset(data: { signerId: $signerId}, where: {id: $datasetId}){id}}
+                """, {
+                    'signerId': iam_integration.uid,
+                    'datasetId': dataset.uid
+                })
+            validation_result = self.execute(
+                """mutation validateDatasetPyApi($id: ID!){validateDataset(where: {id : $id}){
+                    valid checks{name, success}}}
+                """, {'id': dataset.uid})
+            if not validation_result['validateDataset']['checks'][0]['success']:
+                raise labelbox.exceptions.LabelboxError(
+                    f"IAMIntegration {validation_result['validateDataset']['checks']['name']} was not successfully added added to the project."
+                )
+        except Exception as e:
+            dataset.delete()
+            raise e
+        return dataset
 
     def create_project(self, **kwargs):
         """ Creates a Project object on the server.
@@ -643,5 +657,3 @@ def create_model(self, name, ontology_id):
             "ontologyId": ontology_id
         })
         return Model(self, result['createModel'])
-
-

labelbox/schema/dataset.py

@@ -44,7 +44,8 @@ class Dataset(DbObject, Updateable, Deletable):
     data_rows = Relationship.ToMany("DataRow", False)
     created_by = Relationship.ToOne("User", False, "created_by")
     organization = Relationship.ToOne("Organization", False)
-    iam_integration = Relationship.ToOne("IAMIntegration", False, "iam_integration", "signer")
+    iam_integration = Relationship.ToOne("IAMIntegration", False,
+                                         "iam_integration", "signer")
 
     def create_data_row(self, **kwargs):
         """ Creates a single DataRow belonging to this dataset.

labelbox/schema/iam_integration.py

@@ -15,6 +15,7 @@ class IAMIntegration(DbObject):
         is_org_default (boolean)
 
     """
+    _DEFAULT = "DEFAULT"
 
     name = Field.String("name")
     created_at = Field.DateTime("created_at")

labelbox/schema/organization.py

@@ -130,20 +130,31 @@ def remove_user(self, user: User):
             updateUser(where: {id: $%s}, data: {deleted: true}) { id deleted }
         }""" % (user_id_param, user_id_param), {user_id_param: user.uid})
 
-
     def get_iam_integrations(self):
         """
         Returns all IAM Integrations for an organization
         """
-        res = self.client.execute("""query getAllIntegrationsPyApi { iamIntegrations {%s} } """ % query.results_query_part(Entity.IAMIntegration))
-        return [Entity.IAMIntegration(self.client, integration_data) for integration_data in res['iamIntegrations']]
+        res = self.client.execute(
+            """query getAllIntegrationsPyApi { iamIntegrations {%s} } """ %
+            query.results_query_part(Entity.IAMIntegration))
+        return [
+            Entity.IAMIntegration(self.client, integration_data)
+            for integration_data in res['iamIntegrations']
+        ]
 
     def get_default_iam_integration(self):
         """
-        Returns the default IAM integration for the organization
+        Returns the default IAM integration for the organization.
+        Will return None if there are no default integrations for the org.
         """
         integrations = self.get_iam_integrations()
-        default_integration = [integration for integration in integrations if integration.is_org_default]
+        default_integration = [
+            integration for integration in integrations
+            if integration.is_org_default
+        ]
         if len(default_integration) > 1:
-            raise ValueError("Found more than one default signer. Please contact Labelbox to resolve")
-        return None if not len(default_integration) else default_integration.pop()
+            raise ValueError(
+                "Found more than one default signer. Please contact Labelbox to resolve"
+            )
+        return None if not len(
+            default_integration) else default_integration.pop()

tests/integration/test_data_row_metadata.py

@@ -209,6 +209,7 @@ def test_delete_non_existent_schema_id(datarow, mdo):
 
 
 @pytest.mark.slow
+@pytest.mark.skip("Test is inconsistent.")
 def test_large_bulk_delete_non_existent_schema_id(big_dataset, mdo):
     deletes = []
     n_fields_start = 0

tests/integration/test_delegated_access.py

@@ -0,0 +1,43 @@
+import requests
+import pytest
+
+
+@pytest.mark.skip("Can only be tested in specific organizations.")
+def test_default_integration(client):
+    # This tests assumes the following:
+    # 1. gcp delegated access is configured to work with utkarsh-da-test-bucket
+    # 2. the integration name is gcp test
+    # 3. This integration is the default
+    ds = client.create_dataset(name="new_ds")
+    dr = ds.create_data_row(
+        row_data=
+        "gs://utkarsh-da-test-bucket/mathew-schwartz-8rj4sz9YLCI-unsplash.jpg")
+    assert requests.get(dr.row_data).status_code == 200
+    assert ds.iam_integration().name == "GCP Test"
+    ds.delete()
+
+
+@pytest.mark.skip("Can only be tested in specific organizations.")
+def test_non_default_integration(client):
+    # This tests assumes the following:
+    # 1. aws delegated access is configured to work with lbox-test-bucket
+    # 2. an integration called aws is available to the org
+    integrations = client.get_organization().get_iam_integrations()
+    integration = [inte for inte in integrations if 'aws' in inte.name][0]
+    assert integration.valid
+    ds = client.create_dataset(iam_integration=integration, name="new_ds")
+    assert ds.iam_integration().name == "aws"
+    dr = ds.create_data_row(
+        row_data=
+        "https://lbox-test-bucket.s3.us-east-1.amazonaws.com/2021_09_08_0hz_Kleki.png"
+    )
+    assert requests.get(dr.row_data).status_code == 200
+    ds.delete()
+
+
+def test_no_integration(client, image_url):
+    ds = client.create_dataset(iam_integration=None, name="new_ds")
+    assert ds.iam_integration() is None
+    dr = ds.create_data_row(row_data=image_url)
+    assert requests.get(dr.row_data).status_code == 200
+    ds.delete()