fix: update post-deployment scripts to use dynamic admin command and improve resource group input handling

Author: Abdul-Microsoft

scripts/post_deployment_setup.ps1

@@ -27,9 +27,14 @@ Write-Host " Post-Deployment Setup"
 Write-Host " Resource Group: $ResourceGroupName"
 Write-Host "=============================================="
 
-# Remove rdbms-connect extension if present (it conflicts with built-in microsoft-entra-admin commands)
+# Remove rdbms-connect extension if present (it conflicts with built-in admin commands)
 az extension remove --name rdbms-connect 2>$null | Out-Null
 
+# Detect whether to use 'microsoft-entra-admin' (newer CLI) or 'ad-admin' (older CLI)
+$pgAdminCmd = "ad-admin"
+$entraCheck = az postgres flexible-server microsoft-entra-admin --help 2>$null
+if ($LASTEXITCODE -eq 0) { $pgAdminCmd = "microsoft-entra-admin" }
+
 # Track resources that need public access restored to Disabled
 $resourcesToRestore = @()
 
@@ -80,6 +85,8 @@ function Restore-NetworkAccess {
 # -------------------------------------------------------
 # STEP 1 — Set Function App Client Key
 # -------------------------------------------------------
+try {
+
 Write-Host ""
 Write-Host "--- Step 1: Set Function App Client Key ---"
 
@@ -145,7 +152,6 @@ else {
         Write-Host "✓ Retrieving function key from Key Vault..."
         $functionKey = az keyvault secret show --vault-name $keyVaultName --name "FUNCTION-KEY" --query "value" -o tsv
         if ($LASTEXITCODE -ne 0 -or -not $functionKey) {
-            Restore-NetworkAccess
             Write-Error "✗ Failed to retrieve 'FUNCTION-KEY' secret from Key Vault '$keyVaultName'."
             exit 1
         }
@@ -171,7 +177,6 @@ else {
                 return ($LASTEXITCODE -eq 0)
             }
             if (-not $keySet) {
-                Restore-NetworkAccess
                 Write-Error "✗ Failed to set function key on '$functionAppName' after retries."
                 exit 1
             }
@@ -234,14 +239,13 @@ else {
     $currentUserUpn = az ad signed-in-user show --query "userPrincipalName" -o tsv 2>$null
     $currentUserOid = az ad signed-in-user show --query "id" -o tsv 2>$null
     if (-not $currentUserUpn -or -not $currentUserOid) {
-        Restore-NetworkAccess
         Write-Error "✗ Could not determine current signed-in user. Ensure you are logged in with 'az login'."
         exit 1
     }
     Write-Host "✓ Current user: $currentUserUpn ($currentUserOid)"
 
     # Ensure current user is a PostgreSQL Entra administrator
-    $existingAdmins = az postgres flexible-server microsoft-entra-admin list --resource-group $ResourceGroupName --server-name $serverName --query "[].objectId" -o tsv 2>$null
+    $existingAdmins = az postgres flexible-server $pgAdminCmd list --resource-group $ResourceGroupName --server-name $serverName --query "[].objectId" -o tsv 2>$null
     $isAdmin = $false
     if ($existingAdmins) {
         foreach ($adminOid in ($existingAdmins -split "`n")) {
@@ -251,7 +255,7 @@ else {
     $addedPgAdmin = $false
     if (-not $isAdmin) {
         Write-Host "✓ Adding current user as PostgreSQL Entra administrator..."
-        $adminOutput = az postgres flexible-server microsoft-entra-admin create `
+        $adminOutput = az postgres flexible-server $pgAdminCmd create `
             --resource-group $ResourceGroupName `
             --server-name $serverName `
             --display-name $currentUserUpn `
@@ -290,7 +294,7 @@ else {
         # Remove temporary PostgreSQL admin if we added it
         if ($addedPgAdmin) {
             Write-Host "✓ Removing temporary PostgreSQL Entra admin for current user..."
-            az postgres flexible-server microsoft-entra-admin delete `
+            az postgres flexible-server $pgAdminCmd delete `
                 --resource-group $ResourceGroupName `
                 --server-name $serverName `
                 --object-id $currentUserOid `
@@ -308,10 +312,12 @@ else {
     Write-Host "✓ PostgreSQL table creation completed."
 }
 
-# -------------------------------------------------------
-# STEP 3 — Restore private networking (if it was enabled)
-# -------------------------------------------------------
-Restore-NetworkAccess
+} finally {
+    # -------------------------------------------------------
+    # STEP 3 — Restore private networking (if it was enabled)
+    # -------------------------------------------------------
+    Restore-NetworkAccess
+}
 
 Write-Host ""
 Write-Host "=============================================="

scripts/post_deployment_setup.sh

@@ -17,22 +17,32 @@ export MSYS_NO_PATHCONV=1
 # Usage: ./scripts/post_deployment_setup.sh <resource-group-name>
 
 if [ -z "$1" ]; then
-    echo "Usage: $0 <resource-group-name>"
-    exit 1
+    read -rp "Enter the resource group name: " RESOURCE_GROUP
+    if [ -z "$RESOURCE_GROUP" ]; then
+        echo "Resource group name is required."
+        exit 1
+    fi
+else
+    RESOURCE_GROUP="$1"
 fi
 
-RESOURCE_GROUP="$1"
-
 echo "=============================================="
 echo " Post-Deployment Setup"
 echo " Resource Group: ${RESOURCE_GROUP}"
 echo "=============================================="
 
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -W 2>/dev/null || pwd)"
 
-# Remove rdbms-connect extension if present (it conflicts with built-in microsoft-entra-admin commands)
+# Remove rdbms-connect extension if present (it conflicts with built-in admin commands)
 az extension remove --name rdbms-connect > /dev/null 2>&1 || true
 
+# Detect whether to use 'microsoft-entra-admin' (newer CLI) or 'ad-admin' (older CLI)
+if az postgres flexible-server microsoft-entra-admin --help > /dev/null 2>&1; then
+    PG_ADMIN_CMD="microsoft-entra-admin"
+else
+    PG_ADMIN_CMD="ad-admin"
+fi
+
 # Track resources that need public access restored to Disabled
 RESTORE_KV_NAME=""
 RESTORE_PG_NAME=""
@@ -48,6 +58,30 @@ restore_network_access() {
     fi
 }
 
+# Global cleanup function — handles PostgreSQL temp resources (if set) and network restore
+cleanup() {
+    # Remove temporary PostgreSQL admin if we added it
+    if [ "$ADDED_PG_ADMIN" = "true" ]; then
+        echo "✓ Removing temporary PostgreSQL Entra admin for current user..."
+        az postgres flexible-server $PG_ADMIN_CMD delete \
+            --resource-group "$RESOURCE_GROUP" \
+            --server-name "$SERVER_NAME" \
+            --object-id "$CURRENT_USER_OID" \
+            --yes 2>/dev/null || true
+    fi
+    # Remove temporary firewall rule if server was discovered
+    if [ -n "$SERVER_NAME" ]; then
+        echo "✓ Removing temporary firewall rule..."
+        az postgres flexible-server firewall-rule delete \
+            --resource-group "$RESOURCE_GROUP" \
+            --name "$SERVER_NAME" \
+            --rule-name "AllowPostDeploySetup" \
+            --yes 2>/dev/null || true
+    fi
+    restore_network_access
+}
+trap cleanup EXIT
+
 # -------------------------------------------------------
 # STEP 1 — Set Function App Client Key
 # -------------------------------------------------------
@@ -112,7 +146,6 @@ else
         FUNCTION_KEY=$(az keyvault secret show --vault-name "$KEY_VAULT_NAME" --name "FUNCTION-KEY" --query "value" -o tsv 2>/dev/null || true)
         if [ -z "$FUNCTION_KEY" ]; then
             echo "✗ ERROR: Failed to retrieve 'FUNCTION-KEY' secret from Key Vault '${KEY_VAULT_NAME}'." >&2
-            restore_network_access
             exit 1
         fi
 
@@ -170,7 +203,6 @@ else
         if [ -n "$PG_ERR" ]; then
             echo "✗ ERROR: Failed to enable public access on PostgreSQL. Cannot proceed." >&2
             echo "  $PG_ERR" >&2
-            restore_network_access
             exit 1
         fi
         RESTORE_PG_NAME="$SERVER_NAME"
@@ -207,13 +239,12 @@ else
     CURRENT_USER_OID=$(az ad signed-in-user show --query "id" -o tsv 2>/dev/null || true)
     if [ -z "$CURRENT_USER_UPN" ] || [ -z "$CURRENT_USER_OID" ]; then
         echo "✗ ERROR: Could not determine current signed-in user. Ensure you are logged in with 'az login'." >&2
-        restore_network_access
         exit 1
     fi
     echo "✓ Current user: ${CURRENT_USER_UPN} (${CURRENT_USER_OID})"
 
     # Ensure current user is a PostgreSQL Entra administrator
-    EXISTING_ADMINS=$(az postgres flexible-server microsoft-entra-admin list --resource-group "$RESOURCE_GROUP" --server-name "$SERVER_NAME" --query "[].objectId" -o tsv 2>/dev/null || true)
+    EXISTING_ADMINS=$(az postgres flexible-server $PG_ADMIN_CMD list --resource-group "$RESOURCE_GROUP" --server-name "$SERVER_NAME" --query "[].objectId" -o tsv 2>/dev/null || true)
     IS_ADMIN=false
     ADDED_PG_ADMIN=false
     if [ -n "$EXISTING_ADMINS" ]; then
@@ -226,7 +257,7 @@ else
     fi
     if [ "$IS_ADMIN" = "false" ]; then
         echo "✓ Adding current user as PostgreSQL Entra administrator..."
-        ADMIN_ERR=$(az postgres flexible-server microsoft-entra-admin create \
+        ADMIN_ERR=$(az postgres flexible-server $PG_ADMIN_CMD create \
             --resource-group "$RESOURCE_GROUP" \
             --server-name "$SERVER_NAME" \
             --display-name "$CURRENT_USER_UPN" \
@@ -244,27 +275,6 @@ else
         echo "✓ Current user is already a PostgreSQL Entra administrator."
     fi
 
-    # Ensure firewall rule cleanup on exit (along with network restore)
-    cleanup() {
-        # Remove temporary PostgreSQL admin if we added it
-        if [ "$ADDED_PG_ADMIN" = "true" ]; then
-            echo "✓ Removing temporary PostgreSQL Entra admin for current user..."
-            az postgres flexible-server microsoft-entra-admin delete \
-                --resource-group "$RESOURCE_GROUP" \
-                --server-name "$SERVER_NAME" \
-                --object-id "$CURRENT_USER_OID" \
-                --yes 2>/dev/null || true
-        fi
-        echo "✓ Removing temporary firewall rule..."
-        az postgres flexible-server firewall-rule delete \
-            --resource-group "$RESOURCE_GROUP" \
-            --name "$SERVER_NAME" \
-            --rule-name "AllowPostDeploySetup" \
-            --yes 2>/dev/null || true
-        restore_network_access
-    }
-    trap cleanup EXIT
-
     # Install Python dependencies
     REQUIREMENTS_FILE="${SCRIPT_DIR}/data_scripts/requirements.txt"
     if [ -f "$REQUIREMENTS_FILE" ]; then
@@ -277,14 +287,6 @@ else
     echo "✓ PostgreSQL table creation completed."
 fi
 
-# -------------------------------------------------------
-# STEP 3 — Restore private networking (if no PostgreSQL trap set it)
-# -------------------------------------------------------
-# If no PostgreSQL server was found, the trap won't fire, so restore here
-if [ -z "$SERVER_FQDN" ]; then
-    restore_network_access
-fi
-
 echo ""
 echo "=============================================="
 echo " Post-Deployment Setup Complete"