18
18
from azure .ai .ml ._utils .utils import _snake_to_camel
19
19
from azure .ai .ml .entities ._workspace .connections .workspace_connection import WorkspaceConnection
20
20
from azure .core .credentials import TokenCredential
21
+ from azure .core .tracing .decorator import distributed_trace
21
22
from azure .ai .ml .entities ._credentials import (
22
23
ApiKeyConfiguration ,
23
24
)
@@ -70,6 +71,7 @@ def _try_fill_api_key(self, connection: WorkspaceConnection) -> None:
70
71
if list_secrets_response .properties .credentials is not None :
71
72
connection .credentials .key = list_secrets_response .properties .credentials .key
72
73
74
+ @distributed_trace
73
75
@monitor_with_activity (ops_logger , "WorkspaceConnections.Get" , ActivityType .PUBLICAPI )
74
76
def get (self , name : str , * , populate_secrets : bool = False , ** kwargs : Dict ) -> WorkspaceConnection :
75
77
"""Get a connection by name.
@@ -98,10 +100,11 @@ def get(self, name: str, *, populate_secrets: bool = False, **kwargs: Dict) -> W
98
100
self ._try_fill_api_key (connection )
99
101
return connection # type: ignore[return-value]
100
102
103
+ @distributed_trace
101
104
@monitor_with_activity (ops_logger , "WorkspaceConnections.CreateOrUpdate" , ActivityType .PUBLICAPI )
102
105
def create_or_update (
103
106
self , workspace_connection : WorkspaceConnection , * , populate_secrets : bool = False , ** kwargs : Any
104
- ) -> Optional [ WorkspaceConnection ] :
107
+ ) -> WorkspaceConnection :
105
108
"""Create or update a connection.
106
109
107
110
:param workspace_connection: Definition of a Workspace Connection or one of its subclasses
@@ -126,8 +129,9 @@ def create_or_update(
126
129
self ._try_fill_api_key (conn )
127
130
return conn
128
131
132
+ @distributed_trace
129
133
@monitor_with_activity (ops_logger , "WorkspaceConnections.Delete" , ActivityType .PUBLICAPI )
130
- def delete (self , name : str ) -> None :
134
+ def delete (self , name : str , ** kwargs : Any ) -> None :
131
135
"""Delete the connection.
132
136
133
137
:param name: Name of the connection.
@@ -138,8 +142,10 @@ def delete(self, name: str) -> None:
138
142
connection_name = name ,
139
143
workspace_name = self ._workspace_name ,
140
144
** self ._scope_kwargs ,
145
+ ** kwargs ,
141
146
)
142
147
148
+ @distributed_trace
143
149
@monitor_with_activity (ops_logger , "WorkspaceConnections.List" , ActivityType .PUBLICAPI )
144
150
def list (
145
151
self ,
0 commit comments