Update datastore rest client to 2023-04-01-preview + OneLake datastore fixes (Azure#31517)

* Fixed serialization issue in OneLakeArtifact and updated datastore service client to support OneLake datastore types. Updated test fixtures using datastores to also use mock updated datastore service clients.

* artifact_type: lakehouse -> artifact_type: lake_house

* Fixed equality comparison.

* Fixed credential getting set to NoneCredentialConfiguration by updating the restclient to the correct version (2023_04_01_preview) used in DatastoreOperations. Also updated other datastore entities to use the 2023_04_01_preview client.

* Added test to validate credential less onelake datastore schema.

* Formatting fixes for pylint.

* Code formatting updates from running black.

* Added docstrings for newly added OneLakeArtifact and LakeHouseArtifact classes.

Author: youngpark

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

@@ -363,7 +363,7 @@ def __init__(
         self._datastores = DatastoreOperations(
             operation_scope=self._operation_scope,
             operation_config=self._operation_config,
-            serviceclient_2022_10_01=self._service_client_10_2022,
+            serviceclient_2023_04_01_preview=self._service_client_04_2023_preview,
             **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.DATASTORE, self._datastores)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_datastore/one_lake.py

@@ -6,19 +6,21 @@
 
 from typing import Any, Dict
 
-from marshmallow import fields, post_load
+from marshmallow import fields, post_load, Schema
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType, OneLakeArtifactType
 
 from azure.ai.ml._schema.core.fields import NestedField, PathAwareSchema, StringTransformedEnum, UnionField
 from azure.ai.ml._utils.utils import camel_to_snake
 
 from .credentials import NoneCredentialsSchema, ServicePrincipalSchema
 
 
-class OneLakeArtifactSchema(PathAwareSchema):
+class OneLakeArtifactSchema(Schema):
     artifact_name = fields.Str(required=True)
-    artifact_type = fields.Str()
+    artifact_type = StringTransformedEnum(
+        allowed_values=OneLakeArtifactType.LAKE_HOUSE, casing_transform=camel_to_snake
+    )
 
 
 class OneLakeSchema(PathAwareSchema):

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

@@ -25,21 +25,23 @@
 )
 from azure.ai.ml._restclient.v2022_05_01.models import ManagedServiceIdentity as RestManagedServiceIdentityConfiguration
 from azure.ai.ml._restclient.v2022_05_01.models import UserAssignedIdentity as RestUserAssignedIdentityConfiguration
-from azure.ai.ml._restclient.v2022_10_01.models import (
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     AccountKeyDatastoreCredentials as RestAccountKeyDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2022_10_01.models import AccountKeyDatastoreSecrets as RestAccountKeyDatastoreSecrets
-from azure.ai.ml._restclient.v2022_10_01.models import (
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    AccountKeyDatastoreSecrets as RestAccountKeyDatastoreSecrets,
+)
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     CertificateDatastoreCredentials as RestCertificateDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2022_10_01.models import CertificateDatastoreSecrets, CredentialsType
-from azure.ai.ml._restclient.v2022_10_01.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
-from azure.ai.ml._restclient.v2022_10_01.models import SasDatastoreCredentials as RestSasDatastoreCredentials
-from azure.ai.ml._restclient.v2022_10_01.models import SasDatastoreSecrets as RestSasDatastoreSecrets
-from azure.ai.ml._restclient.v2022_10_01.models import (
+from azure.ai.ml._restclient.v2023_04_01_preview.models import CertificateDatastoreSecrets, CredentialsType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
+from azure.ai.ml._restclient.v2023_04_01_preview.models import SasDatastoreCredentials as RestSasDatastoreCredentials
+from azure.ai.ml._restclient.v2023_04_01_preview.models import SasDatastoreSecrets as RestSasDatastoreSecrets
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     ServicePrincipalDatastoreCredentials as RestServicePrincipalDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2022_10_01.models import (
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     ServicePrincipalDatastoreSecrets as RestServicePrincipalDatastoreSecrets,
 )
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/_on_prem.py

@@ -8,8 +8,8 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2022_10_01_preview.models import Datastore as DatastoreData, DatastoreType
-from azure.ai.ml._restclient.v2022_10_01_preview.models import HdfsDatastore as RestHdfsDatastore
+from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData, DatastoreType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import HdfsDatastore as RestHdfsDatastore
 from azure.ai.ml._schema._datastore._on_prem import HdfsSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/_on_prem_credentials.py

@@ -4,7 +4,7 @@
 
 from base64 import b64encode
 
-from azure.ai.ml._restclient.v2022_10_01_preview import models as model_preview
+from azure.ai.ml._restclient.v2023_04_01_preview import models as model_preview
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.entities._credentials import NoneCredentialConfiguration
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/adls_gen1.py

@@ -7,8 +7,10 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2022_10_01.models import AzureDataLakeGen1Datastore as RestAzureDatalakeGen1Datastore
-from azure.ai.ml._restclient.v2022_10_01.models import Datastore as DatastoreData, DatastoreType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    AzureDataLakeGen1Datastore as RestAzureDatalakeGen1Datastore,
+)
+from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData, DatastoreType
 from azure.ai.ml._schema._datastore.adls_gen1 import AzureDataLakeGen1Schema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
 from azure.ai.ml.entities._credentials import CertificateConfiguration, ServicePrincipalConfiguration

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/azure_storage.py

@@ -8,10 +8,12 @@
 from typing import Dict, Optional, Union
 
 from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata
-from azure.ai.ml._restclient.v2022_10_01.models import AzureBlobDatastore as RestAzureBlobDatastore
-from azure.ai.ml._restclient.v2022_10_01.models import AzureDataLakeGen2Datastore as RestAzureDataLakeGen2Datastore
-from azure.ai.ml._restclient.v2022_10_01.models import AzureFileDatastore as RestAzureFileDatastore
-from azure.ai.ml._restclient.v2022_10_01.models import Datastore as DatastoreData, DatastoreType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureBlobDatastore as RestAzureBlobDatastore
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    AzureDataLakeGen2Datastore as RestAzureDataLakeGen2Datastore,
+)
+from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureFileDatastore as RestAzureFileDatastore
+from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData, DatastoreType
 from azure.ai.ml._schema._datastore import AzureBlobSchema, AzureDataLakeGen2Schema, AzureFileSchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
 from azure.ai.ml.entities._credentials import (

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/datastore.py

@@ -9,8 +9,7 @@
 from pathlib import Path
 from typing import IO, Any, AnyStr, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2022_10_01.models import Datastore as DatastoreData, DatastoreType
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType as DatastoreTypePreview
+from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData, DatastoreType
 from azure.ai.ml._utils.utils import camel_to_snake, dump_yaml_to_file
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, CommonYamlFields
 from azure.ai.ml.entities._credentials import NoneCredentialConfiguration
@@ -122,7 +121,7 @@ def _load(
             ds_type = AzureDataLakeGen1Datastore
         elif type == camel_to_snake(DatastoreType.AZURE_DATA_LAKE_GEN2):
             ds_type = AzureDataLakeGen2Datastore
-        elif type == camel_to_snake(DatastoreTypePreview.ONE_LAKE):
+        elif type == camel_to_snake(DatastoreType.ONE_LAKE):
             ds_type = OneLakeDatastore
         # disable unless preview release
         # elif type == camel_to_snake(DatastoreTypePreview.HDFS):
@@ -167,7 +166,7 @@ def _from_rest_object(cls, datastore_resource: DatastoreData) -> "Datastore":
             return AzureBlobDatastore._from_rest_object(datastore_resource)
         if datastore_type == DatastoreType.AZURE_FILE:
             return AzureFileDatastore._from_rest_object(datastore_resource)
-        if datastore_type == DatastoreTypePreview.ONE_LAKE:
+        if datastore_type == DatastoreType.ONE_LAKE:
             return OneLakeDatastore._from_rest_object(datastore_resource)
         # disable unless preview release
         # elif datastore_type == DatastoreTypePreview.HDFS:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/one_lake.py

@@ -4,23 +4,54 @@
 
 # pylint: disable=protected-access,no-member
 
+from abc import ABC
 from pathlib import Path
 from typing import Dict, Optional, Union
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import (
     OneLakeDatastore as RestOneLakeDatastore,
     Datastore as DatastoreData,
     DatastoreType,
-    OneLakeArtifact,
+    LakeHouseArtifact as RestLakeHouseArtifact,
 )
 from azure.ai.ml._schema._datastore.one_lake import OneLakeSchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
 from azure.ai.ml.entities._credentials import NoneCredentialConfiguration, ServicePrincipalConfiguration
 from azure.ai.ml.entities._datastore.datastore import Datastore
 from azure.ai.ml.entities._datastore.utils import from_rest_datastore_credentials
+from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin
 from azure.ai.ml.entities._util import load_from_dict
 
 
+class OneLakeArtifact(RestTranslatableMixin, DictMixin, ABC):
+    """OneLake artifact (data source) backing the OneLake workspace.
+
+    :param artifact_name: OneLake artifact name/GUID. ex) 01234567-abcd-1234-5678-012345678901
+    :type artifact_name: str
+    :param artifact_type: OneLake artifact type. Only LakeHouse artifacts are currently supported.
+    :type artifact_type: str
+    """
+
+    def __init__(self, artifact_name: str, artifact_type: Optional[str] = None):
+        super().__init__()
+        self.artifact_name = artifact_name
+        self.artifact_type = artifact_type
+
+
+class LakeHouseArtifact(OneLakeArtifact):
+    """LakeHouse artifact type for OneLake.
+
+    :param artifact_name: OneLake LakeHouse artifact name/GUID. ex) 01234567-abcd-1234-5678-012345678901
+    :type artifact_name: str
+    """
+
+    def __init__(self, artifact_name: str):
+        super(LakeHouseArtifact, self).__init__(artifact_name=artifact_name, artifact_type="lake_house")
+
+    def _to_datastore_rest_object(self) -> RestLakeHouseArtifact:
+        return RestLakeHouseArtifact(artifact_name=self.artifact_name)
+
+
 class OneLakeDatastore(Datastore):
     """OneLake datastore that is linked to an Azure ML workspace.
 
@@ -40,8 +71,7 @@ class OneLakeDatastore(Datastore):
     :type properties: dict[str, str]
     :param credentials: Credentials to use to authenticate against OneLake.
     :type credentials: Union[
-        ~azure.ai.ml.entities.ServicePrincipalConfiguration, ~azure.ai.ml.entities.NoneCredentialConfiguration
-        ]
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration, ~azure.ai.ml.entities.NoneCredentialConfiguration]
     :param kwargs: A dictionary of additional configuration parameters.
     :type kwargs: dict
     """
@@ -70,7 +100,7 @@ def __init__(
     def _to_rest_object(self) -> DatastoreData:
         one_lake_ds = RestOneLakeDatastore(
             credentials=self.credentials._to_datastore_rest_object(),
-            artifact=self.artifact,
+            artifact=RestLakeHouseArtifact(artifact_name=self.artifact["artifact_name"]),
             one_lake_workspace_name=self.one_lake_workspace_name,
             endpoint=self.endpoint,
             description=self.description,
@@ -88,7 +118,7 @@ def _from_rest_object(cls, datastore_resource: DatastoreData):
         return OneLakeDatastore(
             name=datastore_resource.name,
             id=datastore_resource.id,
-            artifact=properties.artifact,
+            artifact=LakeHouseArtifact(artifact_name=properties.artifact.artifact_name),
             one_lake_workspace_name=properties.one_lake_workspace_name,
             endpoint=properties.endpoint,
             credentials=from_rest_datastore_credentials(properties.credentials),
@@ -100,7 +130,8 @@ def __eq__(self, other) -> bool:
         return (
             super().__eq__(other)
             and self.one_lake_workspace_name == other.one_lake_workspace_name
-            and self.artifact == other.artifact
+            and self.artifact.artifact_type == other.artifact["artifact_type"]
+            and self.artifact.artifact_name == other.artifact["artifact_name"]
             and self.endpoint == other.endpoint
         )
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/utils.py

@@ -6,8 +6,7 @@
 
 from typing import Optional
 
-from azure.ai.ml._restclient.v2022_10_01_preview import models as models_preview
-from azure.ai.ml._restclient.v2022_10_01 import models
+from azure.ai.ml._restclient.v2023_04_01_preview import models
 from azure.ai.ml.entities._credentials import (
     AccountKeyConfiguration,
     CertificateConfiguration,
@@ -38,11 +37,11 @@ def from_rest_datastore_credentials(rest_credentials: models.DatastoreCredential
 
 
 def _from_rest_datastore_credentials_preview(
-    rest_credentials: models_preview.DatastoreCredentials,
+    rest_credentials: models.DatastoreCredentials,
 ) -> Optional[BaseKerberosCredentials]:
-    if isinstance(rest_credentials, models_preview.KerberosKeytabCredentials):
+    if isinstance(rest_credentials, models.KerberosKeytabCredentials):
         return KerberosKeytabCredentials._from_rest_object(rest_credentials)
-    if isinstance(rest_credentials, models_preview.KerberosPasswordCredentials):
+    if isinstance(rest_credentials, models.KerberosPasswordCredentials):
         return KerberosPasswordCredentials._from_rest_object(rest_credentials)
 
     return None