Add dataset.set_iam_integration() to select/deselect integrations

Author: paulnoirel

libs/labelbox/src/labelbox/schema/dataset.py

@@ -31,6 +31,7 @@
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.task import Task
 from labelbox.schema.user import User
+from labelbox.schema.iam_integration import IAMIntegration
 
 logger = logging.getLogger(__name__)
 
@@ -912,3 +913,80 @@ def _convert_items_to_upsert_format(self, _items):
             }  # remove None values
             _upsert_items.append(DataRowUpsertItem(payload=item, id=key))
         return _upsert_items
+    
+    def set_iam_integration(self, iam_integration: Union[str, IAMIntegration, None] = None) -> Optional[IAMIntegration]:
+          """
+          Sets the IAM integration for the dataset. IAM integration is used to sign URLs for data row assets.
+
+          >>>    # Get all IAM integrations
+          >>>    iam_integrations = client.get_organization().get_iam_integrations()
+          >>>    
+          >>>    # Get IAM integration id
+          >>>    iam_integration_id = [integration.uid for integration
+          >>>      in iam_integrations
+          >>>      if integration.name == "My S3 integration"][0]
+          >>>
+          >>>   # Set IAM integration for integration id
+          >>>   dataset.set_iam_integration(iam_integration_id)
+          >>>
+          >>>    # Get IAM integration object
+          >>>    iam_integration = [integration.uid for integration
+          >>>      in iam_integrations
+          >>>      if integration.name == "My S3 integration"][0]
+          >>>
+          >>>   # Set IAM integration for integration object
+          >>>   dataset.set_iam_integration(iam_integration)
+          >>>
+          >>>   # Unset IAM integration
+          >>>   dataset.set_iam_integration()
+          """
+
+          # Unset IAM integration if iam_integration is None
+          if iam_integration is None:
+              query = """mutation DetachSignerPyApi($id: ID!) {
+                  clearSignerForDataset(where: {id: $id}) {
+                      id
+                      signer {
+                      id
+                      }
+                  }
+                  }"""
+              response = self.client.execute(query, {"id": self.uid})
+              
+              if response:
+                 return response["clearSignerForDataset"]["signer"]
+              else:
+                raise lb.exceptions.LabelboxError("Can't unset IAM integration")
+              
+          else:
+
+              if isinstance(iam_integration, IAMIntegration):
+                  iam_integration_id = iam_integration.uid
+              else:
+                  iam_integration_id = iam_integration
+
+              query = """mutation AttachSignerPyApi($signerId: ID!, $datasetId: ID!) {
+                  setSignerForDataset(data: {signerId: $signerId}, where: {id: $datasetId}) {
+                      id
+                      signer {
+                      id
+                      }
+                  }
+                  }"""
+              response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid})
+              
+              # Return IAM Integration object if  
+              if response:
+                try:
+
+                  iam_integration_id = response.get("setSignerForDataset", {}).get("signer", {})["id"]
+                  return [integration for integration 
+                          in self.client.get_organization().get_iam_integrations() 
+                          if integration.uid == iam_integration_id][0]
+                except:
+                  raise LabelboxError(f"Can't retrieve IAM integration {iam_integration_id}")
+                
+              else:
+                raise LabelboxError(f"Can't set IAM integration {iam_integration_id}")
+
+          return response

libs/labelbox/tests/integration/test_delegated_access.py

@@ -2,6 +2,7 @@
 
 import requests
 import pytest
+import uuid
 
 from labelbox import Client
 
@@ -77,3 +78,52 @@ def test_no_default_integration(client):
     ds = client.create_dataset(name="new_ds")
     assert ds.iam_integration() is None
     ds.delete()
+
+
+@pytest.mark.skip(
+    reason=
+    "Google credentials are being updated for this test, disabling till it's all sorted out"
+)
+@pytest.mark.skipif(not os.environ.get("DA_GCP_LABELBOX_API_KEY"),
+                    reason="DA_GCP_LABELBOX_API_KEY not found")
+def test_integration_change():
+    """
+    This test is based on test_non_default_integration() and 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
+
+    Currently tests against:
+    Org ID: cl26d06tk0gch10901m7jeg9v
+    Email: jtso+aws_sdk_tests@labelbox.com
+    """
+    client = Client(api_key=os.environ.get("DA_GCP_LABELBOX_API_KEY"))
+    integrations = client.get_organization().get_iam_integrations()
+
+    # Prepare dataset with an existing integration
+    integration = [
+        integration for integration 
+        in integrations 
+        if 'aws-da-test-bucket' in integration.name][0]
+
+    ds = client.create_dataset(iam_integration=integration, name=f"integration_change-{uuid.uuid4()}")
+
+    # Test unset integration
+    ds.set_iam_integration()
+    assert ds.iam_integration() is None
+
+    # Test set integration with object
+    new_integration = ds.set_iam_integration(integration)
+    assert new_integration == integration
+
+    # Test set integration with integration id
+    integration_id = [
+        integration.uid for integration 
+        in integrations 
+        if 'aws-da-test-bucket' in integration.name][0]
+    
+    new_integration = ds.set_iam_integration(integration_id)
+    assert new_integration == integration
+
+    # Cleaning
+    ds.delete()