1- from typing import List
1+ from typing import Dict
2+ from typing import Optional
23
4+ from lib .core .conditions import Condition
5+ from lib .core .conditions import CONDITION_EQ as EQ
36from lib .core .entities import FolderEntity
47from lib .core .entities import IntegrationEntity
58from lib .core .entities import ProjectEntity
9+ from lib .core .entities .integrations import IntegrationTypeEnum
10+ from lib .core .enums import ProjectType
611from lib .core .exceptions import AppException
712from lib .core .reporter import Reporter
813from lib .core .response import Response
@@ -25,6 +30,11 @@ def execute(self) -> Response:
2530
2631
2732class AttachIntegrations (BaseReportableUseCase ):
33+ MULTIMODAL_INTEGRATIONS = [
34+ IntegrationTypeEnum .DATABRICKS ,
35+ IntegrationTypeEnum .SNOWFLAKE ,
36+ ]
37+
2838 def __init__ (
2939 self ,
3040 reporter : Reporter ,
@@ -33,46 +43,139 @@ def __init__(
3343 service_provider : BaseServiceProvider ,
3444 integration : IntegrationEntity ,
3545 folder_path : str = None ,
46+ query : Optional [str ] = None ,
47+ item_name_column : Optional [str ] = None ,
48+ custom_item_name : Optional [str ] = None ,
49+ component_mapping : Optional [Dict [str , str ]] = None ,
3650 ):
37-
3851 super ().__init__ (reporter )
3952 self ._project = project
4053 self ._folder = folder
4154 self ._integration = integration
4255 self ._service_provider = service_provider
4356 self ._folder_path = folder_path
57+ self ._query = query
58+ self ._item_name_column = item_name_column
59+ self ._custom_item_name = custom_item_name
60+ self ._component_mapping = component_mapping
61+ self ._options = {} # using only for Databricks and Snowflake
62+ self ._item_category_column = None
4463
4564 @property
4665 def _upload_path (self ):
4766 return f"{ self ._project .name } { f'/{ self ._folder .name } if self ._folder .name != 'root' else '' }
4867
49- def execute (self ) -> Response :
50- integrations : List [
51- IntegrationEntity
52- ] = self ._service_provider .integrations .list ().data .integrations
53- integration_name_lower = self ._integration .name .lower ()
54- integration = next (
55- (i for i in integrations if i .name .lower () == integration_name_lower ), None
68+ def validate_integration (self ):
69+ # TODO add support in next iterations
70+ if self ._integration .type == IntegrationTypeEnum .SNOWFLAKE :
71+ raise AppException (
72+ "Attaching items is not supported with Snowflake integration."
73+ )
74+
75+ if self ._integration .type in self .MULTIMODAL_INTEGRATIONS :
76+ if self ._project .type != ProjectType .MULTIMODAL :
77+ raise AppException (
78+ f"{ self ._integration .name }
79+ )
80+
81+ def validate_options_for_multimodal_integration (self ):
82+ if self ._integration .type in self .MULTIMODAL_INTEGRATIONS :
83+ if self ._item_name_column and self ._custom_item_name :
84+ raise AppException (
85+ "‘item_name_column and custom_item_name cannot be used simultaneously."
86+ )
87+
88+ if not self ._item_name_column and not self ._custom_item_name :
89+ raise AppException (
90+ "Either item_name_column or custom_item_name is required."
91+ )
92+
93+ if not all ((self ._query , self ._component_mapping )):
94+ raise AppException (
95+ f"{ self ._integration .name }
96+ )
97+
98+ category_setting : bool = bool (
99+ next (
100+ (
101+ setting .value
102+ for setting in self ._service_provider .projects .list_settings (
103+ self ._project
104+ ).data
105+ if setting .attribute == "CategorizeItems"
106+ ),
107+ None ,
108+ )
109+ )
110+ if (
111+ not category_setting
112+ and "_item_category" in self ._component_mapping .values ()
113+ ):
114+ raise AppException (
115+ "Item Category must be enabled for a project to use _item_category"
116+ )
117+
118+ item_category_column = next (
119+ (
120+ k
121+ for k , v in self ._component_mapping .items ()
122+ if v == "_item_category"
123+ ),
124+ None ,
125+ )
126+ if item_category_column :
127+ self ._item_category_column = self ._component_mapping .pop (
128+ item_category_column
129+ )
130+
131+ sa_components = [
132+ c .name .lower ()
133+ for c in self ._service_provider .annotation_classes .list (
134+ condition = Condition ("project_id" , self ._project .id , EQ )
135+ ).data
136+ ]
137+
138+ for i in self ._component_mapping .values ():
139+ if i .lower () not in sa_components :
140+ raise AppException (
141+ f"Component mapping contains invalid component ID: `{ i }
142+ )
143+
144+ def generate_options_for_multimodal_integration (self ):
145+ self ._options ["query" ] = self ._query
146+ self ._options ["item_name" ] = (
147+ self ._custom_item_name if self ._custom_item_name else self ._item_name_column
56148 )
57- if integration :
149+ self ._options ["prefix" ] = True if self ._custom_item_name else False
150+ self ._options ["column_class_map" ] = self ._component_mapping
151+ if self ._item_category_column :
152+ self ._options ["item_category" ] = self ._item_category_column
153+
154+ def execute (self ) -> Response :
155+ if self .is_valid ():
156+ if self ._integration .type in self .MULTIMODAL_INTEGRATIONS :
157+ self .generate_options_for_multimodal_integration ()
158+
58159 self .reporter .log_info (
59160 "Attaching file(s) from "
60- f"{ integration .root } { f'/{ self ._folder_path } if self ._folder_path else '' }
161+ f"{ self . _integration .root } { f'/{ self ._folder_path } if self ._folder_path else '' }
61162 f"to { self ._upload_path }
62163 )
63- attached = self ._service_provider .integrations .attach_items (
164+
165+ attache_response = self ._service_provider .integrations .attach_items (
64166 project = self ._project ,
65167 folder = self ._folder ,
66- integration = integration ,
67- folder_name = self ._folder_path ,
168+ integration = self ._integration ,
169+ folder_name = self ._folder_path
170+ if self ._integration .type not in self .MULTIMODAL_INTEGRATIONS
171+ else None ,
172+ options = self ._options if self ._options else None ,
68173 )
69- if not attached :
174+ if not attache_response . ok :
70175 self ._response .errors = AppException (
71176 f"An error occurred for { self ._integration .name }
72177 "\n - The bucket exists."
73178 "\n - The connection is valid."
74179 "\n - The path to a specified directory is correct."
75180 )
76- else :
77- self ._response .errors = AppException ("Integration not found." )
78- return self ._response
181+ return self ._response
0 commit comments