First

sdk/ai/azure-ai-projects/azure/ai/projects/aio/_patch.py

@@ -6,6 +6,7 @@
 
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
+import os
 from typing import List, Any, TYPE_CHECKING
 from typing_extensions import Self
 from azure.core.credentials_async import AsyncTokenCredential
@@ -17,6 +18,21 @@
     # pylint: disable=unused-import,ungrouped-imports
     from azure.ai.agents.aio import AgentsClient
 
+_console_logging_enabled: bool = os.environ.get("ENABLE_AZURE_AI_PROJECTS_CONSOLE_LOGGING", "False").lower() in (
+    "true",
+    "1",
+    "yes",
+)
+if _console_logging_enabled:
+    import sys
+    import logging
+
+    azure_logger = logging.getLogger("azure")
+    azure_logger.setLevel(logging.DEBUG)
+    azure_logger.addHandler(logging.StreamHandler(stream=sys.stdout))
+    identity_logger = logging.getLogger("azure.identity")
+    identity_logger.setLevel(logging.ERROR)
+
 
 class AIProjectClient(AIProjectClientGenerated):  # pylint: disable=too-many-instance-attributes
     """AIProjectClient.
@@ -55,6 +71,8 @@ class AIProjectClient(AIProjectClientGenerated):  # pylint: disable=too-many-ins
 
     def __init__(self, endpoint: str, credential: AsyncTokenCredential, **kwargs: Any) -> None:
 
+        kwargs.setdefault("logging_enable", _console_logging_enabled)
+
         self._kwargs = kwargs.copy()
         self._patched_user_agent = _patch_user_agent(self._kwargs.pop("user_agent", None))
 

sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_datasets_async.py

@@ -8,7 +8,7 @@
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
 import logging
-from typing import Any, Tuple
+from typing import Any, Tuple, Optional
 from pathlib import Path
 from azure.storage.blob.aio import ContainerClient
 from azure.core.tracing.decorator_async import distributed_trace_async
@@ -43,64 +43,42 @@ async def _create_dataset_and_get_its_container_client(
         self,
         name: str,
         input_version: str,
+        connection_name: Optional[str] = None,
     ) -> Tuple[ContainerClient, str]:
 
         pending_upload_response: PendingUploadResponse = await self.pending_upload(
             name=name,
             version=input_version,
-            body=PendingUploadRequest(pending_upload_type=PendingUploadType.BLOB_REFERENCE),
+            body=PendingUploadRequest(
+                pending_upload_type=PendingUploadType.BLOB_REFERENCE,
+                connection_name=connection_name,
+            ),
         )
         output_version: str = input_version
 
         if not pending_upload_response.blob_reference:
-            raise ValueError("Blob reference for consumption is not present")
-        if not pending_upload_response.blob_reference.credential.type:
-            raise ValueError("Credential type is not present")
-        if pending_upload_response.blob_reference.credential.type != CredentialType.SAS:
-            raise ValueError("Credential type is not SAS")
-        if not pending_upload_response.blob_reference.blob_uri:
-            raise ValueError("Blob URI is not present or empty")
-
-        if logger.getEffectiveLevel() == logging.DEBUG:
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.pending_upload_id = %s.",
-                pending_upload_response.pending_upload_id,
-            )
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.pending_upload_type = %s.",
-                pending_upload_response.pending_upload_type,
-            )  # == PendingUploadType.BLOB_REFERENCE
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.blob_uri = %s.",
-                pending_upload_response.blob_reference.blob_uri,
-            )  # Hosted on behalf of (HOBO) not visible to the user. If the form of: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.storage_account_arm_id = %s.",
-                pending_upload_response.blob_reference.storage_account_arm_id,
-            )  # /subscriptions/<>/resourceGroups/<>/Microsoft.Storage/accounts/<>
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.credential.sas_uri = %s.",
-                pending_upload_response.blob_reference.credential.sas_uri,
-            )
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.credential.type = %s.",
-                pending_upload_response.blob_reference.credential.type,
-            )  # == CredentialType.SAS
+            raise ValueError("Blob reference is not present")
+        if not pending_upload_response.blob_reference.credential:
+            raise ValueError("SAS credential are not present")
+        if not pending_upload_response.blob_reference.credential.sas_uri:
+            raise ValueError("SAS URI is missing or empty")
 
         # For overview on Blob storage SDK in Python see:
         # https://learn.microsoft.com/azure/storage/blobs/storage-quickstart-blobs-python
         # https://learn.microsoft.com/azure/storage/blobs/storage-blob-upload-python
 
         # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.containerclient?view=azure-python#azure-storage-blob-containerclient-from-container-url
         return (
-            await ContainerClient.from_container_url(
-                container_url=pending_upload_response.blob_reference.blob_uri,  # Of the form: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
+            ContainerClient.from_container_url(
+                container_url=pending_upload_response.blob_reference.credential.sas_uri,  # Of the form: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
             ),
             output_version,
         )
 
     @distributed_trace_async
-    async def upload_file(self, *, name: str, version: str, file_path: str, **kwargs: Any) -> DatasetVersion:
+    async def upload_file(
+        self, *, name: str, version: str, file_path: str, connection_name: Optional[str] = None, **kwargs: Any
+    ) -> DatasetVersion:
         """Upload file to a blob storage, and create a dataset that references this file.
         This method uses the `ContainerClient.upload_blob` method from the azure-storage-blob package
         to upload the file. Any keyword arguments provided will be passed to the `upload_blob` method.
