fix: added subscription selection in scripts, added retry logic for uploading documents to blob

Author: Harsh-Microsoft

infra/scripts/copy_kb_files.sh

@@ -46,12 +46,14 @@ if [ -z "$role_assignment" ]; then
     MSYS_NO_PATHCONV=1 az role assignment create --assignee $signed_user_id --role "Storage Blob Data Contributor" --scope $storage_account_resource_id --output none
     if [ $? -eq 0 ]; then
         echo "Role assignment completed successfully."
+        sleep 5
         retries=3
         while [ $retries -gt 0 ]; do
             # Check if the role assignment was successful
             role_assignment_check=$(MSYS_NO_PATHCONV=1 az role assignment list --assignee $signed_user_id --role "Storage Blob Data Contributor" --scope $storage_account_resource_id --query "[].roleDefinitionId" -o tsv)
             if [ -n "$role_assignment_check" ]; then
                 echo "Role assignment verified successfully."
+                sleep 5
                 break
             else
                 echo "Role assignment not found, retrying..."
@@ -93,6 +95,27 @@ unzip -o $zipUrl1 -d infra/data/"$extractedFolder1"
 # unzip /mnt/azscripts/azscriptinput/"$zipFileName2" -d /mnt/azscripts/azscriptinput/"$extractedFolder2"
 
 # Using az storage blob upload-batch to upload files with managed identity authentication, as the az storage fs directory upload command is not working with managed identity authentication.
-echo "Uploading files to Azure Storage"
+echo "Uploading files to Azure Blob Storage"
 az storage blob upload-batch --account-name "$storageAccount" --destination "$fileSystem"/"$extractedFolder1" --source infra/data/"$extractedFolder1" --auth-mode login --pattern '*' --overwrite --output none
+if [ $? -ne 0 ]; then
+    retries=3
+    sleepTime=10
+    echo "Error: Failed to upload files to Azure Blob Storage. Retrying upload...($((4 - retries)) of 3)"
+    while [ $retries -gt 0 ]; do
+        sleep $sleepTime
+        az storage blob upload-batch --account-name "$storageAccount" --destination "$fileSystem"/"$extractedFolder1" --source infra/data/"$extractedFolder1" --auth-mode login --pattern '*' --overwrite --output none
+        if [ $? -eq 0 ]; then
+            echo "Files uploaded successfully to Azure Blob Storage."
+            break
+        else
+            ((retries--))
+            echo "Retrying upload... ($((4 - retries)) of 3)"
+            sleepTime=$((sleepTime * 2))
+            sleep $sleepTime
+        fi
+    done
+    exit 1
+else
+    echo "Files uploaded successfully to Azure Blob Storage."
+fi
 # az storage blob upload-batch --account-name "$storageAccount" --destination data/"$extractedFolder2" --source /mnt/azscripts/azscriptinput/"$extractedFolder2" --auth-mode login --pattern '*' --overwrite

infra/scripts/process_sample_data.sh

@@ -34,12 +34,72 @@ if [ -z "$aiFoundryName" ]; then
     aiFoundryName=$(azd env get-value AI_FOUNDRY_NAME)
 fi
 
+azSubscriptionId=$(azd env get-value AZURE_SUBSCRIPTION_ID)
+
 # Check if all required arguments are provided
 if [ -z "$storageAccount" ] || [ -z "$fileSystem" ] || [ -z "$keyvaultName" ] || [ -z "$cosmosDbAccountName" ] || [ -z "$resourceGroupName" ] || [ -z "$aiFoundryName" ]; then
     echo "Usage: $0 <storageAccount> <storageContainerName> <keyvaultName> <cosmosDbAccountName> <resourceGroupName> <aiFoundryName>"
     exit 1
 fi
 
+# Authenticate with Azure
+if az account show &> /dev/null; then
+    echo "Already authenticated with Azure."
+else
+    if [ -n "$managedIdentityClientId" ]; then
+        # Use managed identity if running in Azure
+        echo "Authenticating with Managed Identity..."
+        az login --identity --client-id ${managedIdentityClientId}
+    else
+        # Use Azure CLI login if running locally
+        echo "Authenticating with Azure CLI..."
+        az login
+    fi
+    echo "Not authenticated with Azure. Attempting to authenticate..."
+fi
+
+#check if user has selected the correct subscription
+currentSubscriptionId=$(az account show --query id -o tsv)
+currentSubscriptionName=$(az account show --query name -o tsv)
+if [ "$currentSubscriptionId" != "$azSubscriptionId" ]; then
+    echo "Current selected subscription is $currentSubscriptionName ( $currentSubscriptionId )."
+    read -rp "Do you want to continue with this subscription?(y/n): " confirmation
+    if [[ "$confirmation" != "y" && "$confirmation" != "Y" ]]; then
+        echo "Fetching available subscriptions..."
+        availableSubscriptions=$(az account list --query "[?state=='Enabled'].[name,id]" --output tsv)
+        while true; do
+            echo ""
+            echo "Available Subscriptions:"
+            echo "========================"
+            echo "$availableSubscriptions" | awk '{printf "%d. %s ( %s )\n", NR, $1, $2}'
+            echo "========================"
+            echo ""
+            read -rp "Enter the number of the subscription (1-$(echo "$availableSubscriptions" | wc -l)) to use: " subscriptionIndex
+            if [[ "$subscriptionIndex" =~ ^[0-9]+$ ]] && [ "$subscriptionIndex" -ge 1 ] && [ "$subscriptionIndex" -le $(echo "$availableSubscriptions" | wc -l) ]; then
+                selectedSubscription=$(echo "$availableSubscriptions" | sed -n "${subscriptionIndex}p")
+                selectedSubscriptionName=$(echo "$selectedSubscription" | cut -f1)
+                selectedSubscriptionId=$(echo "$selectedSubscription" | cut -f2)
+
+                # Set the selected subscription
+                if  az account set --subscription "$selectedSubscriptionId"; then
+                    echo "Switched to subscription: $selectedSubscriptionName ( $selectedSubscriptionId )"
+                    break
+                else
+                    echo "Failed to switch to subscription: $selectedSubscriptionName ( $selectedSubscriptionId )."
+                fi
+            else
+                echo "Invalid selection. Please try again."
+            fi
+        done
+    else
+        echo "Proceeding with the current subscription: $currentSubscriptionName ( $currentSubscriptionId )"
+        az account set --subscription "$currentSubscriptionId"
+    fi
+else
+    echo "Proceeding with the subscription: $currentSubscriptionName ( $currentSubscriptionId )"
+    az account set --subscription "$currentSubscriptionId"
+fi
+
 # Call add_cosmosdb_access.sh
 echo "Running add_cosmosdb_access.sh"
 bash infra/scripts/add_cosmosdb_access.sh "$resourceGroupName" "$cosmosDbAccountName" "$managedIdentityClientId"