From b96a6226a2f77cc4aa8b7f44e00ecee42fb0887d Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Mon, 20 May 2024 19:40:36 +0100 Subject: [PATCH 1/8] Add dataset.set_iam_integration() to select/deselect integrations --- libs/labelbox/src/labelbox/schema/dataset.py | 78 +++++++++++++++++++ .../integration/test_delegated_access.py | 50 ++++++++++++ 2 files changed, 128 insertions(+) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index 90a201028..be7deccfc 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/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 diff --git a/libs/labelbox/tests/integration/test_delegated_access.py b/libs/labelbox/tests/integration/test_delegated_access.py index 3f025a1ab..cb58c9071 100644 --- a/libs/labelbox/tests/integration/test_delegated_access.py +++ b/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() \ No newline at end of file From f85de736ee97a4d81eb2fc3836e0289250b9a225 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Mon, 20 May 2024 19:50:41 +0100 Subject: [PATCH 2/8] Fix typo --- libs/labelbox/src/labelbox/schema/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index be7deccfc..62186722f 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -934,7 +934,7 @@ def set_iam_integration(self, iam_integration: Union[str, IAMIntegration, None] >>> in iam_integrations >>> if integration.name == "My S3 integration"][0] >>> - >>> # Set IAM integration for integration object + >>> # Set IAM integration for IAMIntegrtion object >>> dataset.set_iam_integration(iam_integration) >>> >>> # Unset IAM integration @@ -956,7 +956,7 @@ def set_iam_integration(self, iam_integration: Union[str, IAMIntegration, None] if response: return response["clearSignerForDataset"]["signer"] else: - raise lb.exceptions.LabelboxError("Can't unset IAM integration") + raise LabelboxError("Can't unset IAM integration") else: From 16753c7da37aed8ad104fff6305c8668a237a5e3 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Tue, 21 May 2024 20:59:36 +0100 Subject: [PATCH 3/8] Split method and rename it. Split tests --- libs/labelbox/src/labelbox/schema/dataset.py | 164 ++++++++++-------- .../integration/test_delegated_access.py | 78 ++++++++- 2 files changed, 164 insertions(+), 78 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index 62186722f..0a54e8b7e 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -914,79 +914,103 @@ def _convert_items_to_upsert_format(self, _items): _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]: - """ + def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> 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 IAMIntegrtion object - >>> dataset.set_iam_integration(iam_integration) - >>> - >>> # Unset IAM integration - >>> dataset.set_iam_integration() + Args: + iam_integration (Union[str, IAMIntegration]): IAM integration object or IAM integration id. + + Returns: + IAMIntegration: IAM integration object. + + Raises: + LabelboxError: If the IAM integration can't be set. + + Examples: + + >>> # 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 IAMIntegrtion object + >>> dataset.set_iam_integration(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 LabelboxError("Can't unset IAM integration") + iam_integration_id = iam_integration.uid if isinstance(iam_integration, IAMIntegration) else iam_integration + + query = """ + mutation SetSignerForDatasetPyApi($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}) + + if not response: + raise LabelboxError(f"Can't set IAM integration {iam_integration_id}") + + 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: - - 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}) + def remove_iam_integration(self) -> None: + """ + Unsets the IAM integration for the dataset. + + Args: + None + + Returns: + None + + Raises: + LabelboxError: If the IAM integration can't be unset. + + Examples: + >>> dataset.remove_iam_integration() + """ + + query = """ + mutation DetachSignerPyApi($id: ID!) { + clearSignerForDataset(where: { id: $id }) { + id + signer { + id + } + } + } + """ + + response = self.client.execute(query, {"id": self.uid}) + + if response: + return response.get("clearSignerForDataset", {}).get("signer") + else: + raise LabelboxError("Can't unset IAM integration") - # 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 + \ No newline at end of file diff --git a/libs/labelbox/tests/integration/test_delegated_access.py b/libs/labelbox/tests/integration/test_delegated_access.py index cb58c9071..edee1564b 100644 --- a/libs/labelbox/tests/integration/test_delegated_access.py +++ b/libs/labelbox/tests/integration/test_delegated_access.py @@ -86,7 +86,7 @@ def test_no_default_integration(client): ) @pytest.mark.skipif(not os.environ.get("DA_GCP_LABELBOX_API_KEY"), reason="DA_GCP_LABELBOX_API_KEY not found") -def test_integration_change(): +def test_add_integration_from_object(): """ This test is based on test_non_default_integration() and assumes the following: @@ -106,24 +106,86 @@ def test_integration_change(): 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 + ds = client.create_dataset(iam_integration=None, name=f"integration_change-{uuid.uuid4()}") # Test set integration with object - new_integration = ds.set_iam_integration(integration) + new_integration = ds.add_iam_integration(integration) assert new_integration == integration + # Cleaning + 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_add_integration_from_uid(): + """ + 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=None, name=f"integration_change-{uuid.uuid4()}") + # 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) + new_integration = ds.add_iam_integration(integration_id) assert new_integration == integration + # Cleaning + 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_remove(): + """ + 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.remove_iam_integration() + assert ds.iam_integration() is None + # Cleaning ds.delete() \ No newline at end of file From c9b93a6694b3067389dc0c5e49704400279294f2 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Tue, 21 May 2024 21:08:11 +0100 Subject: [PATCH 4/8] Fix indentation --- .../tests/integration/test_delegated_access.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/libs/labelbox/tests/integration/test_delegated_access.py b/libs/labelbox/tests/integration/test_delegated_access.py index edee1564b..1592319d2 100644 --- a/libs/labelbox/tests/integration/test_delegated_access.py +++ b/libs/labelbox/tests/integration/test_delegated_access.py @@ -100,13 +100,13 @@ def test_add_integration_from_object(): 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 + # Prepare dataset with no integration integration = [ integration for integration in integrations if 'aws-da-test-bucket' in integration.name][0] - ds = client.create_dataset(iam_integration=None, name=f"integration_change-{uuid.uuid4()}") + ds = client.create_dataset(iam_integration=None, name=f"integration_add_obj-{uuid.uuid4()}") # Test set integration with object new_integration = ds.add_iam_integration(integration) @@ -135,13 +135,13 @@ def test_add_integration_from_uid(): 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 + # Prepare dataset with no integration integration = [ integration for integration in integrations if 'aws-da-test-bucket' in integration.name][0] - ds = client.create_dataset(iam_integration=None, name=f"integration_change-{uuid.uuid4()}") + ds = client.create_dataset(iam_integration=None, name=f"integration_add_id-{uuid.uuid4()}") # Test set integration with integration id integration_id = [ @@ -155,7 +155,7 @@ def test_add_integration_from_uid(): # Cleaning ds.delete() - @pytest.mark.skip( +@pytest.mark.skip( reason= "Google credentials are being updated for this test, disabling till it's all sorted out" ) @@ -181,7 +181,7 @@ def test_integration_remove(): in integrations if 'aws-da-test-bucket' in integration.name][0] - ds = client.create_dataset(iam_integration=integration, name=f"integration_change-{uuid.uuid4()}") + ds = client.create_dataset(iam_integration=integration, name=f"integration_remove-{uuid.uuid4()}") # Test unset integration ds.remove_iam_integration() From b177143d5176f46d4dd43bf72ca9bad3952555e7 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Wed, 22 May 2024 00:48:03 +0100 Subject: [PATCH 5/8] Fix exceptions --- libs/labelbox/src/labelbox/schema/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index 0a54e8b7e..d897ee54b 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -968,7 +968,7 @@ def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IA response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid}) if not response: - raise LabelboxError(f"Can't set IAM integration {iam_integration_id}") + raise ResourceNotFoundError(IAMIntegration, {"signerId": iam_integration_id, "datasetId": self.uid}) try: iam_integration_id = response.get("setSignerForDataset", {}).get("signer", {})["id"] @@ -1011,6 +1011,6 @@ def remove_iam_integration(self) -> None: if response: return response.get("clearSignerForDataset", {}).get("signer") else: - raise LabelboxError("Can't unset IAM integration") + raise ResourceNotFoundError(IAMIntegration, {"datasetId": self.uid}) \ No newline at end of file From bb837163edbb5611958df807e2dab0d74575bf69 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Wed, 22 May 2024 22:30:21 +0100 Subject: [PATCH 6/8] Clean remove_iam_integration() --- libs/labelbox/src/labelbox/schema/dataset.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index d897ee54b..edf2d4c94 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -999,18 +999,15 @@ def remove_iam_integration(self) -> None: mutation DetachSignerPyApi($id: ID!) { clearSignerForDataset(where: { id: $id }) { id - signer { - id - } } } """ response = self.client.execute(query, {"id": self.uid}) - if response: - return response.get("clearSignerForDataset", {}).get("signer") - else: - raise ResourceNotFoundError(IAMIntegration, {"datasetId": self.uid}) + if not response: + raise ResourceNotFoundError(Dataset, {"id": self.uid}) + + \ No newline at end of file From 7c51d6edd9afcb56d42a1589469705e902a59293 Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Thu, 23 May 2024 00:28:30 +0100 Subject: [PATCH 7/8] Remove final space --- libs/labelbox/src/labelbox/schema/dataset.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index edf2d4c94..eda30f34d 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -1007,7 +1007,4 @@ def remove_iam_integration(self) -> None: if not response: raise ResourceNotFoundError(Dataset, {"id": self.uid}) - - - - \ No newline at end of file + \ No newline at end of file From 66e98b582de5db19e915d8bd21144427d7bd1ada Mon Sep 17 00:00:00 2001 From: paulnoirel <87332996+paulnoirel@users.noreply.github.com> Date: Thu, 23 May 2024 00:46:46 +0100 Subject: [PATCH 8/8] Reformat add_iam_integration() --- libs/labelbox/src/labelbox/schema/dataset.py | 85 ++++++++++---------- 1 file changed, 43 insertions(+), 42 deletions(-) diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py index eda30f34d..0bc5b74b1 100644 --- a/libs/labelbox/src/labelbox/schema/dataset.py +++ b/libs/labelbox/src/labelbox/schema/dataset.py @@ -915,43 +915,43 @@ def _convert_items_to_upsert_format(self, _items): return _upsert_items def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IAMIntegration: - """ - Sets the IAM integration for the dataset. IAM integration is used to sign URLs for data row assets. + """ + Sets the IAM integration for the dataset. IAM integration is used to sign URLs for data row assets. - Args: - iam_integration (Union[str, IAMIntegration]): IAM integration object or IAM integration id. - - Returns: - IAMIntegration: IAM integration object. - - Raises: - LabelboxError: If the IAM integration can't be set. - - Examples: - - >>> # 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 IAMIntegrtion object - >>> dataset.set_iam_integration(iam_integration) - """ - - iam_integration_id = iam_integration.uid if isinstance(iam_integration, IAMIntegration) else iam_integration - - query = """ + Args: + iam_integration (Union[str, IAMIntegration]): IAM integration object or IAM integration id. + + Returns: + IAMIntegration: IAM integration object. + + Raises: + LabelboxError: If the IAM integration can't be set. + + Examples: + + >>> # 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 IAMIntegrtion object + >>> dataset.set_iam_integration(iam_integration) + """ + + iam_integration_id = iam_integration.uid if isinstance(iam_integration, IAMIntegration) else iam_integration + + query = """ mutation SetSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) { setSignerForDataset( data: { signerId: $signerId } @@ -963,19 +963,20 @@ def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IA } } } - """ + """ - response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid}) + response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid}) - if not response: + if not response: raise ResourceNotFoundError(IAMIntegration, {"signerId": iam_integration_id, "datasetId": self.uid}) - - try: + + 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: + except: raise LabelboxError(f"Can't retrieve IAM integration {iam_integration_id}") def remove_iam_integration(self) -> None: