fix to_dict (#35501)

Author: MilesHolland

sdk/ml/azure-ai-ml/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_ce8aa03671"
+  "Tag": "python/ml/azure-ai-ml_bcde27db64"
 }

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

@@ -17,6 +17,7 @@
     OpenAIConnectionSchema,
     SerpConnectionSchema,
     ServerlessConnectionSchema,
+    OneLakeArtifactSchema,
 )
 
 __all__ = [
@@ -32,4 +33,5 @@
     "OpenAIConnectionSchema",
     "SerpConnectionSchema",
     "ServerlessConnectionSchema",
+    "OneLakeArtifactSchema",
 ]

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

@@ -6,7 +6,7 @@
 
 from marshmallow import fields, post_load
 
-from azure.ai.ml._restclient.v2023_06_01_preview.models import ConnectionCategory
+from azure.ai.ml._restclient.v2024_04_01_preview.models import ConnectionCategory
 from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField
 from azure.ai.ml._schema.core.resource import ResourceSchema
 from azure.ai.ml._schema.job import CreationContextSchema
@@ -19,6 +19,8 @@
     UsernamePasswordConfigurationSchema,
     AccessKeyConfigurationSchema,
     ApiKeyConfigurationSchema,
+    AadCredentialConfigurationSchema,
+    NoneCredentialConfigurationSchema,
 )
 from azure.ai.ml._utils.utils import camel_to_snake
 from azure.ai.ml.constants._common import ConnectionTypes
@@ -60,6 +62,8 @@ class ConnectionSchema(ResourceSchema):
             NestedField(AccessKeyConfigurationSchema),
             NestedField(ApiKeyConfigurationSchema),
             NestedField(AccountKeyConfigurationSchema),
+            NestedField(AadCredentialConfigurationSchema),
+            NestedField(NoneCredentialConfigurationSchema),
         ],
         required=False,
         load_default=NoneCredentialConfiguration(),
@@ -71,7 +75,8 @@ class ConnectionSchema(ResourceSchema):
     def make(self, data, **kwargs):
         from azure.ai.ml.entities import Connection
 
-        # Replace ALDS gen 2 empty default with AAD over None
+        # Most non-subclassed connections default to a none credential if none
+        # is provided. ALDS Gen 2 connections default to AAD with this code.
         if (
             data.get("type", None) == ConnectionTypes.AZURE_DATA_LAKE_GEN_2
             and data.get("credentials", None) == NoneCredentialConfiguration()

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

@@ -5,6 +5,8 @@
 # pylint: disable=unused-argument
 
 from marshmallow import fields, post_load
+from marshmallow.exceptions import ValidationError
+from marshmallow.decorators import pre_load
 
 from azure.ai.ml._restclient.v2024_04_01_preview.models import ConnectionCategory
 from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField
@@ -15,6 +17,7 @@
     SasTokenConfigurationSchema,
     ServicePrincipalConfigurationSchema,
     AccountKeyConfigurationSchema,
+    AadCredentialConfigurationSchema,
 )
 from azure.ai.ml.entities import AadCredentialConfiguration
 from .connection import ConnectionSchema
@@ -30,6 +33,7 @@ class AzureBlobStoreConnectionSchema(ConnectionSchema):
         [
             NestedField(SasTokenConfigurationSchema),
             NestedField(AccountKeyConfigurationSchema),
+            NestedField(AadCredentialConfigurationSchema),
         ],
         required=False,
         load_default=AadCredentialConfiguration(),
@@ -52,13 +56,34 @@ class MicrosoftOneLakeConnectionSchema(ConnectionSchema):
     type = StringTransformedEnum(
         allowed_values=ConnectionCategory.AZURE_ONE_LAKE, casing_transform=camel_to_snake, required=True
     )
