From 6bb57f007514fb57dac450cfe3ae479f779a228f Mon Sep 17 00:00:00 2001
From: Eric Hare <ericrhare@gmail.com>
Date: Mon, 6 Nov 2023 07:29:23 -0800
Subject: [PATCH] Abstract away the auth header param

---
 astrapy/db.py  | 12 ++++++++----
 astrapy/ops.py |  5 +++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/astrapy/db.py b/astrapy/db.py
index a82f81f2..70f0c471 100644
--- 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
index 249c7d4c..f574aebf 100644
--- a/astrapy/ops.py
+++ b/astrapy/ops.py
@@ -21,11 +21,12 @@
 
 
 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
@@ -33,7 +34,7 @@ def _ops_request(self, method, path, options=None, json_data=None):
         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,