Merge pull request #256 from tigergraph/GML-1889-fix-getToken

fix getToken

Author: parkererickson-tg

pyTigerGraph/__init__.py

@@ -1,5 +1,5 @@
 from pyTigerGraph.pyTigerGraph import TigerGraphConnection
 
-__version__ = "1.7.2"
+__version__ = "1.7.3"
 
 __license__ = "Apache 2"

pyTigerGraph/pyTigerGraphAuth.py

@@ -222,52 +222,35 @@ def getToken(self, secret: str = None, setToken: bool = True, lifetime: int = No
         if self.version:
             s, m, i = self.version.split(".")
         success = False
-
-        if secret and (int(s) < 3 or (int(s) == 3 and int(m) < 5)):
+        
+        if not(secret) or self.graphname:
+            if self.graphname:
+                _json = {"graph": self.graphname}
             try:
-                # /gsql/v1/tokens endpoint only supported on version >=4.1 and replaced /requesttoken
-                _json = {"secret": secret, "graph": self.graphname}
-                if lifetime:
-                    _json["lifetime"] = str(lifetime)
-                res = requests.request("POST", self.gsUrl +
-                    "/gsql/v1/tokens", verify=False, json=_json, headers={"X-User-Agent": "pyTigerGraph"})
+                res = self._post(self.restppUrl+"/requesttoken", authMode="pwd", data=str(_json), resKey="results")
+                mainVer = 3
 
-                # if /gsql/v1/tokens endpoint doesn't exist then try old endpoint
-                if res.status_code == 404:
-                    res = requests.request("GET", self.restppUrl +
-                        "/requesttoken?secret=" + secret +
-                        ("&lifetime=" + str(lifetime) if lifetime else ""), verify=False)
-                    mainVer = 3 # Can't use _verGreaterThan4_0 to check version since you need to set a token for that
-                else:
-                    mainVer = 4
-                res = json.loads(res.text)
-
-                if not res["error"]:
-                    success = True
+            # The old endpoint doesn't exist (since on TigerGraph Ver >=4.1). Use new endpoint
+            # have to handle it in this order since _req changes the url to the new url path if first request fails
             except Exception as e:
-                raise e
-        elif not(success) and not(secret):
-            _json = {"graph": self.graphname}
-            try:
-                res = self._post(self.gsUrl +
-                    "/gsql/v1/tokens", data=_json, authMode="pwd", jsonData=True, resKey=None)
-                mainVer = 4
-                
-            # The new endpoint doesn't exist (since on TigerGraph Ver <4.1). Use old endpoint
-            except requests.exceptions.HTTPError as e:
-                if e.response.status_code == 404:    
-                    res = self._post(self.restppUrl+"/requesttoken", authMode="pwd", data=str({"graph": self.graphname}), resKey="results")
-                    mainVer = 3
-                if e.response.status_code == 400:
-                    raise TigerGraphException("Error requesting token. Check if the connection's graphname is correct.", 400)
-                else:
-                    raise e
+                try:
+                    res = self._post(self.gsUrl + "/gsql/v1/tokens",
+                                     data=_json,
+                                     authMode="pwd",
+                                     jsonData=True,
+                                     resKey=None)
+                    mainVer = 4
+                except requests.exceptions.HTTPError as e:
+                    if e.response.status_code == 404:
+                        raise TigerGraphException(
+                            "Error requesting token. Check if the connection's graphname is correct and that REST authentication is enabled.",
+                            404
+                        )
+                    else:
+                        raise e
+                pass
             success = True
-        elif not(success) and (int(s) < 3 or (int(s) == 3 and int(m) < 5)):
-            raise TigerGraphException("Cannot request a token with username/password for versions < 3.5.")
-
-
-        if not success and mainVer == 3:
+        elif secret:
             try:
                 data = {"secret": secret}
 
@@ -276,8 +259,24 @@ def getToken(self, secret: str = None, setToken: bool = True, lifetime: int = No
 
                 res = json.loads(requests.post(self.restppUrl + "/requesttoken",
                     data=json.dumps(data), verify=False).text)
-            except Exception as e:
-                raise e
+                success = True
+                mainVer = 3
+            except Exception as e: # case of less than version 3.5
+                try:
+                    res = requests.request("GET", self.restppUrl +
+                    "/requesttoken?secret=" + secret +
+                    ("&lifetime=" + str(lifetime) if lifetime else ""), verify=False)
+                    mainVer = 3 # Can't use _verGreaterThan4_0 to check version since you need to set a token for that
+                    
+                    res = json.loads(res.text)
+
+                    if not res["error"]:
+                        success = True
+                except:
+                    raise e
+        else:
+            raise TigerGraphException("Cannot request a token with username/password for versions < 3.5.")
+            
 
 
         if not res.get("error"):
@@ -511,4 +510,5 @@ def deleteToken(self, secret, token=None, skipNA=True) -> bool:
 
             return True
 
+
         raise TigerGraphException(res["message"], (res["code"] if "code" in res else None))

pyTigerGraph/pyTigerGraphBase.py

@@ -218,21 +218,24 @@ def _verify_jwt_token_support(self):
             logger.debug(f"Using auth header: {self.authHeader}") 
             version = self.getVer()
             logger.info(f"Database version: {version}")
-
+            '''
             # Check JWT support for GSQL server
             if self._versionGreaterThan4_0():
-                logger.debug(f"Attempting to get auth info with URL: {self.gsUrl + '/gsql/v1/auth/simple'}")
-                self._get(f"{self.gsUrl}/gsql/v1/auth/simple", authMode="token", resKey=None)    
+                logger.debug(f"Attempting to get auth info with URL: {self.gsUrl + '/gsql/v1/tokens/check'}")
+                res = self._post(f"{self.gsUrl}/gsql/v1/tokens/check", authMode="token", resKey=None, data={"token": self.jwtToken}, jsonData=True)
+                if "error" in res and res["error"]:
+                    raise TigerGraphException(res["message"], (res["code"] if "code" in res else None))  
             else:
                 logger.debug(f"Attempting to get auth info with URL: {self.gsUrl + '/gsqlserver/gsql/simpleauth'}")
                 self._get(f"{self.gsUrl}/gsqlserver/gsql/simpleauth", authMode="token", resKey=None)
+            '''
         except requests.exceptions.ConnectionError as e:
             logger.error(f"Connection error: {e}.")
-            raise RuntimeError(f"Connection error: {e}.") from e
+            raise TigerGraphException("Connection error: "+str(e))
         except Exception as e:
-            message = "The JWT token might be invalid or expired or DB version doesn't support JWT token. Please generate new JWT token or switch to API token or username/password."
+            message = "The JWT token might be invalid or expired or DB version doesn't support JWT token. Please generate new JWT token or switch to API token or username/password. Error: "+str(e)
             logger.error(f"Error occurred: {e}. {message}")
-            raise RuntimeError(message) from e
+            raise TigerGraphException(message) 
 
     def _locals(self, _locals: dict) -> str:
         del _locals["self"]
@@ -345,7 +348,7 @@ def _req(self, method: str, url: str, authMode: str = "token", headers: dict = N
             # This block should only be called once. When using 4.x, using port 9000 should fail so self.restppurl will change to host:14240/restpp
             # ----
             # Changes port to gsql port, adds /restpp to end to url, tries again, saves changes if successful
-            if "/restpp" not in url or self.tgCloud:
+            if self.restppPort in url and "/gsql" not in url and ("/restpp" not in url or self.tgCloud):
                 newRestppUrl = self.host + ":"+self.gsPort+"/restpp"
                 # In tgcloud /restpp can already be in the restpp url. We want to extract everything after the port or /restpp
                 if self.tgCloud:

tests/test_jwtAuth.py

@@ -4,6 +4,7 @@
 
 from pyTigerGraphUnitTest import make_connection
 from pyTigerGraph import TigerGraphConnection
+from pyTigerGraph.pyTigerGraphException import TigerGraphException
 
 
 class TestJWTTokenAuth(unittest.TestCase):
@@ -43,7 +44,7 @@ def _requestJWTToken(self):
 
 
     def _test_jwtauth_3_9(self):
-        with self.assertRaises(RuntimeError) as context:
+        with self.assertRaises(TigerGraphException) as context:
             TigerGraphConnection(
                 host=self.conn.host,
                 jwtToken="fake.JWT.Token"
@@ -79,7 +80,7 @@ def _test_jwtauth_4_1_success(self):
 
 
     def _test_jwtauth_4_1_fail(self):
-        with self.assertRaises(RuntimeError) as context:
+        with self.assertRaises(TigerGraphException) as context:
             TigerGraphConnection(
                 host=self.conn.host,
                 jwtToken="invalid.JWT.Token"

tests/test_pyTigerGraphAuth.py

@@ -5,7 +5,7 @@
 from pyTigerGraph.pyTigerGraphException import TigerGraphException
 
 
-class test_pyTigerGraphPath(unittest.TestCase):
+class test_pyTigerGraphAuth(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         cls.conn = make_connection()
@@ -61,27 +61,37 @@ def test_04_dropSecret(self):
     def test_05_getToken(self):
         res = self.conn.createSecret("secret5", True)
         token = self.conn.getToken(res["secret5"])
-        self.assertIsInstance(token, tuple)
+        if isinstance(token, str): # handle plaintext tokens from TG 3.x
+            self.assertIsInstance(token, str)
+        else:
+            self.assertIsInstance(token, tuple)
         self.conn.dropSecret("secret5")
 
+    '''
     def test_06_refreshToken(self):
         # TG 4.x does not allow refreshing tokens
-        self.conn.getToken(self.conn.createSecret())
         if self.conn._versionGreaterThan4_0(): 
             with self.assertRaises(TigerGraphException) as tge:
                 self.conn.refreshToken("secret1")
             self.assertEqual("Refreshing tokens is only supported on versions of TigerGraph <= 4.0.0.", tge.exception.message)
         else:
             res = self.conn.createSecret("secret6", True)
             token = self.conn.getToken(res["secret6"])
-            refreshed = self.conn.refreshToken(res["secret6"], token[0])
-            self.assertIsInstance(refreshed, tuple)
-            self.conn.dropSecret("secret6")
-
+            if isinstance(token, str): # handle plaintext tokens from TG 3.x
+                refreshed = self.conn.refreshToken(res["secret6"], token)
+                self.assertIsInstance(refreshed, str)
+            else:
+                refreshed = self.conn.refreshToken(res["secret6"], token[0])
+                self.assertIsInstance(refreshed, tuple)
+                self.conn.dropSecret("secret6")
+    '''
     def test_07_deleteToken(self):
         res = self.conn.createSecret("secret7", True)
         token = self.conn.getToken(res["secret7"])
-        self.assertTrue(self.conn.deleteToken(res["secret7"], token[0]))
+        if isinstance(token, str): # handle plaintext tokens from TG 3.x
+            self.assertTrue(self.conn.deleteToken(res["secret7"], token))
+        else:
+            self.assertTrue(self.conn.deleteToken(res["secret7"], token[0]))
         self.conn.dropSecret("secret7")
 
 if __name__ == '__main__':