[ml] Fix datalake gen2 datastore type annotation (#34669)

* Remove Optional type annotation from get() methods

* Remove duplicate overloads

* Update type annotations

---------

Co-authored-by: Anna Tisch <[email protected]>

Author: diondrapeck

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

@@ -17,7 +17,12 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models import 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 AccountKeyConfiguration, SasTokenConfiguration
+from azure.ai.ml.entities._credentials import (
+    AccountKeyConfiguration,
+    CertificateConfiguration,
+    SasTokenConfiguration,
+    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._util import load_from_dict
@@ -244,8 +249,11 @@ class AzureDataLakeGen2Datastore(Datastore):
     :param protocol: Protocol to use to connect with the Azure storage account
     :type protocol: str
     :param credentials: Credentials to use for Azure ML workspace to connect to the storage.
-    :type credentials: Union[~azure.ai.ml.entities.AccountKeyConfiguration,
-        ~azure.ai.ml.entities.SasTokenConfiguration]
+    :type credentials: Union[
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration,
+        ~azure.ai.ml.entities.CertificateConfiguration
+
+        ]
     :param properties: The asset property dictionary.
     :type properties: dict[str, str]
     :param kwargs: A dictionary of additional configuration parameters.
@@ -263,7 +271,7 @@ def __init__(
         endpoint: str = _get_storage_endpoint_from_metadata(),
         protocol: str = HTTPS,
         properties: Optional[Dict] = None,
-        credentials: Optional[Union[AccountKeyConfiguration, SasTokenConfiguration]] = None,
+        credentials: Optional[Union[ServicePrincipalConfiguration, CertificateConfiguration]] = None,
         **kwargs: Any
     ):
         kwargs[TYPE] = DatastoreType.AZURE_DATA_LAKE_GEN2

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

@@ -13,7 +13,13 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models import 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
+from azure.ai.ml.entities._credentials import (
+    AccountKeyConfiguration,
+    CertificateConfiguration,
+    NoneCredentialConfiguration,
+    SasTokenConfiguration,
+    ServicePrincipalConfiguration,
+)
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 from azure.ai.ml.entities._resource import Resource
 from azure.ai.ml.entities._util import find_type_in_override
@@ -28,7 +34,14 @@ class Datastore(Resource, RestTranslatableMixin, ABC):
     :param description: Description of the resource.
     :type description: str
     :param credentials: Credentials to use for Azure ML workspace to connect to the storage.
-    :type credentials: Union[ServicePrincipalSection, CertificateSection]
+    :type credentials: Optional[Union[
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration,
+        ~azure.ai.ml.entities.CertificateConfiguration,
+        ~azure.ai.ml.entities.NoneCredentialConfiguration,
+        ~azure.ai.ml.entities.AccountKeyConfiguration,
+        ~azure.ai.ml.entities.SasTokenConfiguration
+
+        ]]
     :param tags: Tag dictionary. Tags can be added, removed, and updated.
     :type tags: dict[str, str]
     :param properties: The asset property dictionary.
@@ -39,7 +52,15 @@ class Datastore(Resource, RestTranslatableMixin, ABC):
 
     def __init__(
         self,
-        credentials: Any,
+        credentials: Optional[
+            Union[
+                ServicePrincipalConfiguration,
+                CertificateConfiguration,
+                NoneCredentialConfiguration,
+                AccountKeyConfiguration,
+                SasTokenConfiguration,
+            ]
+        ],
         name: Optional[str] = None,
         description: Optional[str] = None,
         tags: Optional[Dict] = None,