def update_cluster_env_vars(session, team_id, cluster_id, updates):
"""
Add or update environment variables in a Mage Pro cluster.
Args:
session (requests.Session): Authenticated session with Mage Cloud API.
team_id (str): Mage team ID.
cluster_id (str): Cluster ID.
updates (dict): Dictionary of environment variable key-value pairs to add or update.
"""
cluster_url = f'https://cloud.mage.ai/api/v1/teams/{team_id}/clusters/{cluster_id}'
# Fetch current cluster config
response = session.get(cluster_url)
response.raise_for_status()
cluster = response.json()['cluster']
config = cluster.get('config', {})
# Load existing environment variables
existing_envs = config.get('extraEnvs', [])
env_map = {env['name']: env['value'] for env in existing_envs}
# Apply updates
env_map.update(updates)
# Convert back to list format
config['extraEnvs'] = [{'name': k, 'value': v} for k, v in env_map.items()]
# Build payload and apply update
update_payload = {
"cluster": {
"config": config
}
}
put_response = session.put(cluster_url, json=update_payload)
put_response.raise_for_status()
print("Cluster environment variables updated successfully.")