-    credentials = NestedField(
-        ServicePrincipalConfigurationSchema, required=False, load_default=AadCredentialConfiguration()
+    credentials = UnionField(
+        [NestedField(ServicePrincipalConfigurationSchema), NestedField(AadCredentialConfigurationSchema)],
+        required=False,
+        load_default=AadCredentialConfiguration(),
     )
-    artifact = NestedField(OneLakeArtifactSchema, required=True)
-
-    endpoint = fields.Str()
-    one_lake_workspace_name = fields.Str()
+    artifact = NestedField(OneLakeArtifactSchema, required=False, allow_none=True)
+
+    endpoint = fields.Str(required=False)
+    one_lake_workspace_name = fields.Str(required=False)
+
+    @pre_load
+    def check_for_target(self, data, **kwargs):
+        target = data.get("target", None)
+        artifact = data.get("artifact", None)
+        endpoint = data.get("endpoint", None)
+        one_lake_workspace_name = data.get("one_lake_workspace_name", None)
+        # If the user is using a target, then they don't need the artifact and one lake workspace name.
+        # This is distinct from when the user set's the 'endpoint' value, which is also used to construct
+        # the target. If the target is already present, then the loaded connection YAML was probably produced
+        # by dumping an extant connection.
+        if target is None:
+            if artifact is None:
+                raise ValidationError("If target is unset, then artifact must be set")
+            if endpoint is None:
+                raise ValidationError("If target is unset, then endpoint must be set")
+            if one_lake_workspace_name is None:
+                raise ValidationError("If target is unset, then one_lake_workspace_name must be set")
+        return data
 
     @post_load
     def make(self, data, **kwargs):

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

@@ -25,6 +25,8 @@
     AccessKeyConfiguration,
     ApiKeyConfiguration,
     AccountKeyConfiguration,
+    AadCredentialConfiguration,
+    NoneCredentialConfiguration,
 )
 
 
@@ -148,3 +150,29 @@ class AccountKeyConfigurationSchema(metaclass=PatchedSchemaMeta):
     def make(self, data: Dict[str, str], **kwargs) -> AccountKeyConfiguration:
         data.pop("type")
         return AccountKeyConfiguration(**data)
+
+
+class AadCredentialConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.AAD,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> AadCredentialConfiguration:
+        data.pop("type")
+        return AadCredentialConfiguration(**data)
+
+
+class NoneCredentialConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.NONE,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> NoneCredentialConfiguration:
+        data.pop("type")
+        return NoneCredentialConfiguration(**data)

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

@@ -162,10 +162,10 @@ class MicrosoftOneLakeConnection(Connection):
     :param endpoint: The endpoint of the connection.
     :type endpoint: str
     :param artifact: The artifact class used to further specify the connection.
-    :type artifact: ~azure.ai.ml.entities.OneLakeArtifact
+    :type artifact: Optional[~azure.ai.ml.entities.OneLakeArtifact]
     :param one_lake_workspace_name: The name, not ID, of the workspace where the One Lake
         resource lives.
-    :type one_lake_workspace_name: str
+    :type one_lake_workspace_name: Optional[str]
     :param credentials: The credentials for authenticating to the blob store. This type of
         connection accepts 3 types of credentials: account key and SAS token credentials,
         or NoneCredentialConfiguration for credential-less connections.
@@ -182,8 +182,8 @@ def __init__(
         self,
         *,
         endpoint: str,
-        artifact: OneLakeConnectionArtifact,
-        one_lake_workspace_name: str,
+        artifact: Optional[OneLakeConnectionArtifact] = None,
+        one_lake_workspace_name: Optional[str] = None,
         **kwargs,
     ):
         kwargs.pop("type", None)  # make sure we never somehow use wrong type
