31
31
from labelbox .schema .identifiable import UniqueId , GlobalKey
32
32
from labelbox .schema .task import Task , DataUpsertTask
33
33
from labelbox .schema .user import User
34
- < << << << HEAD
35
34
from labelbox .schema .iam_integration import IAMIntegration
36
- == == == =
37
35
from labelbox .schema .internal .data_row_create_upsert import (DataRowItemBase ,
38
36
DataRowUpsertItem ,
39
37
DataRowCreateItem )
40
38
from labelbox .schema .internal .data_row_uploader import DataRowUploader
41
39
from labelbox .schema .internal .datarow_upload_constants import (
42
40
MAX_DATAROW_PER_API_OPERATION , FILE_UPLOAD_THREAD_COUNT , UPSERT_CHUNK_SIZE )
43
- > >> >> >> 58e48 ccf (Refactor upsert code so that it can be reused for create )
44
41
45
42
logger = logging .getLogger (__name__ )
46
43
@@ -644,11 +641,13 @@ def _exec_upsert_data_rows(
644
641
task ._user = self .client .get_user ()
645
642
return task
646
643
647
- def add_iam_integration (self , iam_integration : Union [str , IAMIntegration ]) -> IAMIntegration :
644
+ def add_iam_integration (
645
+ self , iam_integration : Union [str ,
646
+ IAMIntegration ]) -> IAMIntegration :
648
647
"""
649
648
Sets the IAM integration for the dataset. IAM integration is used to sign URLs for data row assets.
650
649
651
- Args:
650
+ Args:
652
651
iam_integration (Union[str, IAMIntegration]): IAM integration object or IAM integration id.
653
652
654
653
Returns:
@@ -679,7 +678,8 @@ def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IA
679
678
>>> dataset.set_iam_integration(iam_integration)
680
679
"""
681
680
682
- iam_integration_id = iam_integration .uid if isinstance (iam_integration , IAMIntegration ) else iam_integration
681
+ iam_integration_id = iam_integration .uid if isinstance (
682
+ iam_integration , IAMIntegration ) else iam_integration
683
683
684
684
query = """
685
685
mutation SetSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) {
@@ -695,20 +695,30 @@ def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IA
695
695
}
696
696
"""
697
697
698
- response = self .client .execute (query , {"signerId" : iam_integration_id , "datasetId" : self .uid })
698
+ response = self .client .execute (query , {
699
+ "signerId" : iam_integration_id ,
700
+ "datasetId" : self .uid
701
+ })
699
702
700
703
if not response :
701
- raise ResourceNotFoundError (IAMIntegration , {"signerId" : iam_integration_id , "datasetId" : self . uid })
702
-
703
- try :
704
- iam_integration_id = response . get ( "setSignerForDataset" , {}). get ( "signer" , {})[ "id" ]
704
+ raise ResourceNotFoundError (IAMIntegration , {
705
+ "signerId" : iam_integration_id ,
706
+ "datasetId" : self . uid
707
+ })
705
708
706
- return [integration for integration
707
- in self .client .get_organization ().get_iam_integrations ()
708
- if integration .uid == iam_integration_id ][0 ]
709
+ try :
710
+ iam_integration_id = response .get ("setSignerForDataset" ,
711
+ {}).get ("signer" , {})["id" ]
712
+
713
+ return [
714
+ integration for integration in
715
+ self .client .get_organization ().get_iam_integrations ()
716
+ if integration .uid == iam_integration_id
717
+ ][0 ]
709
718
except :
710
- raise LabelboxError (f"Can't retrieve IAM integration { iam_integration_id } " )
711
-
719
+ raise LabelboxError (
720
+ f"Can't retrieve IAM integration { iam_integration_id } " )
721
+
712
722
def remove_iam_integration (self ) -> None :
713
723
"""
714
724
Unsets the IAM integration for the dataset.
@@ -738,4 +748,3 @@ def remove_iam_integration(self) -> None:
738
748
739
749
if not response :
740
750
raise ResourceNotFoundError (Dataset , {"id" : self .uid })
741
-
0 commit comments