Dev/sso login ux

graphistry/__init__.py

@@ -46,7 +46,10 @@
     ArrowFileUploader,
     PyGraphistry,
     from_igraph,
-    from_cugraph
+    from_cugraph,
+    databricks_sso_wait_for_token,
+    databricks_sso_login_init,
+    databricks_sso_login
 )
 
 from graphistry.compute import (

graphistry/pygraphistry.py

@@ -406,6 +406,8 @@ def refresh(token=None, fail_silent=False):
                 logger.debug("2. @PyGraphistry refresh :relogin")
                 if isinstance(e, TokenExpireException):
                     print("Token is expired, you need to relogin")                    
+                    PyGraphistry._config['api_token'] = None
+                    PyGraphistry._is_authenticated = False 
                 return PyGraphistry.relogin()
 
             if not fail_silent:
@@ -2448,6 +2450,72 @@ def _handle_api_response(response):
 
 
 
+    # Databricks Dashboard SSO helper functions
+
+    class DatabricksHelper():
+
+        @staticmethod
+        def sso_repeat_get_token(repeat, wait):
+            time.sleep(wait)
+
+            for i in range(repeat):
+                token = PyGraphistry.sso_get_token()
+                if token:
+                    return token
+                time.sleep(2)
+
+            return
+
+        @staticmethod
+        def databricks_sso_wait_for_token(repeat=5, wait=20):
+            from IPython.display import display, HTML
+            if not PyGraphistry.api_token():
+                msg_html = '<br /><strong> .... </strong>'
+                if not PyGraphistry.DatabricksHelper.sso_repeat_get_token(repeat, wait):
+                    msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat} .... </strong>'
+                    raise Exception(f"Failed to get token after {repeat} retries")
+                msg_html = f'{msg_html}<br /><strong>Got token</strong>'
+                display(HTML(msg_html))
+            else:
+                display(HTML('<br /><strong>Token is valid, no waiting required.</strong>'))
+
+        @staticmethod
+        def databricks_sso_login_init(wait=20):
+            from IPython.display import display, HTML, clear_output
+
+            def init_login(wait):
+                clear_output()
+                token = PyGraphistry.api_token()
+                if token:
+                    is_valid = PyGraphistry.verify_token()
+                    if not is_valid:
+                        print("***********token not valid, refresh token*****************")
+                        display(HTML('<br /><strong>Refresh token ....</strong>'))
+                        try:
+                            PyGraphistry.refresh()
+                        except Exception:
+                            pass
+
+                    else:
+                        print("Token is still valid")
+                        display(HTML('<br /><strong>Token is still valid ....</strong>'))
+
+                else:
+                    print("***********Prepare to sign in*****************")            
+
+                    msg_html = f'<br /><strong>Prepare to sign in ....</strong><br><strong>Please Login with the link appear later. Waiting for success login for {wait} seconds, please login within {wait} seconds....</strong><br /><strong>Please close the browser tab and come back to dashboard....</strong>'
+
+                    display(HTML(msg_html))
+            init_login(wait)
+
+
+        @staticmethod
+        def databricks_sso_login(server="hub.graphistry.com", org_name=None, idp_name=None, retry=5, wait=20):
+            from IPython.display import clear_output
+            clear_output()
+            if not PyGraphistry.api_token():
+                PyGraphistry.register(api=3, protocol="https", server=server, is_sso_login=True, org_name=org_name, idp_name=idp_name, sso_timeout=None, sso_opt_into_type="display")
+
 
 client_protocol_hostname = PyGraphistry.client_protocol_hostname
 store_token_creds_in_memory = PyGraphistry.store_token_creds_in_memory
@@ -2501,6 +2569,10 @@ def _handle_api_response(response):
 personal_key_secret = PyGraphistry.personal_key_secret
 switch_org = PyGraphistry.switch_org
 
+# databricks dashboard helper functions
+databricks_sso_wait_for_token = PyGraphistry.DatabricksHelper.databricks_sso_wait_for_token
+databricks_sso_login_init = PyGraphistry.DatabricksHelper.databricks_sso_login_init
+databricks_sso_login = PyGraphistry.DatabricksHelper.databricks_sso_login
 
 
 class NumpyJSONEncoder(json.JSONEncoder):