Add new workspace connection types. (Azure#31295)

* changes

* changelog

* redo recordings

* fix error in automatic recording creation

* api key support

* surface new credential class

* comments

Author: MilesHolland

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -5,6 +5,8 @@
 
 ### Features Added
 - Added support of features that are known into the future/at forecast time for dnn in AutoML Forecasting jobs.
+- Added support for new workspace connection types: azure_open_ai, cognitive_search, and cognitive_service.
+- Added support for new credential type: ApiKeyConfiguration.
 - Added support of `download` for component operations.
 
 ### Bugs Fixed
@@ -18,6 +20,7 @@
 
 - `azure-ai-ml` now performs all file i/o on `utf-8` encoded files per Azure SDK guidance. 
   (instead of the default behavior for python < 3.15, which uses locale specific encodings)
+- Removed references to deprecated "feature_store" workspace connection type.
 
 ## 1.9.0 (2023-07-25)
 
@@ -28,7 +31,6 @@
 
 - Improved the output when printing a workspace object to be more clean and readable.
 - Log level of unknown field notifications for pipeline nodes raised from INFO to WARNING.
-
 ## 1.8.0 (2023-06-12)
 
 ### Features Added

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -27,6 +27,7 @@
 from azure.ai.ml._restclient.v2020_09_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient092020DataplanePreview,
 )
+
 from azure.ai.ml._restclient.v2022_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022022Preview
 from azure.ai.ml._restclient.v2022_05_01 import AzureMachineLearningWorkspaces as ServiceClient052022
 from azure.ai.ml._restclient.v2022_10_01 import AzureMachineLearningWorkspaces as ServiceClient102022
@@ -305,6 +306,13 @@ def __init__(
             **kwargs,
         )
 
+        self._service_client_06_2023_preview = ServiceClient062023Preview(
+            credential=self._credential,
+            subscription_id=self._operation_scope._subscription_id,
+            base_url=base_url,
+            **kwargs,
+        )
+
         self._workspaces = WorkspaceOperations(
             self._operation_scope,
             self._service_client_06_2023_preview,
@@ -335,7 +343,7 @@ def __init__(
         self._workspace_connections = WorkspaceConnectionsOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_04_2023_preview,
+            self._service_client_06_2023_preview,
             self._operation_container,
             self._credential,
         )

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/connections/workspace_connection.py

@@ -6,7 +6,7 @@
 
 from marshmallow import fields, post_load
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import ConnectionCategory
+from azure.ai.ml._restclient.v2023_06_01_preview.models import ConnectionCategory
 from azure.ai.ml._schema.core.fields import ArmStr, NestedField, StringTransformedEnum, UnionField
 from azure.ai.ml._schema.core.schema import PathAwareSchema
 from azure.ai.ml._schema.job import CreationContextSchema
@@ -31,13 +31,15 @@ class WorkspaceConnectionSchema(PathAwareSchema):
             ConnectionCategory.GIT,
             ConnectionCategory.CONTAINER_REGISTRY,
             ConnectionCategory.PYTHON_FEED,
-            ConnectionCategory.FEATURE_STORE,
             ConnectionCategory.S3,
             ConnectionCategory.SNOWFLAKE,
             ConnectionCategory.AZURE_SQL_DB,
             ConnectionCategory.AZURE_SYNAPSE_ANALYTICS,
             ConnectionCategory.AZURE_MY_SQL_DB,
             ConnectionCategory.AZURE_POSTGRES_DB,
+            ConnectionCategory.AZURE_OPEN_AI,
+            ConnectionCategory.COGNITIVE_SERVICE,
+            ConnectionCategory.COGNITIVE_SEARCH,
         ],
         casing_transform=camel_to_snake,
         required=True,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py

@@ -50,6 +50,7 @@
     UserIdentityConfiguration,
     UsernamePasswordConfiguration,
     AccessKeyConfiguration,
