From 67af06acf9bae7853d039ce8f2e6d7bd14513432 Mon Sep 17 00:00:00 2001
From: conradoverta <53195685+conradoverta@users.noreply.github.com>
Date: Mon, 12 Jun 2023 14:29:55 -0700
Subject: [PATCH] feat: Allow specifying org id in client (#3878)

* feat: Allow specifying org id in client

* review feedback
---
 client/verta/verta/client.py      |  4 ++++
 client/verta/verta/credentials.py | 28 ++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/client/verta/verta/client.py b/client/verta/verta/client.py
index 9df83dbe54..e55a040af6 100644
--- a/client/verta/verta/client.py
+++ b/client/verta/verta/client.py
@@ -94,6 +94,8 @@ class Client(object):
         Whether to print extra verbose information to aid in debugging.
     extra_auth_headers : dict, default {}
         Extra headers to include on requests, like to permit traffic through a restrictive application load balancer
+    organization_id : str, optional
+        (alpha) Organization to use for the client calls. If not provided, the default organization will be used.
     _connect : str, default True
         Whether to connect to server (``False`` for unit tests).
 
@@ -142,6 +144,7 @@ def __init__(
         extra_auth_headers={},
         jwt_token=None,
         jwt_token_sig=None,
+        organization_id=None,
         _connect=True,
     ):
         self._load_config()
@@ -170,6 +173,7 @@ def __init__(
             dev_key=dev_key,
             jwt_token=jwt_token,
             jwt_token_sig=jwt_token_sig,
+            organization_id=organization_id, # TODO: add organization_name as parameter and resolve that
         )
         self._workspace = self._get_with_fallback(None, env_var="VERTA_WORKSPACE")
 
diff --git a/client/verta/verta/credentials.py b/client/verta/verta/credentials.py
index 381354b848..9c1c7a55b1 100644
--- a/client/verta/verta/credentials.py
+++ b/client/verta/verta/credentials.py
@@ -46,21 +46,24 @@ class EmailCredentials(Credentials):
         A user email address.
     dev_key : str
         A dev key to use for authentication.
+    organization_id : str, optional
+        An organization ID to use for authentication.
     """
 
     EMAIL_ENV = "VERTA_EMAIL"
     DEV_KEY_ENV = "VERTA_DEV_KEY"
 
-    def __init__(self, email, dev_key):
+    def __init__(self, email, dev_key, organization_id=None):
         self.email = email
         self.dev_key = dev_key
+        self.organization_id = organization_id
 
     def export_env_vars_to_os(self):
         os.environ[self.EMAIL_ENV] = self.email
         os.environ[self.DEV_KEY_ENV] = self.dev_key
 
     def headers(self):
-        return {
+        headers = {
             "source": "PythonClient",
             "email": self.email,
             "developer_key": self.dev_key,
@@ -68,10 +71,13 @@ def headers(self):
             # https://www.nginx.com/resources/wiki/start/topics/tutorials/config_pitfalls#missing-disappearing-http-headers
             "developer-key": self.dev_key,
         }
+        if self.organization_id:
+            headers["organization-id"] = self.organization_id
+        return headers
 
     def __repr__(self):
         key = self.dev_key[:8] + re.sub(r"[^-]", "*", self.dev_key[8:])
-        return "EmailCredentials({}, {})".format(self.email, key)
+        return "EmailCredentials({}, {}, {})".format(self.email, key, self.organization_id)
 
     @classmethod
     def load_from_os_env(cls):
@@ -92,14 +98,17 @@ class JWTCredentials(Credentials):
         A jwt token.
     jwt_token_sig : str
         A jwt token signature.
+    organization_id : str, optional
+        An organization ID to use for authentication.
     """
 
     JWT_TOKEN_ENV = "VERTA_JWT_TOKEN"
     JWT_TOKEN_SIG_ENV = "VERTA_JWT_TOKEN_SIG"
 
-    def __init__(self, jwt_token, jwt_token_sig):
+    def __init__(self, jwt_token, jwt_token_sig, organization_id=None):
         self.jwt_token = jwt_token
         self.jwt_token_sig = jwt_token_sig
+        self.organization_id = organization_id
 
     def export_env_vars_to_os(self):
         os.environ[self.JWT_TOKEN_ENV] = self.jwt_token
@@ -113,16 +122,19 @@ def headers(self):
             # without underscore, for NGINX support
             # https://www.nginx.com/resources/wiki/start/topics/tutorials/config_pitfalls#missing-disappearing-http-headers
             "bearer-access-token": self.jwt_token,
+            "organization-id": self.organization_id or "",
         }
         if self.jwt_token_sig:
             headers["bearer_access_token_sig"] = headers[
                 "bearer-access-token-sig"
             ] = self.jwt_token_sig
+        if self.organization_id:
+            headers["organization-id"] = self.organization_id
         return headers
 
     def __repr__(self):
         token = self.jwt_token[:8] + re.sub(r"[^-]", "*", self.jwt_token[8:])
-        return "JWTCredentials({}, {})".format(token, self.jwt_token_sig)
+        return "JWTCredentials({}, {}, {})".format(token, self.jwt_token_sig, self.organization_id)
 
     @classmethod
     def load_from_os_env(cls):
@@ -152,11 +164,11 @@ def load_from_os_env():
     return credentials
 
 
-def _build(email=None, dev_key=None, jwt_token=None, jwt_token_sig=None):
+def _build(email=None, dev_key=None, jwt_token=None, jwt_token_sig=None, organization_id=None):
     if email and dev_key:
-        return EmailCredentials(email, dev_key)
+        return EmailCredentials(email, dev_key, organization_id=organization_id)
     elif jwt_token:
-        return JWTCredentials(jwt_token, jwt_token_sig)
+        return JWTCredentials(jwt_token, jwt_token_sig, organization_id=organization_id)
     elif email or dev_key:
         raise ValueError("`email` and `dev_key` must be provided together")
     else: