added token to init, created setup_authentication function, adjusted request fuction

gitworkflows · web-flow · commit 1e0889d45319 · 2025-04-14T17:30:56.000+06:00
Signed-off-by: gitworkflows &lt;118260833+gitworkflows@users.noreply.github.com&gt;
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -47,6 +47,8 @@ class InfluxDBClient(object):
     :type username: str
     :param password: password of the user, defaults to 'root'
     :type password: str
+    :param token: token for access.  Overrides username/password
+    :type token: str
     :param pool_size: urllib3 connection pool size, defaults to 10.
     :type pool_size: int
     :param database: database name to connect to, defaults to None
@@ -101,6 +103,7 @@ def __init__(self,
                  port=8086,
                  username='root',
                  password='root',
+                 token=None,
                  database=None,
                  ssl=False,
                  verify_ssl=False,
@@ -125,6 +128,7 @@ def __init__(self,
         self._database = database
         self._timeout = timeout
         self._retries = retries
+        self._token = token
 
         self._verify_ssl = verify_ssl
 
@@ -134,7 +138,7 @@ def __init__(self,
         if not session:
             session = requests.Session()
 
-        self._session = session
+        self._session = self.setup_authenticaiton(session, self._token, self._username, self._password)
         adapter = _SocketOptionsAdapter(
             pool_connections=int(pool_size),
             pool_maxsize=int(pool_size),
@@ -258,6 +262,16 @@ def from_dsn(cls, dsn, **kwargs):
 
         return cls(**init_args)
 
+    def setup_authenticaiton(self, session, token=None, username=None, password=None):
+
+        if token:
+            session.headers.update({"Authorization": f"Token {token}"})
+            return session
+
+        session.auth = (username, password)
+
+        return session
+
     def switch_database(self, database):
         """Change the client's database.
 
@@ -266,16 +280,22 @@ def switch_database(self, database):
         """
         self._database = database
 
-    def switch_user(self, username, password):
-        """Change the client's username.
+    def switch_user(self, token=None, username=None, password=None):
+        """Change the client's user.
 
+        :param token: the token of the user to switch to
+        :type token: str
         :param username: the username to switch to
         :type username: str
         :param password: the password for the username
         :type password: str
         """
-        self._username = username
-        self._password = password
+        if token:
+            self._session.headers.update({"Authorization": f"Token {token}"})
+            self._session.auth = None
+
+        self._session.auth = (username, password)
+        self._session.headers.pop("Authorization", None)
 
     def request(self, url, method='GET', params=None, data=None, stream=False,
                 expected_response_code=200, headers=None):
@@ -336,14 +356,9 @@ def request(self, url, method='GET', params=None, data=None, stream=False,
         _try = 0
         while retry:
             try:
-                if "Authorization" in headers:
-                    auth = (None, None)
-                else:
-                    auth = (self._username, self._password)
                 response = self._session.request(
                     method=method,
                     url=url,
-                    auth=auth if None not in auth else None,
                     params=params,
                     data=data,
                     stream=stream,