+    ApiKeyConfiguration
 )
 from ._datastore.adls_gen1 import AzureDataLakeGen1Datastore
 from ._datastore.azure_storage import AzureBlobDatastore, AzureDataLakeGen2Datastore, AzureFileDatastore
@@ -397,6 +398,7 @@
     "Route",
     "AccessKeyConfiguration",
     "AlertNotification",
+    "ApiKeyConfiguration",
     "MonitorDefinition",
     "MonitorInputData",
     "MonitorSchedule",

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_credentials.py

@@ -42,6 +42,7 @@
 from azure.ai.ml._restclient.v2022_10_01.models import (
     ServicePrincipalDatastoreSecrets as RestServicePrincipalDatastoreSecrets,
 )
+
 from azure.ai.ml._restclient.v2023_04_01_preview.models import AmlToken as RestAmlToken
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ConnectionAuthType
 from azure.ai.ml._restclient.v2023_04_01_preview.models import IdentityConfiguration as RestJobIdentityConfiguration
@@ -52,10 +53,15 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     WorkspaceConnectionAccessKey as RestWorkspaceConnectionAccessKey,
 )
+from azure.ai.ml._restclient.v2023_06_01_preview.models import ConnectionAuthType
+from azure.ai.ml._restclient.v2023_06_01_preview.models import (
+    WorkspaceConnectionApiKey as WorkspaceConnectionApiKey,
+)
 from azure.ai.ml._utils.utils import camel_to_snake, snake_to_pascal
 from azure.ai.ml.constants._common import CommonYamlFields, IdentityType
 from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin, YamlTranslatableMixin
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, JobException, ValidationErrorType, ValidationException
+from azure.ai.ml._utils._experimental import experimental
 
 
 class _BaseIdentityConfiguration(ABC, DictMixin, RestTranslatableMixin):
@@ -776,3 +782,38 @@ def __eq__(self, other: object) -> bool:
         if not isinstance(other, AccessKeyConfiguration):
             return NotImplemented
         return self.access_key_id == other.access_key_id and self.secret_access_key == other.secret_access_key
+
+@experimental
+class ApiKeyConfiguration(RestTranslatableMixin, DictMixin):
+    """Api Key Credentials.
+
+    :param key: API key id
+    :type key: str
+    """
+
+    def __init__(
+        self,
+        *,
+        key: str,
+    ):
+        super().__init__()
+        self.type = camel_to_snake(ConnectionAuthType.API_KEY)
+        self.key = key
+
+    def _to_workspace_connection_rest_object(self) -> WorkspaceConnectionApiKey:
+        return WorkspaceConnectionApiKey(
+            key=self.key,
+        )
+
+    @classmethod
+    def _from_workspace_connection_rest_object(
+        cls, obj: Optional[WorkspaceConnectionApiKey]
+    ) -> "AccessKeyConfiguration":
+        return cls(
+            key=obj.key if obj is not None and obj.key else None,
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, AccessKeyConfiguration):
+            return NotImplemented
+        return self.key == other.key

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/connections/workspace_connection.py

