azure blob base changes

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

@@ -13,11 +13,13 @@
     AzureAIServiceConnection,
     GitHubConnection,
     CustomConnection,
+    AzureBlobStoreConnection,
 )
 from .mlindex import Index
 from .project import Project
 from .data import Data
 from .configs import AzureOpenAIModelConfiguration
+from azure.ai.ml.entities import ApiKeyConfiguration
 
 __all__ = [
     "BaseConnection",
@@ -31,5 +33,7 @@
     "AzureOpenAIModelConfiguration",
     "GitHubConnection",
     "CustomConnection",
+    "AzureBlobStoreConnection",
+    "ApiKeyConfiguration"
 ]
 

sdk/ai/azure-ai-resources/azure/ai/resources/entities/base_connection.py

@@ -10,8 +10,9 @@
 from azure.ai.ml.entities._credentials import ApiKeyConfiguration
 from azure.core.credentials import TokenCredential
 from azure.ai.ml._restclient.v2023_06_01_preview.models import ConnectionCategory
-
-
+from azure.ai.ml.constants._common import (
+    WorkspaceConnectionTypes,
+)
 class BaseConnection:
     """A connection to a specific external AI system. This is a base class and should not be
     instantiated directly. Use the child classes that are specialized for connections to different services.
@@ -46,7 +47,8 @@ def __init__(
         **kwargs,
     ):
         # Sneaky short-circuit to allow direct v2 WS conn injection without any
-        # polymorphic required argument hassles.
+        # polymorphic-argument hassles.
+        # See _from_v2_workspace_connection for usage.
         if kwargs.pop("make_empty", False):
             return
 
@@ -87,6 +89,9 @@ def _get_ai_connection_class_from_type(cls,  conn_type: str):
             AzureOpenAIConnection,
             AzureAISearchConnection,
             AzureAIServiceConnection,
+            GitHubConnection,
+            CustomConnection,
+            AzureBlobStoreConnection
         )
 
         cat = camel_to_snake(conn_type).lower()
@@ -96,6 +101,15 @@ def _get_ai_connection_class_from_type(cls,  conn_type: str):
             return AzureAISearchConnection
         elif cat == camel_to_snake(ConnectionCategory.COGNITIVE_SERVICE).lower():
             return AzureAIServiceConnection
+        elif cat == camel_to_snake(ConnectionCategory.GIT).lower():
+            return GitHubConnection
+        elif cat == camel_to_snake("azure_blob").lower():
+            return AzureBlobStoreConnection
+        elif (cat == camel_to_snake(ConnectionCategory.CUSTOM_KEYS).lower()
+            or cat == camel_to_snake(WorkspaceConnectionTypes.CUSTOM).lower()):
+            # Accept both custom_keys and custom as conversion type in case of
+            # improper input.
+            return CustomConnection
         return BaseConnection
 
 

sdk/ai/azure-ai-resources/azure/ai/resources/entities/connection_subtypes.py

@@ -25,8 +25,6 @@ class AzureOpenAIConnection(BaseConnection):
     :type name: str
     :param target: The URL or ARM resource ID of the external resource.
     :type target: str
-    :param tags: Tag dictionary. Tags can be added, removed, and updated.
-    :type tags: dict
     :param credentials: The credentials for authenticating the external resource.
     :type credentials: ~azure.ai.ml.entities.ApiKeyConfiguration
     :param api_version: The api version that this connection was created for.
@@ -36,6 +34,8 @@ class AzureOpenAIConnection(BaseConnection):
     :param is_shared: For connections created for a project, this determines if the connection
         is shared amongst other connections with that project's parent AI resource. Defaults to True.
     :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
     """
 
     def __init__(
@@ -66,8 +66,8 @@ def api_version(self) -> Optional[str]:
     def api_version(self, value: str) -> None:
         """Set the API version of the connection.
 
-        :return: the API version of the connection.
-        :rtype: str
+        :param value: the API version of the connection.
+        :type value: str
         """
         self._workspace_connection.tags[CONNECTION_API_VERSION_KEY] = value
 
@@ -86,8 +86,8 @@ def api_type(self) -> Optional[str]:
     def api_type(self, value: str) -> None:
         """Set the API type of the connection.
 
-        :return: the API type of the connection.
-        :rtype: str
+        :param value: the API type of the connection.
+        :type value: str
         """
         self._workspace_connection.tags[CONNECTION_API_TYPE_KEY] = value
 
@@ -154,15 +154,15 @@ class AzureAISearchConnection(BaseConnection):
     :type name: str
     :param target: The URL or ARM resource ID of the external resource.
     :type target: str
-    :param tags: Tag dictionary. Tags can be added, removed, and updated.
-    :type tags: dict
     :param credentials: The credentials for authenticating the external resource.
     :type credentials: ~azure.ai.ml.entities.ApiKeyConfiguration
     :param api_version: The api version that this connection was created for. Only applies to certain connection types.
     :type api_version: Optional[str]
     :param is_shared: For connections created for a project, this determines if the connection
         is shared amongst other connections with that project's parent AI resource. Defaults to True.
     :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
     """
 
     def __init__(
@@ -192,8 +192,8 @@ def api_version(self) -> Optional[str]:
     def api_version(self, value: str) -> None:
         """Set the API version of the connection.
 
-        :return: the API version of the connection.
-        :rtype: str
+        :param value: the API version of the connection.
+        :type value: str
         """
         self._workspace_connection.tags[CONNECTION_API_VERSION_KEY] = value
 
@@ -224,8 +224,6 @@ class AzureAIServiceConnection(BaseConnection):
     :type name: str
     :param target: The URL or ARM resource ID of the external resource.
     :type target: str
-    :param tags: Tag dictionary. Tags can be added, removed, and updated.
-    :type tags: dict
     :param credentials: The credentials for authenticating the external resource.
     :type credentials: ~azure.ai.ml.entities.ApiKeyConfiguration
     :param api_version: The api version that this connection was created for.
@@ -236,6 +234,8 @@ class AzureAIServiceConnection(BaseConnection):
     :param is_shared: For connections created for a project, this determines if the connection
         is shared amongst other connections with that project's parent AI resource. Defaults to True.
     :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
     """
 
     def __init__(
@@ -265,8 +265,8 @@ def api_version(self) -> Optional[str]:
     def api_version(self, value: str) -> None:
         """Set the API version of the connection.
 
-        :return: the API version of the connection.
-        :rtype: str
+        :param value: the API version of the connection.
+        :type value: str
         """
         self._workspace_connection.tags[CONNECTION_API_VERSION_KEY] = value
 
@@ -285,8 +285,8 @@ def kind(self) -> Optional[str]:
     def kind(self, value: str) -> None:
         """Set the kind of the connection.
 
-        :return: the kind of the connection.
-        :rtype: str
+        :param value: the kind of the connection.
+        :type value: str
         """
         self._workspace_connection.tags[CONNECTION_KIND_KEY] = value
 
@@ -297,13 +297,13 @@ class GitHubConnection(BaseConnection):
     :type name: str
     :param target: The URL or ARM resource ID of the external resource.
     :type target: str
-    :param tags: Tag dictionary. Tags can be added, removed, and updated.
-    :type tags: dict
     :param credentials: The credentials for authenticating the external resource.
     :type credentials: ~azure.ai.ml.entities.ApiKeyConfiguration
     :param is_shared: For connections created for a project, this determines if the connection
         is shared amongst other connections with that project's parent AI resource. Defaults to True.
     :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
     """
 
     def __init__(
@@ -316,20 +316,91 @@ def __init__(
         kwargs.pop("type", None)  # make sure we never somehow use wrong type
         super().__init__(target=target, type="git", credentials=credentials, **kwargs)
 
+class AzureBlobStoreConnection(BaseConnection):
+    """A Connection to an Azure Blob Datastore.
+
+    :param name: Name of the connection.
+    :type name: str
+    :param target: The URL or ARM resource ID of the resource.
+    :type target: str
+    :param credentials: ?????
+    :type credentials: ?????
+    :param is_shared: For connections created for a project, this determines if the connection
+        is shared amongst other connections with that project's parent AI resource. Defaults to True.
+    :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
+    """
+
+    def __init__(
+        self,
+        *,
+        target: str,
+        credentials = None,
+        container_name: str,
+        account_name: str,
+        **kwargs,
+    ) -> None:
+        kwargs.pop("type", None)  # make sure we never somehow use wrong type
+        super().__init__(
+            target=target,
+            type="azure_blob",
+            credentials=credentials,
+            container_name=container_name,
+            account_name=account_name,
+            **kwargs
+        )
+
+    @property
+    def container_name(self) -> str:
+        """The container name of the connection.
+
+        :return: the container name of the connection.
+        :rtype: str
+        """
+        return self._workspace_connection.container_name
+
+    @container_name.setter
+    def container_name(self, value: str) -> None:
+        """Set the container name of the connection.
+
+        :param value: the new container name of the connection.
+        :type value: str
+        """
+        self._workspace_connection.container_name = value
+
+    @property
+    def account_name(self) -> str:
+        """The account name of the connection.
+
+        :return: the account name of the connection.
+        :rtype: str
+        """
+        return self._workspace_connection.account_name
+
+    @account_name.setter
+    def account_name(self, value: str) -> None:
+        """Set the account name of the connection.
+
+        :param value: the new account name of the connection.
+        :type value: str
+        """
+        self._workspace_connection.account_name = value
+
 class CustomConnection(BaseConnection):
     """A Connection to system that's not encapsulated by other connection types.
 
     :param name: Name of the connection.
     :type name: str
     :param target: The URL or ARM resource ID of the external resource.
     :type target: str
-    :param tags: Tag dictionary. Tags can be added, removed, and updated.
-    :type tags: dict
     :param credentials: The credentials for authenticating the external resource.
     :type credentials: ~azure.ai.ml.entities.ApiKeyConfiguration
     :param is_shared: For connections created for a project, this determines if the connection
         is shared amongst other connections with that project's parent AI resource. Defaults to True.
     :type is_shared: bool
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict
     """
     def __init__(
         self,

sdk/ai/azure-ai-resources/azure/ai/resources/operations/_connection_operations.py

@@ -42,7 +42,7 @@ def list(self, connection_type: Optional[str] = None, scope: str = OperationScop
         client = self._resource_ml_client if scope == OperationScope.AI_RESOURCE else self._project_ml_client
         return [
             BaseConnection._from_v2_workspace_connection(conn)
-            for conn in client._workspace_connections.list(connection_type=connection_type)
+            for conn in client._workspace_connections.list(connection_type=connection_type, params={"includeAll": "true"})
         ]
 
     @distributed_trace
@@ -58,11 +58,11 @@ def get(self, name: str, scope: str = OperationScope.AI_RESOURCE, **kwargs) -> B
         """
         client = self._resource_ml_client if scope == OperationScope.AI_RESOURCE else self._project_ml_client
         workspace_connection = client._workspace_connections.get(name=name, **kwargs)
-        connection = BaseConnection._from_v2_workspace_connection(workspace_connection)
+        connection = BaseConnection._from_v2_workspace_connection(workspace_connection, params={"includeAll": "true"})
 
         # It's by design that both API and V2 SDK don't include the secrets from API response, the following
         # code fills the gap when possible
-        if not connection.credentials.key:
+        if not (connection.credentials and connection.credentials.key):
             list_secrets_response = client.connections._operation.list_secrets(
                 connection_name=name,
                 resource_group_name=client.resource_group_name,

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

@@ -62,3 +62,22 @@ def make(self, data, **kwargs):
         from azure.ai.ml.entities import AzureAIServiceWorkspaceConnection
 
         return AzureAIServiceWorkspaceConnection(**data)
+
+
+# pylint: disable-next=name-too-long
+class AzureBlobStoreWorkspaceConnectionSchema(WorkspaceConnectionSchema):
+    # type and credentials limited
+    type = StringTransformedEnum(
+        allowed_values=ConnectionCategory.COGNITIVE_SERVICE, casing_transform=camel_to_snake, required=True
+    )
+    credentials = NestedField(ApiKeyConfigurationSchema)
+
+    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 AzureBlobStoreWorkspaceConnection
+
+        return AzureBlobStoreWorkspaceConnection(**data)

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -186,6 +186,8 @@
 CONNECTION_API_VERSION_KEY = "ApiVersion"
 CONNECTION_API_TYPE_KEY = "ApiType"
 CONNECTION_KIND_KEY = "Kind"
+CONNECTION_CONTAINER_NAME_KEY = "ContainerName"
+CONNECTION_ACCOUNT_NAME_KEY = "AccountName"
 
 
 class DefaultOpenEncoding:
@@ -793,7 +795,7 @@ class InputOutputModes:
 class WorkspaceConnectionTypes:
     """Names for workspace connection types that are different from that underlying api enum values."""
 
-    CUSTOM = "custom"
+    CUSTOM = "custom" # Corresponds to "custom_keys".
 
 
 class LegacyAssetTypes: