Testing NERSC superfacility API

nersc_api_test.py

import requests
import sys

from authlib.integrations.requests_client import OAuth2Session
from authlib.oauth2.rfc7523 import PrivateKeyJWT


TOKEN_URL = "https://oidc.nersc.gov/c2id/token"


def get_token(client_id, private_key):

    session = OAuth2Session(
        client_id, 
        private_key, 
        PrivateKeyJWT(TOKEN_URL),
        grant_type="client_credentials",
        token_endpoint=TOKEN_URL,
    )
    return session.fetch_token()


def main():
    client_id = sys.argv[1]
    private_key_file = sys.argv[2]
    with open(private_key_file) as f:
       private_key = f.read().strip()
    token = get_token(client_id, private_key)
    print(token)
    token = token["access_token"]
    r = requests.post(
        "https://api.nersc.gov/api/v1.2/utilities/command/dtns",
        data={"executable": "bash -c pwd"},
        headers={"accept": "application/json", "Authorization": token})
    print(r.json())
    taskid = r.json()["task_id"]
    r = requests.get(
        "https://api.nersc.gov/api/v1.2/tasks/" + taskid,
        headers={"accept": "application/json", "Authorization": token})
    print(r.json())
    

if __name__ == "__main__":
    main()