needs to be tested

Author: Matt Sokoloff

labelbox/client.py

@@ -1,6 +1,7 @@
 # type: ignore
 from datetime import datetime, timezone
 import json
+from labelbox.schema.iam_integration import IAMIntegration
 import logging
 import mimetypes
 import os
@@ -519,7 +520,27 @@ def create_dataset(self, **kwargs):
             InvalidAttributeError: If the Dataset type does not contain
                 any of the attribute names given in kwargs.
         """
-        return self._create(Dataset, kwargs)
+        dataset = self._create(Dataset, kwargs)
+        #iam_integration = kwargs.get("iam_integration")
+        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 {iam_integration}")
+            try:
+                import time
+                time.sleep(5)
+                print(iam_integration.uid)
+                res = self.execute("""mutation setSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) { setSignerForDataset(data: { signerId: $signerId}, where: {id: $datasetId}){id}} """, {'signerId' : iam_integration.uid, 'datasetId' : dataset.uid})
+                #result = self.execute("""mutation validateDatasetPyApi($id: ID!){validateDataset(where: {id : $id}){valid checks{name, success}} }""", {'id' : dataset.uid})
+                # TODO: I am not sure what to do with this check...
+                #if not result['validateDataset']['valid']:
+                #    raise labelbox.exceptions.LabelboxError(f"Signer was unsuccessfully added to the dataset. {result}")
+            except Exception as e:
+                breakpoint()
+                dataset.delete()
+                raise e
+        return dataset
+
 
     def create_project(self, **kwargs):
         """ Creates a Project object on the server.
@@ -617,3 +638,5 @@ def create_model(self, name, ontology_id):
             "ontologyId": ontology_id
         })
         return Model(self, result['createModel'])
+
+

labelbox/schema/__init__.py

@@ -18,3 +18,4 @@
 import labelbox.schema.user
 import labelbox.schema.webhook
 import labelbox.schema.data_row_metadata
+import labelbox.schema.iam_integration

labelbox/schema/dataset.py

@@ -1,3 +1,4 @@
+from labelbox.schema import iam_integration
 from labelbox import utils
 import os
 import json
@@ -43,6 +44,7 @@ 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")
 
     def create_data_row(self, **kwargs):
         """ Creates a single DataRow belonging to this dataset.

labelbox/schema/organization.py

@@ -2,7 +2,7 @@
 
 from labelbox.exceptions import LabelboxError
 from labelbox import utils
-from labelbox.orm.db_object import DbObject, experimental, query
+from labelbox.orm.db_object import DbObject, experimental, query, Entity
 from labelbox.orm.model import Field, Relationship
 from labelbox.schema.invite import Invite, InviteLimit, ProjectRole
 from labelbox.schema.user import User
@@ -129,3 +129,21 @@ def remove_user(self, user: User):
             """mutation DeleteMemberPyApi($%s: ID!) {
             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']]
+
+    def get_default_iam_integration(self):
+        """
+        Returns the default IAM integration for the organization
+        """
+        integrations = self.get_iam_integrations()
+        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()