@@ -111,6 +89,9 @@ async def upload_file(self, *, name: str, version: str, file_path: str, **kwargs
         :paramtype version: str
         :keyword file_path: The file name (including optional path) to be uploaded. Required.
         :paramtype file_path: str
+        :keyword connection_name: The name of an Azure Storage Account connection, where the file should be uploaded.
+         If not specified, the default Azure Storage Account connection will be used. Optional.
+        :paramtype connection_name: str
         :return: The created dataset version.
         :rtype: ~azure.ai.projects.models.DatasetVersion
         :raises ~azure.core.exceptions.HttpResponseError: If an error occurs during the HTTP request.
@@ -123,7 +104,9 @@ async def upload_file(self, *, name: str, version: str, file_path: str, **kwargs
             raise ValueError("The provided file is actually a folder. Use method `upload_folder` instead")
 
         container_client, output_version = await self._create_dataset_and_get_its_container_client(
-            name=name, input_version=version
+            name=name,
+            input_version=version,
+            connection_name=connection_name,
         )
 
         async with container_client:
@@ -138,18 +121,21 @@ async def upload_file(self, *, name: str, version: str, file_path: str, **kwargs
                 )
 
                 # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.containerclient?view=azure-python#azure-storage-blob-containerclient-upload-blob
-                with await container_client.upload_blob(name=blob_name, data=data, **kwargs) as blob_client:
+                async with await container_client.upload_blob(name=blob_name, data=data, **kwargs) as blob_client:
 
                     logger.debug("[upload_file] Done uploading")
 
+                    file_dataset_version = FileDatasetVersion(
+                        # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.blobclient?view=azure-python#azure-storage-blob-blobclient-url
+                        # Per above doc the ".url" contains SAS token... should this be stripped away?
+                        data_uri=blob_client.url,  # "<account>.blob.windows.core.net/<container>/<file_name>"
+                    )
+                    file_dataset_version.is_reference = True  # TODO: Update TypeSpec to make this writable.
+
                     dataset_version = await self.create_or_update(
                         name=name,
                         version=output_version,
-                        body=FileDatasetVersion(
-                            # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.blobclient?view=azure-python#azure-storage-blob-blobclient-url
-                            # Per above doc the ".url" contains SAS token... should this be stripped away?
-                            data_uri=blob_client.url,  # "<account>.blob.windows.core.net/<container>/<file_name>"
-                        ),
+                        body=file_dataset_version,
                     )
 
         return dataset_version

sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_datasets.py

@@ -8,7 +8,7 @@
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
 import logging
-from typing import Any, Tuple
+from typing import Any, Tuple, Optional
 from pathlib import Path
 from azure.storage.blob import ContainerClient
 from azure.core.tracing.decorator import distributed_trace
@@ -42,49 +42,25 @@ def _create_dataset_and_get_its_container_client(
         self,
         name: str,
         input_version: str,
+        connection_name: Optional[str] = None,
     ) -> Tuple[ContainerClient, str]:
 
         pending_upload_response: PendingUploadResponse = self.pending_upload(
             name=name,
             version=input_version,
-            body=PendingUploadRequest(pending_upload_type=PendingUploadType.BLOB_REFERENCE),
+            body=PendingUploadRequest(
+                pending_upload_type=PendingUploadType.BLOB_REFERENCE,
+                connection_name=connection_name,
+            ),
         )
         output_version: str = input_version
 
         if not pending_upload_response.blob_reference:
-            raise ValueError("Blob reference for consumption is not present")
-        if not pending_upload_response.blob_reference.credential.type:
-            raise ValueError("Credential type is not present")
-        if pending_upload_response.blob_reference.credential.type != CredentialType.SAS:
-            raise ValueError("Credential type is not SAS")
-        if not pending_upload_response.blob_reference.blob_uri:
-            raise ValueError("Blob URI is not present or empty")
-
-        if logger.getEffectiveLevel() == logging.DEBUG:
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.pending_upload_id = %s.",
-                pending_upload_response.pending_upload_id,
-            )
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.pending_upload_type = %s.",
-                pending_upload_response.pending_upload_type,
-            )  # == PendingUploadType.BLOB_REFERENCE
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.blob_uri = %s.",
-                pending_upload_response.blob_reference.blob_uri,
-            )  # Hosted on behalf of (HOBO) not visible to the user. If the form of: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.storage_account_arm_id = %s.",
-                pending_upload_response.blob_reference.storage_account_arm_id,
-            )  # /subscriptions/<>/resourceGroups/<>/Microsoft.Storage/accounts/<>
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.credential.sas_uri = %s.",
-                pending_upload_response.blob_reference.credential.sas_uri,
-            )
-            logger.debug(
-                "[_create_dataset_and_get_its_container_client] pending_upload_response.blob_reference.credential.type = %s.",
-                pending_upload_response.blob_reference.credential.type,
-            )  # == CredentialType.SAS
+            raise ValueError("Blob reference is not present")
+        if not pending_upload_response.blob_reference.credential:
+            raise ValueError("SAS credential are not present")
+        if not pending_upload_response.blob_reference.credential.sas_uri:
+            raise ValueError("SAS URI is missing or empty")
 
         # For overview on Blob storage SDK in Python see:
         # https://learn.microsoft.com/azure/storage/blobs/storage-quickstart-blobs-python