@@ -9,18 +9,19 @@
 from pathlib import Path
 from typing import IO, Any, AnyStr, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2022_01_01_preview.models import (
+from azure.ai.ml._restclient.v2023_06_01_preview.models import (
     ManagedIdentityAuthTypeWorkspaceConnectionProperties,
     NoneAuthTypeWorkspaceConnectionProperties,
     PATAuthTypeWorkspaceConnectionProperties,
     SASAuthTypeWorkspaceConnectionProperties,
     ServicePrincipalAuthTypeWorkspaceConnectionProperties,
     UsernamePasswordAuthTypeWorkspaceConnectionProperties,
+    ApiKeyAuthWorkspaceConnectionProperties
 )
-from azure.ai.ml._restclient.v2022_01_01_preview.models import (
+from azure.ai.ml._restclient.v2023_06_01_preview.models import (
     WorkspaceConnectionPropertiesV2BasicResource as RestWorkspaceConnection,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.v2023_06_01_preview.models import (
     ConnectionAuthType,
     AccessKeyAuthTypeWorkspaceConnectionProperties,
 )
@@ -35,6 +36,7 @@
     ServicePrincipalConfiguration,
     UsernamePasswordConfiguration,
     AccessKeyConfiguration,
+    ApiKeyConfiguration
 )
 from azure.ai.ml.entities._resource import Resource
 from azure.ai.ml.entities._system_data import SystemData
@@ -54,12 +56,14 @@ class WorkspaceConnection(Resource):
     :type credentials: Union[
         ~azure.ai.ml.entities.PatTokenConfiguration, ~azure.ai.ml.entities.SasTokenConfiguration,
         ~azure.ai.ml.entities.UsernamePasswordConfiguration, ~azure.ai.ml.entities.ManagedIdentityConfiguration
-        ~azure.ai.ml.entities.ServicePrincipalConfiguration, ~azure.ai.ml.entities.AccessKeyConfiguration
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration, ~azure.ai.ml.entities.AccessKeyConfiguration,
+        ~azure.ai.ml.entities.ApiKeyConfiguration
         ]
     :param type: The category of external resource for this connection.
     :type type: The type of workspace connection, possible values are: [
         "git", "python_feed", "container_registry", "feature_store", "s3", "snowflake",
-         "azure_sql_db", "azure_synapse_analytics", "azure_my_sql_db", "azure_postgres_db"
+         "azure_sql_db", "azure_synapse_analytics", "azure_my_sql_db", "azure_postgres_db",
+         "azure_open_ai", "cognitive_search", "cognitive_service"
           ]
     """
 
@@ -203,6 +207,8 @@ def _from_rest_object(cls, rest_obj: RestWorkspaceConnection) -> "WorkspaceConne
             credentials = AccessKeyConfiguration._from_workspace_connection_rest_object(properties.credentials)
         if properties.auth_type == ConnectionAuthType.SERVICE_PRINCIPAL:
             credentials = ServicePrincipalConfiguration._from_workspace_connection_rest_object(properties.credentials)
+        if properties.auth_type == ConnectionAuthType.API_KEY:
+            credentials = ApiKeyConfiguration._from_workspace_connection_rest_object(properties.credentials)
 
         workspace_connection = WorkspaceConnection(
             id=rest_obj.id,
@@ -235,6 +241,8 @@ def _to_rest_object(self) -> RestWorkspaceConnection:
             workspace_connection_properties_class = SASAuthTypeWorkspaceConnectionProperties
         elif auth_type == camel_to_snake(ConnectionAuthType.SERVICE_PRINCIPAL):
             workspace_connection_properties_class = ServicePrincipalAuthTypeWorkspaceConnectionProperties
+        elif auth_type == camel_to_snake(ConnectionAuthType.API_KEY):
+            workspace_connection_properties_class = ApiKeyAuthWorkspaceConnectionProperties
         elif auth_type is None:
             workspace_connection_properties_class = NoneAuthTypeWorkspaceConnectionProperties
 

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_connections_operations.py

@@ -6,7 +6,7 @@
 
 from typing import Dict, Iterable, Optional
 
-from azure.ai.ml._restclient.v2022_05_01 import AzureMachineLearningWorkspaces as ServiceClient052022
+from azure.ai.ml._restclient.v2023_06_01_preview import AzureMachineLearningWorkspaces as ServiceClient062023Preview
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -34,7 +34,7 @@ def __init__(
         self,
         operation_scope: OperationScope,
         operation_config: OperationConfig,
-        service_client: ServiceClient052022,
+        service_client: ServiceClient062023Preview,
         all_operations: OperationsContainer,
         credentials: Optional[TokenCredential] = None,
         **kwargs: Dict,
@@ -79,7 +79,7 @@ def create_or_update(self, workspace_connection, **kwargs) -> WorkspaceConnectio
         response = self._operation.create(
             workspace_name=self._workspace_name,
             connection_name=workspace_connection.name,
-            parameters=rest_workspace_connection,
+            body=rest_workspace_connection,
             **self._scope_kwargs,
             **kwargs,
         )