Merge pull request #69 from RelationalAI/pv-override-auth0-audience

add config option for overriding auth0 audience

Author: salamehsameera

railib/api.py

@@ -136,11 +136,12 @@ class Permission(str, Enum):
 # Context contains the state required to make rAI API calls.
 class Context(rest.Context):
     def __init__(self, host: str = None, port: str = None, scheme: str = None,
-                 region: str = None, credentials=None):
+                 region: str = None, credentials=None, audience: str = None):
         super().__init__(region=region, credentials=credentials)
         self.host = host
         self.port = port or "443"
         self.scheme = scheme or "https"
+        self.audience = audience
 
 
 # Construct a URL from the given context and path.

railib/config.py

@@ -84,7 +84,7 @@ def read(fname: str = "~/.rai/config", profile: str = "default"):
         raise Exception(f"can't find file: {path}")
     data = _read_config_profile(path, profile)
     creds = _read_credentials(data, path)
-    _keys = ["host", "port", "region", "scheme"]
+    _keys = ["host", "port", "region", "scheme", "audience"]
     result = {k: v for k, v in data.items() if k in _keys}
     result["credentials"] = creds
     return result

railib/rest.py

@@ -131,7 +131,7 @@ def _request_access_token(ctx: Context, url: str) -> AccessToken:
     creds = ctx.credentials
     assert type(creds) == ClientCredentials
     # ensure the audience contains the protocol scheme
-    audience = f"https://{_get_host(url)}"
+    audience = ctx.audience or f"https://{_get_host(url)}"
     headers = {
         "Accept": "application/json",
         "Content-Type": "application/json",