@@ -93,13 +69,15 @@ def _create_dataset_and_get_its_container_client(
         # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.containerclient?view=azure-python#azure-storage-blob-containerclient-from-container-url
         return (
             ContainerClient.from_container_url(
-                container_url=pending_upload_response.blob_reference.blob_uri,  # Of the form: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
+                container_url=pending_upload_response.blob_reference.credential.sas_uri  # Of the form: "https://<account>.blob.core.windows.net/<container>?<sasToken>"
             ),
             output_version,
         )
 
     @distributed_trace
-    def upload_file(self, *, name: str, version: str, file_path: str, **kwargs: Any) -> DatasetVersion:
+    def upload_file(
+        self, *, name: str, version: str, file_path: str, connection_name: Optional[str] = None, **kwargs: Any
+    ) -> DatasetVersion:
         """Upload file to a blob storage, and create a dataset that references this file.
         This method uses the `ContainerClient.upload_blob` method from the azure-storage-blob package
         to upload the file. Any keyword arguments provided will be passed to the `upload_blob` method.
@@ -110,6 +88,9 @@ def upload_file(self, *, name: str, version: str, file_path: str, **kwargs: Any)
         :paramtype version: str
         :keyword file_path: The file name (including optional path) to be uploaded. Required.
         :paramtype file_path: str
+        :keyword connection_name: The name of an Azure Storage Account connection, where the file should be uploaded.
+         If not specified, the default Azure Storage Account connection will be used. Optional.
+        :paramtype connection_name: str
         :return: The created dataset version.
         :rtype: ~azure.ai.projects.models.DatasetVersion
         :raises ~azure.core.exceptions.HttpResponseError: If an error occurs during the HTTP request.
@@ -122,7 +103,7 @@ def upload_file(self, *, name: str, version: str, file_path: str, **kwargs: Any)
             raise ValueError("The provided file is actually a folder. Use method `upload_folder` instead")
 
         container_client, output_version = self._create_dataset_and_get_its_container_client(
-            name=name, input_version=version
+            name=name, input_version=version, connection_name=connection_name
         )
 
         with container_client:
@@ -141,14 +122,17 @@ def upload_file(self, *, name: str, version: str, file_path: str, **kwargs: Any)
 
                     logger.debug("[upload_file] Done uploading")
 
+                    file_dataset_version = FileDatasetVersion(
+                        # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.blobclient?view=azure-python#azure-storage-blob-blobclient-url
+                        # Per above doc the ".url" contains SAS token... should this be stripped away?
+                        data_uri=blob_client.url,  # "<account>.blob.windows.core.net/<container>/<file_name>"
+                    )
+                    file_dataset_version.is_reference = True  # TODO: Update TypeSpec to make this writable.
+
                     dataset_version = self.create_or_update(
                         name=name,
                         version=output_version,
-                        body=FileDatasetVersion(
-                            # See https://learn.microsoft.com/python/api/azure-storage-blob/azure.storage.blob.blobclient?view=azure-python#azure-storage-blob-blobclient-url
-                            # Per above doc the ".url" contains SAS token... should this be stripped away?
-                            data_uri=blob_client.url,  # "<account>.blob.windows.core.net/<container>/<file_name>"
-                        ),
+                        body=file_dataset_version,
                     )
 
         return dataset_version