@@ -56,59 +56,84 @@ def convert_number_type(cls, agate_table: agate.Table, col_idx: int) -> str:
5656 def convert_datetime_type (cls , agate_table : agate .Table , col_idx : int ) -> str :
5757 return "timestamp"
5858
59+ @classmethod
60+ def parse_lf_response (
61+ cls ,
62+ response : Dict [str , Any ],
63+ database : str ,
64+ table : Optional [str ],
65+ columns : Optional [List [str ]],
66+ lf_tags : Dict [str , str ],
67+ ) -> str :
68+ failures = response .get ("Failures" , [])
69+ tbl_appendix = f".{ table } if table else ""
70+ columns_appendix = f" for columns { columns } if columns else ""
71+ msg_appendix = tbl_appendix + columns_appendix
72+ if failures :
73+ base_msg = f"Failed to add LF tags: { lf_tags } { database } + msg_appendix
74+ for failure in failures :
75+ tag = failure .get ("LFTag" , {}).get ("TagKey" )
76+ error = failure .get ("Error" , {}).get ("ErrorMessage" )
77+ logger .error (f"Failed to set { tag } { database } + msg_appendix + f" - { error } )
78+ raise DbtRuntimeError (base_msg )
79+ return f"Added LF tags: { lf_tags } { database } + msg_appendix
80+
81+ @classmethod
82+ def lf_tags_columns_is_valid (cls , lf_tags_columns : Dict [str , Dict [str , List [str ]]]) -> Optional [bool ]:
83+ if not lf_tags_columns :
84+ return False
85+ for tag_key , tag_config in lf_tags_columns .items ():
86+ if isinstance (tag_config , Dict ):
87+ for tag_value , columns in tag_config .items ():
88+ if not isinstance (columns , List ):
89+ raise DbtRuntimeError (f"Not a list: { columns } + "Expected format: ['c1', 'c2']" )
90+ else :
91+ raise DbtRuntimeError (f"Not a dict: { tag_config } + "Expected format: {'tag_value': ['c1', 'c2']}" )
92+ return True
93+
5994 # TODO: Add more lf-tag unit tests when moto supports lakeformation
6095 # moto issue: https://github.com/getmoto/moto/issues/5964
6196 @available
62- def add_lf_tags (self , database : str , table : str = None , lf_tags : Dict [str , str ] = None ):
97+ def add_lf_tags (
98+ self ,
99+ database : str ,
100+ table : str = None ,
101+ lf_tags : Optional [Dict [str , str ]] = None ,
102+ lf_tags_columns : Optional [Dict [str , Dict [str , List [str ]]]] = None ,
103+ ):
63104 conn = self .connections .get_thread_connection ()
64105 client = conn .handle
65106
66107 lf_tags = lf_tags or conn .credentials .lf_tags
67- if not lf_tags :
68- logger .debug ("No LF tags configured" )
69- return
70-
71- resource = {
72- "Database" : {"Name" : database },
73- }
74108
75- if table :
76- resource = {
77- "Table" : {
78- "DatabaseName" : database ,
79- "Name" : table ,
80- }
81- }
82-
83- with boto3_client_lock :
84- lf_client = client .session .client (
85- "lakeformation" , region_name = client .region_name , config = get_boto3_config ()
86- )
109+ if not lf_tags and not lf_tags_columns :
110+ logger .debug ("No LF tags configured" )
111+ else :
112+ with boto3_client_lock :
113+ lf_client = client .session .client (
114+ "lakeformation" , region_name = client .region_name , config = get_boto3_config ()
115+ )
87116
88- response = lf_client .add_lf_tags_to_resource (
89- Resource = resource ,
90- LFTags = [
91- {
92- "TagKey" : key ,
93- "TagValues" : [
94- value ,
95- ],
96- }
97- for key , value in lf_tags .items ()
98- ],
99- )
117+ if lf_tags :
118+ resource = {"Database" : {"Name" : database }}
119+ if table :
120+ resource = {"Table" : {"DatabaseName" : database , "Name" : table }}
100121
101- failures = response .get ("Failures" , [])
102- tbl_appendix = f".{ table } if table else ""
103- if failures :
104- base_msg = f"Failed to add LF tags: { lf_tags } { database } + tbl_appendix
105- for failure in failures :
106- tag = failure .get ("LFTag" , {}).get ("TagKey" )
107- error = failure .get ("Error" , {}).get ("ErrorMessage" )
108- logger .error (f"Failed to set { tag } { database } + tbl_appendix + f" - { error } )
109- raise DbtRuntimeError (base_msg )
110- else :
111- logger .debug (f"Added LF tags: { lf_tags } { database } + tbl_appendix )
122+ response = lf_client .add_lf_tags_to_resource (
123+ Resource = resource , LFTags = [{"TagKey" : key , "TagValues" : [value ]} for key , value in lf_tags .items ()]
124+ )
125+ logger .debug (self .parse_lf_response (response , database , table , None , lf_tags ))
126+
127+ if self .lf_tags_columns_is_valid (lf_tags_columns ):
128+ for tag_key , tag_config in lf_tags_columns .items ():
129+ for tag_value , columns in tag_config .items ():
130+ response = lf_client .add_lf_tags_to_resource (
131+ Resource = {
132+ "TableWithColumns" : {"DatabaseName" : database , "Name" : table , "ColumnNames" : columns }
133+ },
134+ LFTags = [{"TagKey" : tag_key , "TagValues" : [tag_value ]}],
135+ )
136+ logger .debug (self .parse_lf_response (response , database , table , columns , {tag_key : tag_value }))
112137
113138 @available
114139 def get_work_group_output_location (self ) -> Optional [str ]:
0 commit comments