@@ -267,11 +267,22 @@ def _load_mapping(self) -> dict[str, Tag]:
267267 pkgs = {p .name : p for p in self .packages }
268268 spec = json .loads (self ._openapi_spec ())
269269 for tag in spec ['tags' ]:
270+ is_account = tag .get ('x-databricks-is-accounts' )
271+ # Unique identifier for the tag. Note that the service name may not be unique
272+ key = 'a' if is_account else 'w'
273+ parent_service = tag .get ('x-databricks-parent-service' )
274+ if parent_service :
275+ # SDK generation removes the "account" prefix from account services
276+ clean_parent_service = parent_service .lower ().removeprefix ("account" )
277+ key = f"{ key } .{ clean_parent_service } "
278+
279+ key = f"{ key } .{ tag ['x-databricks-service' ]} " .lower ()
280+
270281 t = Tag (name = tag ['name' ],
271282 service = tag ['x-databricks-service' ],
272283 is_account = tag .get ('x-databricks-is-accounts' , False ),
273284 package = pkgs [tag ['x-databricks-package' ]])
274- mapping [tag [ 'name' ] ] = t
285+ mapping [key ] = t
275286 return mapping
276287
277288 @staticmethod
@@ -360,7 +371,7 @@ def service_docs(self, client_inst, client_prefix: str) -> list[ServiceDoc]:
360371 service_name = service_name ,
361372 class_name = class_name ,
362373 doc = class_doc ,
363- tag = self ._get_tag_name (service_inst .__class__ .__name__ , service_name ),
374+ tag = self ._get_tag_name (service_inst .__class__ .__name__ , client_prefix , service_name ),
364375 methods = self .class_methods (service_inst ),
365376 property = self .class_properties (service_inst )))
366377 return all
@@ -399,13 +410,13 @@ def write_dataclass_docs(self):
399410
400411 { all } ''' )
401412
402- def _get_tag_name (self , class_name , service_name ) -> Tag :
413+ def _get_tag_name (self , class_name , client_prefix , service_name ) -> Tag :
403414 if class_name [- 3 :] == 'Ext' :
404415 # ClustersExt, DbfsExt, WorkspaceExt, but not ExternalLocations
405416 class_name = class_name .replace ('Ext' , 'API' )
406417 class_name = class_name [:- 3 ]
407- for tag_name , t in self .mapping .items ():
408- if t . service . lower () == str (class_name ).lower ():
418+ for key , t in self .mapping .items ():
419+ if key == f' { client_prefix } . { str (class_name ).lower ()} ' :
409420 return t
410421 raise KeyError (f'Cannot find { class_name } / { service_name } tag' )
411422
0 commit comments