Add AAD to workload configs (#42389)

tvaron3 · web-flow · commit 51e03f878c2f · 2025-08-18T12:33:11.000-07:00
* update workloads to use aad

* update workloads to use aad

* improve setup scripts

* revert bicep change

* fix analyze

* fix upserts

* fix upserts

* made new script an executable

* fix path to dev_requirements

* fix setup script

* fix setup script

* fix setup script

* fix setup script

* staggered startup

* fix setup script

* fix setup script
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/dev.md b/sdk/cosmos/azure-cosmos/tests/workloads/dev.md
@@ -9,25 +9,20 @@ file when run. These logs are named in this format `<file name>-<process id>-<da
    - 32 GB RAM
    - Ubuntu
    - Accelerated networking
+1. Give the VM necessary [permissions](https://learn.microsoft.com/azure/cosmos-db/nosql/how-to-grant-data-plane-access?tabs=built-in-definition%2Ccsharp&pivots=azure-interface-cli) to access the Cosmos DB account if using AAD (Optional). 
 1. Create an Azure App Insights Resource (Optional)
 1. Fork and clone this repository
 1. Go to azure cosmos folder
    - `cd azure-sdk-for-python/sdk/cosmos/azure-cosmos`
 1. Install the required packages and create virtual environment
-   - `sudo apt-get update`
-   - `sudo apt-get install python3-pip`
-   - `sudo apt-get install python3.12-venv`
-   - `python3 -m venv azure-cosmosdb-sdk-environment`
+   - `setup_env.sh` 
    - `source azure-cosmosdb-sdk-environment/bin/activate`
-   - `pip install -r dev_requirements.txt`
-   - `pip install azure-monitor-opentelemetry`
 1. Checkout the branch with the changes to test. 
 1. Install azure-cosmos
    - `pip install .`
 1. Go to workloads folder
     - `cd tests/workloads`
 1. Fill out relevant configs in `workload_configs.py`: key, host, etc
-1. Install envoy proxy https://www.envoyproxy.io/docs/envoy/latest/start/install
 1. Go to envoy folder and generate envoy configuration file using template. Template files are in `envoy/templates` directory. `<account_name>` is your Cosmos DB account name.
     - `cd envoy`
     - `./generate_envoy_config.sh <template_file_path> <output_envoy_config_file> <account_name> <write_region> <read_region>`
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/get-database-account-call.py b/sdk/cosmos/azure-cosmos/tests/workloads/get-database-account-call.py
@@ -1,12 +1,9 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
 import os
-import sys
 
 from workload_utils import create_logger
-from workload_configs import COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS
-
-sys.path.append(r"./")
+from workload_configs import COSMOS_URI, COSMOS_CREDENTIAL, PREFERRED_LOCATIONS
 
 from azure.cosmos.aio import CosmosClient as AsyncClient
 import asyncio
@@ -16,7 +13,7 @@
 
 
 async def run_workload(client_id: str):
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY, preferred_locations=PREFERRED_LOCATIONS,
+    async with AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, preferred_locations=PREFERRED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=logger,
                            user_agent=client_id + "-" + datetime.now().strftime(
                                "%Y%m%d-%H%M%S")) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/initial-setup.py b/sdk/cosmos/azure-cosmos/tests/workloads/initial-setup.py
@@ -1,15 +1,12 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
 import os
-import sys
 
 from azure.cosmos import PartitionKey, ThroughputProperties
-from workload_utils import create_logger
+from workload_utils import create_logger, create_random_item
 from workload_configs import COSMOS_URI, COSMOS_KEY, PREFERRED_LOCATIONS, COSMOS_CONTAINER, COSMOS_DATABASE, \
     NUMBER_OF_LOGICAL_PARTITIONS, PARTITION_KEY, THROUGHPUT
 
-sys.path.append(r"./")
-
 from azure.cosmos.aio import CosmosClient as AsyncClient
 import asyncio
 
@@ -18,11 +15,16 @@
 async def write_item_concurrently_initial(container, num_upserts):
     tasks = []
     for i in range(num_upserts):
-        tasks.append(container.upsert_item({"id": "test-" + str(i), "pk": "pk-" + str(i)}))
+        item = create_random_item()
+        item["id"] = "test-" + str(i)
+        item["pk"] = "pk-" + str(i)
+        tasks.append(container.upsert_item(item))
     await asyncio.gather(*tasks)
 
 
 async def run_workload(client_id: str):
+    # Key always needs to be used for the initial setup to create the database and container as aad for control plane
+    # operations using the dataplane sdk is not supported.
     async with AsyncClient(COSMOS_URI, COSMOS_KEY, preferred_locations=PREFERRED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=logger,
                            user_agent=str(client_id) + "-" + datetime.now().strftime("%Y%m%d-%H%M%S")) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_proxy_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_proxy_workload.py
@@ -9,7 +9,7 @@
 async def run_workload(client_id, client_logger):
     session = create_custom_session()
     async with AsyncClient(COSMOS_URI,
-                           COSMOS_KEY,
+                           COSMOS_CREDENTIAL,
                            multiple_write_locations=USE_MULTIPLE_WRITABLE_LOCATIONS,
                            preferred_locations=PREFERRED_LOCATIONS,
                            transport=AioHttpTransport(session=session, session_owner=False),
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_proxy_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_proxy_workload.py
@@ -11,7 +11,7 @@
 async def run_workload(client_id, client_logger):
     session = create_custom_session()
     async with AsyncClient(COSMOS_URI,
-                           COSMOS_KEY,
+                           COSMOS_CREDENTIAL,
                            multiple_write_locations=USE_MULTIPLE_WRITABLE_LOCATIONS,
                            preferred_locations=PREFERRED_LOCATIONS,
                            transport=AioHttpTransport(session=session, session_owner=False),
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_sync_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_sync_workload.py
@@ -10,7 +10,7 @@
 def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    with CosmosClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+    with CosmosClient(COSMOS_URI, COSMOS_CREDENTIAL, connection_policy=connectionPolicy,
                       preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                       enable_diagnostics_logging=True, logger=client_logger,
                       user_agent=get_user_agent(client_id)) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_with_incorrect_client_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_with_incorrect_client_workload.py
@@ -11,7 +11,7 @@
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    client = AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+    client = AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, connection_policy=connectionPolicy,
                          preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                          enable_diagnostics_logging=True, logger=client_logger,
                          user_agent=get_user_agent(client_id))
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload.py
@@ -13,7 +13,7 @@
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+    async with AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, connection_policy=connectionPolicy,
                            preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=client_logger,
                            user_agent=get_user_agent(client_id)) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/r_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/r_workload.py
@@ -11,7 +11,7 @@
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+    async with AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, connection_policy=connectionPolicy,
                            preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=client_logger,
                            user_agent=get_user_agent(client_id)) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/run_workloads.sh b/sdk/cosmos/azure-cosmos/tests/workloads/run_workloads.sh
@@ -6,9 +6,22 @@ fi
 
 num_runs=$1
 
+echo "[Info] Installing azure-cosmos package..."
+pip install ../../.
+if [ $? -ne 0 ]; then
+    echo "[Error] Failed to install azure-cosmos. Exiting."
+    exit 2
+fi
+echo "[Info] azure-cosmos installed successfully."
+
 # Loop over each Python file in the current directory ending with _workload.py
 for file in ./*_workload.py; do
     for (( i=0; i<num_runs; i++ )); do
         python3 "$file" &
+        # slow start up
+        sleep 1
     done
 done
+
+wait
+echo "[Info] All workloads started successfully."
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/setup_env.sh b/sdk/cosmos/azure-cosmos/tests/workloads/setup_env.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+# cspell:disable
+# setup_env.sh - Automates Azure Cosmos SDK scale testing environment setup
+# Usage: bash setup_env.sh
+
+set -e
+
+# 1. System update and install dependencies
+echo "[Step 1] System update and install dependencies: started."
+sudo apt-get update
+sudo apt-get install -y python3-pip python3.12-venv wget gnupg
+sudo apt install neovim
+curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
+echo "[Step 1] System update and install dependencies: completed."
+
+# 2. Create and activate Python virtual environment
+echo "[Step 2] Create and activate Python virtual environment: started."
+python3 -m venv azure-cosmosdb-sdk-environment
+source azure-cosmosdb-sdk-environment/bin/activate
+echo "[Step 2] Create and activate Python virtual environment: completed."
+
+# 3. Install Python requirements
+echo "[Step 3] Install Python requirements: started."
+pip install --upgrade pip
+cd ../../
+pip install -r dev_requirements.txt
+pip install azure-monitor-opentelemetry
+# azure core needs to be reinstalled to avoid cannot find Spanning Error from open telemetry
+pip install ../../core/azure-core
+echo "[Step 3] Install Python requirements: completed."
+
+# 4. Install the current azure-cosmos package
+echo "[Step 4] Install the current azure-cosmos package: started."
+pip install .
+cd tests/workloads
+echo "[Step 4] Install the current azure-cosmos package: completed."
+
+# 5. Install Envoy proxy
+echo "[Step 5] Install Envoy proxy: started."
+
+if command -v envoy &> /dev/null; then
+    echo "Envoy proxy is already installed. Version: $(envoy --version)"
+else
+    wget -O- https://apt.envoyproxy.io/signing.key | sudo gpg --dearmor -o /etc/apt/keyrings/envoy-keyring.gpg
+    echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/envoy-keyring.gpg] https://apt.envoyproxy.io focal main" | sudo tee /etc/apt/sources.list.d/envoy.list
+    sudo apt-get update
+    sudo apt-get install -y envoy
+    envoy --version
+fi
+# create logs directory for envoy
+cd envoy
+mkdir logs
+cd ..
+
+echo "[Step 5] Install Envoy proxy: completed."
+
+# 6. Manual steps
+echo "[Step 6] Manual steps: started."
+cat << EOF
+
+Manual steps remaining:
+----------------------
+1. Fill out relevant configs in tests/workloads/workload_configs.py (key, host, etc).
+2. Generate Envoy config:
+   cd envoy
+   ./generate_envoy_config.sh <template_file_path> <output_envoy_config_file> <account_name> <write_region> <read_region>
+3. Start Envoy:
+   envoy -c <envoy_config_file>.yaml --log-level debug --log-path logs/debug.txt
+4. Run initial setup workload:
+   cd ../tests/workloads
+   python3 initial-setup.py
+5. Run scale workloads:
+   ./run_workloads.sh <number of clients per workload>
+
+Refer to dev.md for more details.
+EOF
+echo "[Step 6] Manual steps: completed."
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/w_proxy_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/w_proxy_workload.py
@@ -11,7 +11,7 @@
 async def run_workload(client_id, client_logger):
     session = create_custom_session()
     async with AsyncClient(COSMOS_URI,
-                           COSMOS_KEY,
+                           COSMOS_CREDENTIAL,
                            multiple_write_locations=USE_MULTIPLE_WRITABLE_LOCATIONS,
                            preferred_locations=PREFERRED_LOCATIONS,
                            transport=AioHttpTransport(session=session, session_owner=False),
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/w_workload.py b/sdk/cosmos/azure-cosmos/tests/workloads/w_workload.py
@@ -14,7 +14,7 @@
 async def run_workload(client_id, client_logger):
     connectionPolicy = documents.ConnectionPolicy()
     connectionPolicy.UseMultipleWriteLocations = USE_MULTIPLE_WRITABLE_LOCATIONS
-    async with AsyncClient(COSMOS_URI, COSMOS_KEY, connection_policy=connectionPolicy,
+    async with AsyncClient(COSMOS_URI, COSMOS_CREDENTIAL, connection_policy=connectionPolicy,
                            preferred_locations=PREFERRED_LOCATIONS, excluded_locations=CLIENT_EXCLUDED_LOCATIONS,
                            enable_diagnostics_logging=True, logger=client_logger,
                            user_agent=get_user_agent(client_id)) as client:
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/workload_configs.py b/sdk/cosmos/azure-cosmos/tests/workloads/workload_configs.py
@@ -2,13 +2,15 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Replace with your Cosmos DB details
 import logging
+from azure.identity import DefaultAzureCredential
 
 PREFERRED_LOCATIONS = []
 CLIENT_EXCLUDED_LOCATIONS = []
 REQUEST_EXCLUDED_LOCATIONS = []
 COSMOS_PROXY_URI = "0.0.0.0"
 COSMOS_URI = ""
 COSMOS_KEY = ""
+COSMOS_CREDENTIAL = COSMOS_KEY if COSMOS_KEY else DefaultAzureCredential()
 COSMOS_CONTAINER = "scale_cont"
 COSMOS_DATABASE = "scale_db"
 USER_AGENT_PREFIX = ""
diff --git a/sdk/cosmos/azure-cosmos/tests/workloads/workload_utils.py b/sdk/cosmos/azure-cosmos/tests/workloads/workload_utils.py
@@ -3,6 +3,7 @@
 import asyncio
 import os
 import random
+import uuid
 from datetime import datetime
 from logging.handlers import RotatingFileHandler
 from aiohttp import ClientSession
@@ -19,22 +20,51 @@ def get_user_agent(client_id):
     prefix = USER_AGENT_PREFIX + "-" if USER_AGENT_PREFIX else ""
     return prefix + str(client_id) + "-" + datetime.now().strftime("%Y%m%d-%H%M%S")
 
-def get_random_item():
+def get_existing_random_item():
     random_int = random.randint(0, NUMBER_OF_LOGICAL_PARTITIONS)
-    return {"id": "test-" + str(random_int), "pk": "pk-" + str(random_int)}
+    item = create_random_item()
+    item["id"] = "test-" + str(random_int)
+    item["pk"] = "pk-" + str(random_int)
+    return item
+
+def create_random_item():
+    paragraph1 = (
+        "Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
+        "Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. "
+        "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. "
+        f"Random ID: {uuid.uuid4()}"
+    )
+    paragraph2 = (
+        "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. "
+        "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. "
+        f" Timestamp: {datetime.utcnow().isoformat()}"
+    )
+    return {
+        "id": "test-" + str(uuid.uuid4()),
+        "pk": "pk-" + str(uuid.uuid4()),
+        "value": random.randint(1, 1000000000),
+        "timestamp": datetime.utcnow().isoformat() + "Z",
+        "flag": random.choice([True, False]),
+        "description": paragraph1 + "\n\n" + paragraph2
+    }
+
+def _get_upsert_item():
+    # 10 percent of the time, create a new item instead of updating an existing one
+    return create_random_item() if random.random() < 0.1 else get_existing_random_item()
 
 def upsert_item(container, excluded_locations, num_upserts):
+    item = _get_upsert_item()
     for _ in range(num_upserts):
         if excluded_locations:
-            container.upsert_item(get_random_item(), etag=None, match_condition=None,
+            container.upsert_item(item, etag=None, match_condition=None,
                                   excluded_locations=excluded_locations)
         else:
-            container.upsert_item(get_random_item(), etag=None, match_condition=None)
+            container.upsert_item(item, etag=None, match_condition=None)
 
 
 def read_item(container, excluded_locations, num_reads):
     for _ in range(num_reads):
-        item = get_random_item()
+        item = get_existing_random_item()
         if excluded_locations:
             container.read_item(item["id"], item[PARTITION_KEY], etag=None, match_condition=None,
                                 excluded_locations=excluded_locations)
@@ -47,7 +77,7 @@ def query_items(container, excluded_locations, num_queries):
 
 
 def perform_query(container, excluded_locations):
-    random_item = get_random_item()
+    random_item = get_existing_random_item()
     if excluded_locations:
         results = container.query_items(query="SELECT * FROM c where c.id=@id and c.pk=@pk",
                                         parameters=[{"name": "@id", "value": random_item["id"]},
@@ -64,18 +94,19 @@ def perform_query(container, excluded_locations):
 async def upsert_item_concurrently(container, excluded_locations, num_upserts):
     tasks = []
     for _ in range(num_upserts):
+        item = _get_upsert_item()
         if excluded_locations:
-            tasks.append(container.upsert_item(get_random_item(), etag=None, match_condition=None,
+            tasks.append(container.upsert_item(item, etag=None, match_condition=None,
                                                excluded_locations=excluded_locations))
         else:
-            tasks.append(container.upsert_item(get_random_item(), etag=None, match_condition=None))
+            tasks.append(container.upsert_item(item, etag=None, match_condition=None))
     await asyncio.gather(*tasks)
 
 
 async def read_item_concurrently(container, excluded_locations, num_reads):
     tasks = []
     for _ in range(num_reads):
-        item = get_random_item()
+        item = get_existing_random_item()
         if excluded_locations:
             tasks.append(container.read_item(item["id"], item[PARTITION_KEY], etag=None, match_condition=None,
                                              excluded_locations=excluded_locations))
@@ -102,7 +133,7 @@ def create_custom_session():
     return session
 
 async def perform_query_concurrently(container, excluded_locations):
-    random_item = get_random_item()
+    random_item = get_existing_random_item()
     if excluded_locations:
         results = container.query_items(query="SELECT * FROM c where c.id=@id and c.pk=@pk",
                                         parameters=[{"name": "@id", "value": random_item["id"]},
