Use REST API to update API policy (#110)

Azure-Samples · Feb 6, 2025 · a532a8d · a532a8d
1 parent 3938935
commit a532a8d
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 19 deletions.
diff --git a/labs/zero-to-production/main.bicep b/labs/zero-to-production/main.bicep
@@ -124,6 +124,7 @@ output applicationInsightsAppId string = appInsightsModule.outputs.appId
 output applicationInsightsName string = appInsightsModule.outputs.applicationInsightsName
 output logAnalyticsWorkspaceId string = lawModule.outputs.customerId
 output apimServiceId string = apimModule.outputs.id
+output apimServiceName string = apimModule.outputs.name
 output apimResourceGatewayURL string = apimModule.outputs.gatewayUrl
 
 #disable-next-line outputs-should-not-contain-secrets

diff --git a/labs/zero-to-production/zero-to-production.ipynb b/labs/zero-to-production/zero-to-production.ipynb
@@ -188,6 +188,7 @@
     "\n",
     "if output.success and output.json_data:\n",
     "    apim_service_id = utils.get_deployment_output(output, 'apimServiceId', 'APIM Service Id')\n",
+    "    apim_service_name = utils.get_deployment_output(output, 'apimServiceName', 'APIM Service Name')\n",
     "    apim_resource_gateway_url = utils.get_deployment_output(output, 'apimResourceGatewayURL', 'APIM API Gateway URL')\n",
     "    apim_subscription1_key = utils.get_deployment_output(output, 'apimSubscription1Key', 'APIM Subscription 1 Key (masked)', True)\n",
     "    apim_subscription2_key = utils.get_deployment_output(output, 'apimSubscription2Key', 'APIM Subscription 2 Key (masked)', True)\n",
@@ -381,16 +382,12 @@
    "outputs": [],
    "source": [
     "policy_xml_file = \"policy-2.xml\"\n",
-    "bicep_parameters_file = \"params-2.json\"\n",
     "\n",
-    "bicep_parameters = utils.create_bicep_params(policy_xml_file, bicep_parameters_file, bicep_parameters, [\n",
-    "    ('{backend-id}', backend_id),\n",
-    "    ('{retry-count}', len(openai_resources) - 1)\n",
-    "])\n",
+    "with open(policy_xml_file, 'r') as file:\n",
+    "    policy_xml = file.read()\n",
+    "    policy_xml = policy_xml.replace('{backend-id}', backend_id).replace('{retry-count}', str(len(openai_resources) - 1))\n",
     "\n",
-    "# Run the deployment\n",
-    "output = utils.run(f\"az deployment group create --name {deployment_name} --resource-group {resource_group_name} --template-file main.bicep --parameters {bicep_parameters_file}\",\n",
-    "    f\"Deployment '{deployment_name}' succeeded\", f\"Deployment '{deployment_name}' failed\")"
+    "utils.update_api_policy(subscription_id, resource_group_name, apim_service_name, \"openai\", policy_xml)"
    ]
   },
   {
@@ -563,18 +560,13 @@
    "outputs": [],
    "source": [
     "policy_xml_file = \"policy-3.xml\"\n",
-    "bicep_parameters_file = \"params-3.json\"\n",
     "tokens_per_minute = 500\n",
     "\n",
-    "bicep_parameters = utils.create_bicep_params(policy_xml_file, bicep_parameters_file, bicep_parameters, [\n",
-    "    ('{backend-id}', backend_id),\n",
-    "    ('{retry-count}', len(openai_resources) - 1),\n",
-    "    ('{tpm}', tokens_per_minute)\n",
-    "])\n",
+    "with open(policy_xml_file, 'r') as file:\n",
+    "    policy_xml = file.read()\n",
+    "    policy_xml = policy_xml.replace('{backend-id}', backend_id).replace('{retry-count}', str(len(openai_resources) - 1)).replace('{tpm}', str(tokens_per_minute))\n",
     "\n",
-    "# Run the deployment\n",
-    "output = utils.run(f\"az deployment group create --name {deployment_name} --resource-group {resource_group_name} --template-file main.bicep --parameters {bicep_parameters_file}\",\n",
-    "    f\"Deployment '{deployment_name}' succeeded\", f\"Deployment '{deployment_name}' failed\")"
+    "utils.update_api_policy(subscription_id, resource_group_name, apim_service_name, \"openai\", policy_xml)"
    ]
   },
   {
@@ -679,7 +671,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

diff --git a/shared/utils.py b/shared/utils.py
@@ -1,4 +1,4 @@
-import datetime, json, os, subprocess, time, traceback
+import datetime, json, os, subprocess, requests, time, traceback
 
 # Define ANSI escape code constants vor clarity in the print commands below
 RESET_FORMATTING = "\x1b[0m"
@@ -200,3 +200,29 @@ def create_bicep_params(policy_xml_filepath, parameters_filepath, bicep_paramete
     print(f"📝 Updated the policy XML in the bicep parameters file '{parameters_filepath}'")
 
     return bicep_parameters
+
+def update_api_policy(subscription_id, resource_group_name, apim_service_name, api_id, policy_xml):
+    # We first need to obtain an access token for the REST API
+    output = run(f"az account get-access-token --resource https://management.azure.com/",
+        f"Successfully obtained access token", f"Failed to obtain access token")
+
+    if output.success and output.json_data:
+        access_token = output.json_data['accessToken']
+
+        print("Updating the API policy...")
+        # https://learn.microsoft.com/en-us/rest/api/apimanagement/api-policy/create-or-update?view=rest-apimanagement-2024-06-01-preview
+        url = f"https://management.azure.com/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.ApiManagement/service/{apim_service_name}/apis/{api_id}/policies/policy?api-version=2024-06-01-preview"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {access_token}"
+        }
+
+        body = {
+            "properties": {
+                "format": "rawxml",
+                "value": policy_xml
+            }
+        }
+
+        response = requests.put(url, headers = headers, json = body)
+        print_response_code(response)