From 558b21af3ef648b5fc6d4903557381964c621fcb Mon Sep 17 00:00:00 2001 From: Jayesh Tanna Date: Mon, 12 Aug 2024 15:59:44 +0530 Subject: [PATCH] making workspace connections prop ga (#36781) (#36849) * making wsconnection object and props ga * resolving pr comments related to ga entity * skipping e2e test --- sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py | 1 - .../_workspace/connections/workspace_connection.py | 7 ++----- .../ml/operations/_workspace_connections_operations.py | 10 ++++++++-- .../tests/feature_set/e2etests/test_feature_set.py | 1 + 4 files changed, 11 insertions(+), 8 deletions(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py index 71b4d322507b..a04f831d0c6a 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py @@ -942,7 +942,6 @@ def feature_store_entities(self) -> FeatureStoreEntityOperations: return self._featurestoreentities @property - @experimental def connections(self) -> WorkspaceConnectionsOperations: """A collection of connection related operations. diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/connections/workspace_connection.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/connections/workspace_connection.py index 873fdcd19376..79e25c3e5d8d 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/connections/workspace_connection.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/connections/workspace_connection.py @@ -422,10 +422,7 @@ def _to_dict(self) -> Dict: return res @classmethod - def _from_rest_object(cls, rest_obj: RestWorkspaceConnection) -> Optional["WorkspaceConnection"]: - if not rest_obj: - return None - + def _from_rest_object(cls, rest_obj: RestWorkspaceConnection) -> "WorkspaceConnection": conn_class = cls._get_entity_class_from_rest_obj(rest_obj) popped_metadata = conn_class._get_required_metadata_fields() @@ -456,7 +453,7 @@ def _from_rest_object(cls, rest_obj: RestWorkspaceConnection) -> Optional["Works # No default in pop, this should fail if we somehow don't get a resource ID rest_kwargs["ai_services_resource_id"] = rest_kwargs.pop(camel_to_snake(CONNECTION_RESOURCE_ID_KEY)) connection = conn_class(**rest_kwargs) - return cast(Optional["WorkspaceConnection"], connection) + return cast(WorkspaceConnection, connection) def _validate(self) -> str: return str(self.name) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_connections_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_connections_operations.py index f5286a737469..c8e0c3dccc1e 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_connections_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_connections_operations.py @@ -18,6 +18,7 @@ from azure.ai.ml._utils.utils import _snake_to_camel from azure.ai.ml.entities._workspace.connections.workspace_connection import WorkspaceConnection from azure.core.credentials import TokenCredential +from azure.core.tracing.decorator import distributed_trace from azure.ai.ml.entities._credentials import ( ApiKeyConfiguration, ) @@ -70,6 +71,7 @@ def _try_fill_api_key(self, connection: WorkspaceConnection) -> None: if list_secrets_response.properties.credentials is not None: connection.credentials.key = list_secrets_response.properties.credentials.key + @distributed_trace @monitor_with_activity(ops_logger, "WorkspaceConnections.Get", ActivityType.PUBLICAPI) def get(self, name: str, *, populate_secrets: bool = False, **kwargs: Dict) -> WorkspaceConnection: """Get a connection by name. @@ -98,10 +100,11 @@ def get(self, name: str, *, populate_secrets: bool = False, **kwargs: Dict) -> W self._try_fill_api_key(connection) return connection # type: ignore[return-value] + @distributed_trace @monitor_with_activity(ops_logger, "WorkspaceConnections.CreateOrUpdate", ActivityType.PUBLICAPI) def create_or_update( self, workspace_connection: WorkspaceConnection, *, populate_secrets: bool = False, **kwargs: Any - ) -> Optional[WorkspaceConnection]: + ) -> WorkspaceConnection: """Create or update a connection. :param workspace_connection: Definition of a Workspace Connection or one of its subclasses @@ -126,8 +129,9 @@ def create_or_update( self._try_fill_api_key(conn) return conn + @distributed_trace @monitor_with_activity(ops_logger, "WorkspaceConnections.Delete", ActivityType.PUBLICAPI) - def delete(self, name: str) -> None: + def delete(self, name: str, **kwargs: Any) -> None: """Delete the connection. :param name: Name of the connection. @@ -138,8 +142,10 @@ def delete(self, name: str) -> None: connection_name=name, workspace_name=self._workspace_name, **self._scope_kwargs, + **kwargs, ) + @distributed_trace @monitor_with_activity(ops_logger, "WorkspaceConnections.List", ActivityType.PUBLICAPI) def list( self, diff --git a/sdk/ml/azure-ai-ml/tests/feature_set/e2etests/test_feature_set.py b/sdk/ml/azure-ai-ml/tests/feature_set/e2etests/test_feature_set.py index 3b93043ee891..4f843920fdc9 100644 --- a/sdk/ml/azure-ai-ml/tests/feature_set/e2etests/test_feature_set.py +++ b/sdk/ml/azure-ai-ml/tests/feature_set/e2etests/test_feature_set.py @@ -61,6 +61,7 @@ def feature_set_validation(fset: FeatureSet): # ---------------------------------------------------------------------------------------------------------------# # NOTE Please enable materialization store on test featurestore 'sdk_vnext_cli_fs' to run this test in live mode. # ---------------------------------------------------------------------------------------------------------------# + @pytest.mark.skip(reason="request header size being too large.") def test_list_materialization_jobs( self, feature_store_client: MLClient, tmp_path: Path, randstr: Callable[[], str] ) -> None: