Abstract away the auth header param

datastax · Nov 6, 2023 · 6bb57f0 · 6bb57f0
1 parent c71e311
commit 6bb57f0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/astrapy/db.py b/astrapy/db.py
@@ -33,6 +33,7 @@ def __init__(
         token=None,
         api_endpoint=None,
         namespace=None,
+        auth_header=DEFAULT_AUTH_HEADER,
     ):
         if astra_db is None:
             if token is None or api_endpoint is None:
@@ -45,13 +46,14 @@ def __init__(
         self.astra_db = astra_db
         self.collection_name = collection_name
         self.base_path = f"{self.astra_db.base_path}/{collection_name}"
+        self.auth_header = auth_header
 
     def _request(self, *args, skip_error_check=False, **kwargs):
         response = make_request(
             *args,
             **kwargs,
             base_url=self.astra_db.base_url,
-            auth_header=DEFAULT_AUTH_HEADER,
+            auth_header=self.auth_header,
             token=self.astra_db.token,
         )
         responsebody = response.json()
@@ -290,6 +292,7 @@ def __init__(
         token=None,
         api_endpoint=None,
         namespace=None,
+        auth_header=DEFAULT_AUTH_HEADER,
     ):
         if token is None or api_endpoint is None:
             raise AssertionError("Must provide token and api_endpoint")
@@ -317,12 +320,15 @@ def __init__(
         # Set the namespace parameter
         self.namespace = namespace
 
+        # Set the default name of the auth header
+        self.auth_header = auth_header
+
     def _request(self, *args, skip_error_check=False, **kwargs):
         response = make_request(
             *args,
             **kwargs,
             base_url=self.base_url,
-            auth_header=DEFAULT_AUTH_HEADER,
+            auth_header=self.auth_header,
             token=self.token,
         )
 
@@ -333,8 +339,6 @@ def _request(self, *args, skip_error_check=False, **kwargs):
         else:
             return responsebody
 
-        return result
-
     def collection(self, collection_name):
         return AstraDBCollection(collection_name=collection_name, astra_db=self)
 

diff --git a/astrapy/ops.py b/astrapy/ops.py
@@ -21,19 +21,20 @@
 
 
 class AstraDBOps:
-    def __init__(self, token, dev_ops_url=None):
+    def __init__(self, token, dev_ops_url=None, auth_header="Authorization"):
         dev_ops_url = dev_ops_url or DEFAULT_DEV_OPS_URL
 
         self.token = "Bearer " + token
         self.base_url = f"https://{dev_ops_url}{DEFAULT_DEV_OPS_PATH_PREFIX}"
+        self.auth_header = auth_header
 
     def _ops_request(self, method, path, options=None, json_data=None):
         options = {} if options is None else options
 
         return make_request(
             base_url=self.base_url,
             method=method,
-            auth_header="Authorization",
+            auth_header=self.auth_header,
             token=self.token,
             json_data=json_data,
             url_params=options,