two version of R2R are here HEAD master

5 files changed, 552 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/__init__.py
new file mode 100644
index 00000000..fa462cfb
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/__init__.py
@@ -0,0 +1,37 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+
+from .workspace_connection import WorkspaceConnectionSchema
+from .connection_subtypes import (
+    AzureBlobStoreConnectionSchema,
+    MicrosoftOneLakeConnectionSchema,
+    AzureOpenAIConnectionSchema,
+    AzureAIServicesConnectionSchema,
+    AzureAISearchConnectionSchema,
+    AzureContentSafetyConnectionSchema,
+    AzureSpeechServicesConnectionSchema,
+    APIKeyConnectionSchema,
+    OpenAIConnectionSchema,
+    SerpConnectionSchema,
+    ServerlessConnectionSchema,
+    OneLakeArtifactSchema,
+)
+
+__all__ = [
+    "WorkspaceConnectionSchema",
+    "AzureBlobStoreConnectionSchema",
+    "MicrosoftOneLakeConnectionSchema",
+    "AzureOpenAIConnectionSchema",
+    "AzureAIServicesConnectionSchema",
+    "AzureAISearchConnectionSchema",
+    "AzureContentSafetyConnectionSchema",
+    "AzureSpeechServicesConnectionSchema",
+    "APIKeyConnectionSchema",
+    "OpenAIConnectionSchema",
+    "SerpConnectionSchema",
+    "ServerlessConnectionSchema",
+    "OneLakeArtifactSchema",
+]
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/connection_subtypes.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/connection_subtypes.py
new file mode 100644
index 00000000..d04b3e76
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/connection_subtypes.py
@@ -0,0 +1,225 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# 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
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.constants._common import ConnectionTypes
+from azure.ai.ml._schema.workspace.connections.one_lake_artifacts import OneLakeArtifactSchema
+from azure.ai.ml._schema.workspace.connections.credentials import (
+    SasTokenConfigurationSchema,
+    ServicePrincipalConfigurationSchema,
+    AccountKeyConfigurationSchema,
+    AadCredentialConfigurationSchema,
+)
+from azure.ai.ml.entities import AadCredentialConfiguration
+from .workspace_connection import WorkspaceConnectionSchema
+
+
+class AzureBlobStoreConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.AZURE_BLOB, casing_transform=camel_to_snake, required=True
+    )
+    credentials = UnionField(
+        [
+            NestedField(SasTokenConfigurationSchema),
+            NestedField(AccountKeyConfigurationSchema),
+            NestedField(AadCredentialConfigurationSchema),
+        ],
+        required=False,
+        load_default=AadCredentialConfiguration(),
+    )
+
+    url = fields.Str()
+
+    account_name = fields.Str(required=True, allow_none=False)
+    container_name = fields.Str(required=True, allow_none=False)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureBlobStoreConnection
+
+        return AzureBlobStoreConnection(**data)
+
+
+class MicrosoftOneLakeConnectionSchema(WorkspaceConnectionSchema):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.AZURE_ONE_LAKE, casing_transform=camel_to_snake, required=True
+    )
+    credentials = UnionField(
+        [NestedField(ServicePrincipalConfigurationSchema), NestedField(AadCredentialConfigurationSchema)],
+        required=False,
+        load_default=AadCredentialConfiguration(),
+    )
+    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):
+        from azure.ai.ml.entities import MicrosoftOneLakeConnection
+
+        return MicrosoftOneLakeConnection(**data)
+
+
+class AzureOpenAIConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.AZURE_OPEN_AI, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=False, allow_none=True)
+    api_version = fields.Str(required=False, allow_none=True)
+
+    azure_endpoint = fields.Str()
+    open_ai_resource_id = fields.Str(required=False, allow_none=True)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureOpenAIConnection
+
+        return AzureOpenAIConnection(**data)
+
+
+class AzureAIServicesConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionTypes.AZURE_AI_SERVICES, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=False, allow_none=True)
+    endpoint = fields.Str()
+    ai_services_resource_id = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureAIServicesConnection
+
+        return AzureAIServicesConnection(**data)
+
+
+class AzureAISearchConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionTypes.AZURE_SEARCH, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=False, allow_none=True)
+    endpoint = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureAISearchConnection
+
+        return AzureAISearchConnection(**data)
+
+
+class AzureContentSafetyConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionTypes.AZURE_CONTENT_SAFETY, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=False, allow_none=True)
+    endpoint = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureContentSafetyConnection
+
+        return AzureContentSafetyConnection(**data)
+
+
+class AzureSpeechServicesConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionTypes.AZURE_SPEECH_SERVICES, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=False, allow_none=True)
+    endpoint = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import AzureSpeechServicesConnection
+
+        return AzureSpeechServicesConnection(**data)
+
+
+class APIKeyConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.API_KEY, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=True)
+    api_base = fields.Str(required=True)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import APIKeyConnection
+
+        return APIKeyConnection(**data)
+
+
+class OpenAIConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.OPEN_AI, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=True)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import OpenAIConnection
+
+        return OpenAIConnection(**data)
+
+
+class SerpConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(allowed_values=ConnectionCategory.SERP, casing_transform=camel_to_snake, required=True)
+    api_key = fields.Str(required=True)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import SerpConnection
+
+        return SerpConnection(**data)
+
+
+class ServerlessConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.SERVERLESS, casing_transform=camel_to_snake, required=True
+    )
+    api_key = fields.Str(required=True)
+    endpoint = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import ServerlessConnection
+
+        return ServerlessConnection(**data)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/credentials.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/credentials.py
new file mode 100644
index 00000000..52213c08
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/credentials.py
@@ -0,0 +1,178 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument
+
+##### DEV NOTE: For some reason, these schemas correlate to the classes defined in ~azure.ai.ml.entities._credentials.
+# There used to be a credentials.py file in ~azure.ai.ml.entities.workspace.connections,
+# but it was, as far as I could tell, never used. So I removed it and added this comment.
+
+from typing import Dict
+
+from marshmallow import fields, post_load
+
+from azure.ai.ml._restclient.v2024_04_01_preview.models import ConnectionAuthType
+from azure.ai.ml._schema.core.fields import StringTransformedEnum
+from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.entities._credentials import (
+    ManagedIdentityConfiguration,
+    PatTokenConfiguration,
+    SasTokenConfiguration,
+    ServicePrincipalConfiguration,
+    UsernamePasswordConfiguration,
+    AccessKeyConfiguration,
+    ApiKeyConfiguration,
+    AccountKeyConfiguration,
+    AadCredentialConfiguration,
+    NoneCredentialConfiguration,
+)
+
+
+class WorkspaceCredentialsSchema(metaclass=PatchedSchemaMeta):
+    type = fields.Str()
+
+
+class PatTokenConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.PAT,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    pat = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> PatTokenConfiguration:
+        data.pop("type")
+        return PatTokenConfiguration(**data)
+
+
+class SasTokenConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.SAS,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    sas_token = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> SasTokenConfiguration:
+        data.pop("type")
+        return SasTokenConfiguration(**data)
+
+
+class UsernamePasswordConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.USERNAME_PASSWORD,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    username = fields.Str()
+    password = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> UsernamePasswordConfiguration:
+        data.pop("type")
+        return UsernamePasswordConfiguration(**data)
+
+
+class ManagedIdentityConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.MANAGED_IDENTITY,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    client_id = fields.Str()
+    resource_id = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> ManagedIdentityConfiguration:
+        data.pop("type")
+        return ManagedIdentityConfiguration(**data)
+
+
+class ServicePrincipalConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.SERVICE_PRINCIPAL,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+
+    client_id = fields.Str()
+    client_secret = fields.Str()
+    tenant_id = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> ServicePrincipalConfiguration:
+        data.pop("type")
+        return ServicePrincipalConfiguration(**data)
+
+
+class AccessKeyConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.ACCESS_KEY,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    access_key_id = fields.Str()
+    secret_access_key = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> AccessKeyConfiguration:
+        data.pop("type")
+        return AccessKeyConfiguration(**data)
+
+
+class ApiKeyConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.API_KEY,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    key = fields.Str()
+
+    @post_load
+    def make(self, data: Dict[str, str], **kwargs) -> ApiKeyConfiguration:
+        data.pop("type")
+        return ApiKeyConfiguration(**data)
+
+
+class AccountKeyConfigurationSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=ConnectionAuthType.ACCOUNT_KEY,
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+    account_key = fields.Str()
+
+    @post_load
+    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)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/one_lake_artifacts.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/one_lake_artifacts.py
new file mode 100644
index 00000000..563a9359
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/one_lake_artifacts.py
@@ -0,0 +1,26 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument
+
+from marshmallow import fields, post_load
+
+from azure.ai.ml._schema.core.fields import StringTransformedEnum
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.constants._common import OneLakeArtifactTypes
+
+from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
+
+
+class OneLakeArtifactSchema(metaclass=PatchedSchemaMeta):
+    type = StringTransformedEnum(
+        allowed_values=OneLakeArtifactTypes.ONE_LAKE, casing_transform=camel_to_snake, required=True
+    )
+    name = fields.Str(required=True)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import OneLakeConnectionArtifact
+
+        return OneLakeConnectionArtifact(**data)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/workspace_connection.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/workspace_connection.py
new file mode 100644
index 00000000..20863a5a
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/workspace/connections/workspace_connection.py
@@ -0,0 +1,86 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument
+
+from marshmallow import fields, post_load
+
+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
+from azure.ai.ml._schema.workspace.connections.credentials import (
+    AccountKeyConfigurationSchema,
+    ManagedIdentityConfigurationSchema,
+    PatTokenConfigurationSchema,
+    SasTokenConfigurationSchema,
+    ServicePrincipalConfigurationSchema,
+    UsernamePasswordConfigurationSchema,
+    AccessKeyConfigurationSchema,
+    ApiKeyConfigurationSchema,
+    AadCredentialConfigurationSchema,
+    NoneCredentialConfigurationSchema,
+)
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.constants._common import ConnectionTypes
+from azure.ai.ml.entities import NoneCredentialConfiguration, AadCredentialConfiguration
+
+
+class WorkspaceConnectionSchema(ResourceSchema):
+    # Inherits name, id, tags, and description fields from ResourceSchema
+    creation_context = NestedField(CreationContextSchema, dump_only=True)
+    type = StringTransformedEnum(
+        allowed_values=[
+            ConnectionCategory.GIT,
+            ConnectionCategory.CONTAINER_REGISTRY,
+            ConnectionCategory.PYTHON_FEED,
+            ConnectionCategory.S3,
+            ConnectionCategory.SNOWFLAKE,
+            ConnectionCategory.AZURE_SQL_DB,
+            ConnectionCategory.AZURE_SYNAPSE_ANALYTICS,
+            ConnectionCategory.AZURE_MY_SQL_DB,
+            ConnectionCategory.AZURE_POSTGRES_DB,
+            ConnectionTypes.CUSTOM,
+            ConnectionTypes.AZURE_DATA_LAKE_GEN_2,
+        ],
+        casing_transform=camel_to_snake,
+        required=True,
+    )
+
+    # Sorta false, some connection types require this field, some don't.
+    # And some rename it... for client familiarity reasons.
+    target = fields.Str(required=False)
+
+    credentials = UnionField(
+        [
+            NestedField(PatTokenConfigurationSchema),
+            NestedField(SasTokenConfigurationSchema),
+            NestedField(UsernamePasswordConfigurationSchema),
+            NestedField(ManagedIdentityConfigurationSchema),
+            NestedField(ServicePrincipalConfigurationSchema),
+            NestedField(AccessKeyConfigurationSchema),
+            NestedField(ApiKeyConfigurationSchema),
+            NestedField(AccountKeyConfigurationSchema),
+            NestedField(AadCredentialConfigurationSchema),
+            NestedField(NoneCredentialConfigurationSchema),
+        ],
+        required=False,
+        load_default=NoneCredentialConfiguration(),
+    )
+
+    is_shared = fields.Bool(load_default=True)
+    metadata = fields.Dict(required=False)
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import WorkspaceConnection
+
+        # 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()
+        ):
+            data["credentials"] = AadCredentialConfiguration()
+        return WorkspaceConnection(**data)