@@ -192,6 +192,12 @@ def __init__(
         # need to worry about data-availability nonsense.
         target = kwargs.pop("target", None)
         if target is None:
+            if artifact is None:
+                raise ValueError("If target is unset, then artifact must be set")
+            if endpoint is None:
+                raise ValueError("If target is unset, then endpoint must be set")
+            if one_lake_workspace_name is None:
+                raise ValueError("If target is unset, then one_lake_workspace_name must be set")
             target = MicrosoftOneLakeConnection._construct_target(endpoint, one_lake_workspace_name, artifact)
         super().__init__(
             target=target,

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

@@ -3,7 +3,6 @@
 # ---------------------------------------------------------
 
 # pylint: disable=protected-access
-
 from typing import Any
 from azure.ai.ml._utils._experimental import experimental
 

sdk/ml/azure-ai-ml/tests/connection/e2etests/test_connections.py

@@ -28,6 +28,7 @@
     OpenAIConnection,
     SerpConnection,
     ServerlessConnection,
+    AccountKeyConfiguration,
 )
 from azure.ai.ml.constants._common import ConnectionTypes
 from azure.core.exceptions import ResourceNotFoundError
@@ -398,26 +399,45 @@ def test_workspace_connection_data_connection_listing(
             {"account_name": storage_account_name},
         ]
         internal_blob_ds = load_datastore(blob_store_file, params_override=params_override)
-        created_datastore = self.datastore_create_get_list(client, internal_blob_ds, random_name)
-        assert isinstance(created_datastore, AzureBlobDatastore)
-        assert created_datastore.container_name == internal_blob_ds.container_name
-        assert created_datastore.account_name == internal_blob_ds.account_name
-        assert created_datastore.credentials.account_key == primary_account_key
-
-        # Make sure that normal list call doesn't include data connection
-        assert internal_blob_ds.name not in [conn.name for conn in client.connections.list()]
-
-        # Make sure that the data connection list call includes the data connection
-        found_datastore_conn = False
-        for conn in client.connections.list(include_data_connections=True):
-            if created_datastore.name == conn.name:
-                assert conn.type == camel_to_snake(ConnectionCategory.AZURE_BLOB)
-                assert isinstance(conn, AzureBlobStoreConnection)
-                found_datastore_conn = True
-        # Ensure that we actually found and validated the data connection.
-        assert found_datastore_conn
-        # delete the data store.
-        client.datastores.delete(random_name)
+
+        created_datastore = None
+        created_connection = None
+        try:
+            created_datastore = self.datastore_create_get_list(client, internal_blob_ds, random_name)
+            assert isinstance(created_datastore, AzureBlobDatastore)
+            assert created_datastore.container_name == internal_blob_ds.container_name
+            assert created_datastore.account_name == internal_blob_ds.account_name
+            assert created_datastore.credentials.account_key == primary_account_key
+
+            # Now that a datastore exists, create a connection to it
+            local_connection = AzureBlobStoreConnection(
+                name=created_datastore.name,
+                url=created_datastore.base_path,
+                account_name=created_datastore.account_name,
+                container_name=created_datastore.container_name,
+                credentials=AccountKeyConfiguration(account_key=created_datastore.credentials.account_key),
+            )
+
+            created_connection = client.connections.create_or_update(connection=local_connection)
+
+            # Make sure that normal list call doesn't include data connection
+            assert internal_blob_ds.name not in [conn.name for conn in client.connections.list()]
+
+            # Make sure that the data connection list call includes the data connection
+            found_datastore_conn = False
+            for conn in client.connections.list(include_data_connections=True):
+                if created_datastore.name == conn.name:
+                    assert conn.type == camel_to_snake(ConnectionCategory.AZURE_BLOB)
+                    assert isinstance(conn, AzureBlobStoreConnection)
+                    found_datastore_conn = True
+            # Ensure that we actually found and validated the data connection.
+            assert found_datastore_conn
+        finally:
+            # Delete resources
+            if created_connection is not None:
+                client.connections.delete(name=created_datastore.name)
+            if created_datastore is not None:
+                client.datastores.delete(random_name)
         with pytest.raises(Exception):
             client.datastores.get(random_name)
 
@@ -549,7 +569,7 @@ def test_azure_open_ai_crud(
         assert created_connection.type == camel_to_snake(ConnectionCategory.AZURE_OPEN_AI)
         assert created_connection.tags is not None
         assert created_connection.tags["hello"] == "world"
-        assert created_connection.api_version == "1.0"
+        assert created_connection.api_version is None
         assert created_connection.open_ai_resource_id == None
 
         with pytest.raises(Exception):