GML-1660 add dbversion to set authMode

tigergraph · May 10, 2024 · 5537745 · 5537745
1 parent e3aac00
commit 5537745
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 6 deletions.
diff --git a/pyTigerGraph/pyTigerGraph.py b/pyTigerGraph/pyTigerGraph.py
@@ -31,9 +31,9 @@ def __init__(self, host: str = "http://127.0.0.1", graphname: str = "MyGraph",
             tgCloud: bool = False, restppPort: Union[int, str] = "9000",
             gsPort: Union[int, str] = "14240", gsqlVersion: str = "", version: str = "",
             apiToken: str = "", useCert: bool = None, certPath: str = None, debug: bool = None,
-            sslPort: Union[int, str] = "443", gcp: bool = False, jwtToken: str = ""):
+            sslPort: Union[int, str] = "443", gcp: bool = False, jwtToken: str = "", dbVersion: str = ""):
         super().__init__(host, graphname, gsqlSecret, username, password, tgCloud, restppPort,
-            gsPort, gsqlVersion, version, apiToken, useCert, certPath, debug, sslPort, gcp, jwtToken)
+            gsPort, gsqlVersion, version, apiToken, useCert, certPath, debug, sslPort, gcp, jwtToken, dbVersion)
 
         self.gds = None
         self.ai = None

diff --git a/pyTigerGraph/pyTigerGraphBase.py b/pyTigerGraph/pyTigerGraphBase.py
@@ -11,7 +11,7 @@
 import warnings
 from typing import Union
 from urllib.parse import urlparse
-
+from distutils.version import StrictVersion
 import requests
 
 from pyTigerGraph.pyTigerGraphException import TigerGraphException
@@ -36,7 +36,7 @@ def __init__(self, host: str = "http://127.0.0.1", graphname: str = "MyGraph",
             tgCloud: bool = False, restppPort: Union[int, str] = "9000",
             gsPort: Union[int, str] = "14240", gsqlVersion: str = "", version: str = "",
             apiToken: str = "", useCert: bool = None, certPath: str = None, debug: bool = None,
-            sslPort: Union[int, str] = "443", gcp: bool = False, jwtToken: str = ""):
+            sslPort: Union[int, str] = "443", gcp: bool = False, jwtToken: str = "", dbVersion: str = ""):
         """Initiate a connection object.
 
         Args:
@@ -78,6 +78,8 @@ def __init__(self, host: str = "http://127.0.0.1", graphname: str = "MyGraph",
                 DEPRECATED. Previously used for connecting to databases provisioned on GCP in TigerGraph Cloud.
             jwtToken:
                 The JWT token generated from customer side for authentication
+            dbVersion:
+                The TigerGraph database version being connected to.
 
         Raises:
             TigerGraphException: In case on invalid URL scheme.
@@ -126,9 +128,10 @@ def __init__(self, host: str = "http://127.0.0.1", graphname: str = "MyGraph",
         else:
             self.authHeader = {"Authorization": "Basic {0}".format(self.base64_credential)}
 
-        # If JWT token is provided, set authMode to "token", and overwrite authMode = "pwd" for GSQL authentication as well
+        # If JWT token is provided, set authMode to "token", and overwrite authMode = "pwd" for GSQL authentication as well if version is newer than 4.1.0
         if jwtToken:
-            self.authMode = "token"
+            if dbVersion and StrictVersion(dbVersion) >= StrictVersion("4.1.0"):
+                self.authMode = "token"
             self.jwtToken = jwtToken
             self.authHeader = {"Authorization": "Bearer " + self.jwtToken}
 

diff --git a/tests/pyTigerGraphUnitTest.py b/tests/pyTigerGraphUnitTest.py
@@ -20,6 +20,8 @@ def make_connection(graphname: str = None):
         "sslPort": "443",
         "tgCloud": False,
         "gcp": False,
+        "jwtToken": "",
+        "dbVersion": ""
     }
 
     path = os.path.dirname(os.path.realpath(__file__))
@@ -42,6 +44,8 @@ def make_connection(graphname: str = None):
         certPath=server_config["certPath"],
         sslPort=server_config["sslPort"],
         gcp=server_config["gcp"],
+        jwtToken=server_config["jwtToken"],
+        dbVersion=server_config["dbVersion"]
     )
     if server_config.get("getToken", False):
         conn.getToken(conn.createSecret())