From 89f3d5bd2f453a97ed2e55f1fe806132dd140ff3 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Thu, 17 Oct 2024 13:32:38 +0800
Subject: [PATCH 01/34] fix (typecheck): fix typecheck issue

---
 graphistry/compute/ast.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/compute/ast.py b/graphistry/compute/ast.py
index d478c23f28..22683c76dd 100644
--- a/graphistry/compute/ast.py
+++ b/graphistry/compute/ast.py
@@ -299,7 +299,7 @@ def from_json(cls, d: dict) -> 'ASTEdge':
             direction=d['direction'] if 'direction' in d else None,
             edge_match=maybe_filter_dict_from_json(d, 'edge_match'),
             hops=d['hops'] if 'hops' in d else None,
-            to_fixed_point=d['to_fixed_point'] if 'to_fixed_point' in d else None,
+            to_fixed_point=d['to_fixed_point'] if 'to_fixed_point' in d else DEFAULT_FIXED_POINT,
             source_node_match=maybe_filter_dict_from_json(d, 'source_node_match'),
             destination_node_match=maybe_filter_dict_from_json(d, 'destination_node_match'),
             source_node_query=d['source_node_query'] if 'source_node_query' in d else None,

From dbad1475296a5e1db1ad2a6eae318be75e6eec30 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Mon, 28 Oct 2024 13:25:22 +0800
Subject: [PATCH 02/34] feat: Add databricks helper class and functions

---
 graphistry/pygraphistry.py | 72 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 0925a6a744..b14262d69f 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2448,6 +2448,74 @@ 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.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(f'<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(f'<br /><strong>Refresh token ....</strong>'))
+                        try:
+                            PyGraphistry.refresh()
+                        except Exception:
+                            pass
+
+                    else:
+                        print("Token is still valid")
+                        display(HTML(f'<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
@@ -2502,6 +2570,10 @@ def _handle_api_response(response):
 switch_org = PyGraphistry.switch_org
 
 
+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):
     def default(self, obj):

From 799e93b0a8b2a93e262c33035a917a44bc0b16cb Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Mon, 28 Oct 2024 20:38:16 +0800
Subject: [PATCH 03/34] fix (refresh token, lint): Fix refresh token and lint
 issue

---
 graphistry/pygraphistry.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index b14262d69f..1395e127b9 100644
--- a/graphistry/pygraphistry.py
+++ b/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:
@@ -2475,7 +2477,7 @@ def databricks_sso_wait_for_token(repeat=5, wait=20):
                 msg_html = f'{msg_html}<br /><strong>Got token</strong>'
                 display(HTML(msg_html))
             else:
-                display(HTML(f'<br /><strong>Token is valid, no waiting required.</strong>'))
+                display(HTML('<br /><strong>Token is valid, no waiting required.</strong>'))
 
         @staticmethod
         def databricks_sso_login_init(wait=20):
@@ -2487,7 +2489,7 @@ def init_login(wait):
                     is_valid = PyGraphistry.verify_token()
                     if not is_valid:
                         print("***********token not valid, refresh token*****************")
-                        display(HTML(f'<br /><strong>Refresh token ....</strong>'))
+                        display(HTML('<br /><strong>Refresh token ....</strong>'))
                         try:
                             PyGraphistry.refresh()
                         except Exception:
@@ -2495,7 +2497,7 @@ def init_login(wait):
 
                     else:
                         print("Token is still valid")
-                        display(HTML(f'<br /><strong>Token is still valid ....</strong>'))
+                        display(HTML('<br /><strong>Token is still valid ....</strong>'))
 
                 else:
                     print("***********Prepare to sign in*****************")            
@@ -2514,9 +2516,6 @@ def databricks_sso_login(server="hub.graphistry.com", org_name=None, idp_name=No
                 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
 server = PyGraphistry.server
@@ -2569,7 +2568,7 @@ def databricks_sso_login(server="hub.graphistry.com", org_name=None, idp_name=No
 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

From e1737d84846cac187c827f4fc5efc954a18e8cd3 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Tue, 29 Oct 2024 12:19:31 +0800
Subject: [PATCH 04/34] fix (lint): fix lint issue

---
 graphistry/pygraphistry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 1395e127b9..62d3a6bb91 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2482,6 +2482,7 @@ def databricks_sso_wait_for_token(repeat=5, wait=20):
         @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()

From caa7809a4a01c07e1c5461d88663067bdad533ea Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Tue, 29 Oct 2024 12:38:19 +0800
Subject: [PATCH 05/34] fix (databricks helper): fix missing function
 definition

---
 graphistry/__init__.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/graphistry/__init__.py b/graphistry/__init__.py
index befef2c1a2..35de6f354d 100644
--- a/graphistry/__init__.py
+++ b/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 (

From 71c5eb17f219ee8eabfeaf861ba42e8c8132a745 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 30 Oct 2024 16:24:44 +0800
Subject: [PATCH 06/34] fix (databricks helpers): fix function call namespace

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 62d3a6bb91..38e777d73d 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2471,7 +2471,7 @@ 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.sso_repeat_get_token(repeat, wait):
+                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>'

From c187ead41bcafb5ee6e0706495b8365286f3e474 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 6 Nov 2024 00:02:29 +0800
Subject: [PATCH 07/34] wip (sso logic): Add SSO flow token verification in
 plot function, which originally just a refresh()

---
 graphistry/PlotterBase.py  |  14 +++-
 graphistry/__init__.py     |   5 +-
 graphistry/pygraphistry.py | 151 +++++++++++++++++++++++--------------
 3 files changed, 108 insertions(+), 62 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index d56d07f6b6..3f8daff49e 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -1403,7 +1403,19 @@ def plot(
             info = PyGraphistry._etl1(dataset)
         elif api_version == 3:
             logger.debug("3. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
-            PyGraphistry.refresh()
+            
+            if not PyGraphistry.api_token() and PyGraphistry.sso_state():  # if it is sso login
+                if in_ipython() or in_databricks() or PyGraphistry._config["sso_opt_into_type"] == 'display':
+                    PyGraphistry.sso_wait_for_token_html_display()
+                    if not PyGraphistry.sso_verify_token_html():
+                        PyGraphistry.sso_wait_for_token_html_display()
+                        PyGraphistry.refresh()
+                else:
+                    PyGraphistry.sso_repeat_get_token()
+                
+                
+            else:  # if not sso mode, just refresh to make sure token is valid
+                PyGraphistry.refresh()
             logger.debug("4. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
 
             dataset = self._plot_dispatch(g, n, name, description, 'arrow', self._style, memoize)
diff --git a/graphistry/__init__.py b/graphistry/__init__.py
index 35de6f354d..d7af05f064 100644
--- a/graphistry/__init__.py
+++ b/graphistry/__init__.py
@@ -47,9 +47,8 @@
     PyGraphistry,
     from_igraph,
     from_cugraph,
-    databricks_sso_wait_for_token,
-    databricks_sso_login_init,
-    databricks_sso_login
+    sso_wait_for_token_html_display,
+    register_databricks_sso,
 )
 
 from graphistry.compute import (
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 38e777d73d..9f015776c9 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -135,6 +135,12 @@ def __reset_token_creds_in_memory():
         PyGraphistry._config["api_key"] = None
         PyGraphistry._is_authenticated = False
 
+    @staticmethod
+    def __reset_sso_variables_in_memory():
+        """Reset the sso related variable in memory, used when switching hosts, switching register method"""
+
+        PyGraphistry._config["sso_state"] = None
+        PyGraphistry._config["sso_opt_into_type"] = None
 
 
     @staticmethod
@@ -572,6 +578,11 @@ def certificate_validation(value=None):
     def set_bolt_driver(driver=None):
         PyGraphistry._config["bolt_driver"] = bolt_util.to_bolt_driver(driver)
 
+    @staticmethod
+    def set_sso_opt_into_type(value: Optional[str]):
+        PyGraphistry._config["sso_opt_into_type"] = value
+
+
     @staticmethod
     def register(
         key: Optional[str] = None,
@@ -709,6 +720,8 @@ def register(
         PyGraphistry.set_bolt_driver(bolt)
         # Reset token creds
         PyGraphistry.__reset_token_creds_in_memory()
+        # Reset sso_state in memory
+        PyGraphistry.__reset_sso_variables_in_memory()
 
         if not (username is None) and not (password is None):
             PyGraphistry.login(username, password, org_name)
@@ -2448,73 +2461,94 @@ def _handle_api_response(response):
             logger.error('Error: %s', response, exc_info=True)
             raise Exception("Unknown Error")
 
+    @staticmethod
+    def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
+        """ Function to repeatly call to obtain the jwt token after sso login
+        Function to obtain the JWT token after SSO login
+        """
+        
+        for _ in range(repeat):
+            token = PyGraphistry.sso_get_token()
+            if token:
+                return token
+            time.sleep(wait)
 
+        return
+    
+    @staticmethod
+    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5):
+        """Get the JWT token for SSO login and display corresponding message in HTML
+        Get the JWT token for SSO login and display corresponding message in HTML
+        """
+        from IPython.display import display, HTML
+        if not PyGraphistry.api_token():
+            msg_html = '<br /><strong> .... </strong>'
+            if not PyGraphistry.sso_repeat_get_token(repeat, wait):
+                msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
+                raise Exception(f"Failed to get token after {repeat*wait} seconds ....")
+            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>'))
 
-    # 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))
+    @staticmethod
+    def sso_verify_token_html(
+        repeat: int = 20,
+        wait: int = 5,
+    ) -> bool:
+        from IPython.display import display, HTML, clear_output
+
+        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:
-                display(HTML('<br /><strong>Token is valid, no waiting required.</strong>'))
+                print("Token is still valid")
+                display(HTML('<br /><strong>Token is still valid ....</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("***********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 {repeat*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))
 
-                    else:
-                        print("Token is still valid")
-                        display(HTML('<br /><strong>Token is still valid ....</strong>'))
+            return False
 
-                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>'
+        return True
 
-                    display(HTML(msg_html))
-            init_login(wait)
-            
 
+    # Databricks Dashboard SSO helper functions
+    class DatabricksHelper():
+        """Helper class for databricks.
+
+        **Helper class to improve the sso login flow**
+
+        """
         @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()
+        def register_databricks_sso(
+            server: Optional[str] = None,
+            org_name: Optional[str] = None,
+            idp_name: Optional[str] = None,
+            **kwargs
+        ):
             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")
+            
+
+        # @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
@@ -2570,9 +2604,10 @@ def databricks_sso_login(server="hub.graphistry.com", org_name=None, idp_name=No
 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
+sso_wait_for_token_html_display = PyGraphistry.sso_wait_for_token_html_display
+sso_verify_token_html = PyGraphistry.sso_verify_token_html
+sso_repeat_get_token = PyGraphistry.sso_repeat_get_token
+register_databricks_sso = PyGraphistry.DatabricksHelper.register_databricks_sso
 
 
 class NumpyJSONEncoder(json.JSONEncoder):

From 412e944009a10c726a892e15f6fb8eae00207d13 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 6 Nov 2024 10:12:08 +0800
Subject: [PATCH 08/34] wip (sso): fix whether to throw exception when fail to
 get token

---
 graphistry/pygraphistry.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 9f015776c9..192dc4b2fb 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2476,7 +2476,7 @@ def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
         return
     
     @staticmethod
-    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5):
+    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silently: bool = True):
         """Get the JWT token for SSO login and display corresponding message in HTML
         Get the JWT token for SSO login and display corresponding message in HTML
         """
@@ -2485,7 +2485,8 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5):
             msg_html = '<br /><strong> .... </strong>'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
                 msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
-                raise Exception(f"Failed to get token after {repeat*wait} seconds ....")
+                if not fail_silently:
+                    raise Exception(f"Failed to get token after {repeat*wait} seconds ....")
             msg_html = f'{msg_html}<br /><strong>Got token</strong>'
             display(HTML(msg_html))
         else:

From 9283e8ad582eecf1541ae23209a8953be5c10064 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 6 Nov 2024 10:31:20 +0800
Subject: [PATCH 09/34] wip (sso): fix message display logic

---
 graphistry/pygraphistry.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 192dc4b2fb..a2982bd393 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2476,7 +2476,7 @@ def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
         return
     
     @staticmethod
-    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silently: bool = True):
+    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silently: bool = False):
         """Get the JWT token for SSO login and display corresponding message in HTML
         Get the JWT token for SSO login and display corresponding message in HTML
         """
@@ -2486,7 +2486,12 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
                 msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
                 if not fail_silently:
-                    raise Exception(f"Failed to get token after {repeat*wait} seconds ....")
+                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the process") 
+                else:
+                    msg_html = f'{msg_html}<br /><strong>Got token</strong>'
+                    display(HTML(msg_html))
+                    return
+
             msg_html = f'{msg_html}<br /><strong>Got token</strong>'
             display(HTML(msg_html))
         else:

From be26deeb1487dbc247a5721192aeecfe0c416c4f Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 6 Nov 2024 11:04:35 +0800
Subject: [PATCH 10/34] fix (sso): fix flow

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index a2982bd393..48ac13f96e 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2486,7 +2486,7 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
                 msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
                 if not fail_silently:
-                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the process") 
+                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") 
                 else:
                     msg_html = f'{msg_html}<br /><strong>Got token</strong>'
                     display(HTML(msg_html))

From 5ea4841889d2163d375365f609c3802fdc882898 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 6 Nov 2024 22:49:22 +0800
Subject: [PATCH 11/34] fix (sso): fix display message flow

---
 graphistry/pygraphistry.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 48ac13f96e..c56049fd8c 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -422,7 +422,7 @@ def refresh(token=None, fail_silent=False):
 
     @staticmethod
     def verify_token(token=None, fail_silent=False) -> bool:
-        """Return True iff current or provided token is still valid"""
+        """Return True if current or provided token is still valid"""
         using_self_token = token is None
         try:
             logger.debug("JWT refresh")
@@ -2476,7 +2476,7 @@ def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
         return
     
     @staticmethod
-    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silently: bool = False):
+    def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False):
         """Get the JWT token for SSO login and display corresponding message in HTML
         Get the JWT token for SSO login and display corresponding message in HTML
         """
@@ -2485,7 +2485,7 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
             msg_html = '<br /><strong> .... </strong>'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
                 msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
-                if not fail_silently:
+                if not fail_silent:
                     raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") 
                 else:
                     msg_html = f'{msg_html}<br /><strong>Got token</strong>'
@@ -2506,6 +2506,7 @@ def sso_verify_token_html(
         from IPython.display import display, HTML, clear_output
 
         clear_output()
+        required_login = False
         token = PyGraphistry.api_token()
         if token:
             is_valid = PyGraphistry.verify_token()
@@ -2515,20 +2516,22 @@ def sso_verify_token_html(
                 try:
                     PyGraphistry.refresh()
                 except Exception:
-                    pass
+                    required_login = True
 
             else:
                 print("Token is still valid")
                 display(HTML('<br /><strong>Token is still valid ....</strong>'))
 
         else:
+            required_login = True
+
+        if required_login:
             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 {repeat*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))
 
-            return False
 
-        return True
+        return not required_login
 
 
     # Databricks Dashboard SSO helper functions

From 8a482b5abb7f2bff0bc75ef9c6b83d98723b3526 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Thu, 7 Nov 2024 21:47:10 +0800
Subject: [PATCH 12/34] fix (sso): Graph is not shown in dashboard due to only
 the first HTML will be displayed.

---
 graphistry/PlotterBase.py  |  4 ++--
 graphistry/pygraphistry.py | 47 ++++++++++++++++++++++++++++++++------
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index 3f8daff49e..4deb76c6ff 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -1406,9 +1406,9 @@ def plot(
             
             if not PyGraphistry.api_token() and PyGraphistry.sso_state():  # if it is sso login
                 if in_ipython() or in_databricks() or PyGraphistry._config["sso_opt_into_type"] == 'display':
-                    PyGraphistry.sso_wait_for_token_html_display()
+                    PyGraphistry.sso_wait_for_token_text_display()
                     if not PyGraphistry.sso_verify_token_html():
-                        PyGraphistry.sso_wait_for_token_html_display()
+                        PyGraphistry.sso_wait_for_token_text_display()
                         PyGraphistry.refresh()
                 else:
                     PyGraphistry.sso_repeat_get_token()
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index c56049fd8c..43dbb47ba2 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -277,7 +277,7 @@ def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type):
         if in_ipython() or in_databricks() or sso_opt_into_type == 'display':  # If run in notebook, just display the HTML
             # from IPython.core.display import HTML
             from IPython.display import display, HTML
-            display(HTML(f'<a href="{auth_url}" target="_blank">Login SSO</a>'))
+            display(HTML(f'<a href="{auth_url}" target="_blank">Login SSO</a><br /><span>Please click the above URL to open browser to login</span>'))
             print("Please click the above URL to open browser to login")
             print(f"If you cannot see the URL, please open browser, browse to this URL: {auth_url}")
             print("Please close browser tab after SSO login to back to notebook")
@@ -2474,7 +2474,36 @@ def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
             time.sleep(wait)
 
         return
-    
+
+    @staticmethod
+    def sso_wait_for_token_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False, display_mode: str = 'text'):
+        if display_mode == 'html':
+            PyGraphistry.sso_wait_for_token_html_display(repeat, wait, fail_silent)
+        else:
+            PyGraphistry.sso_wait_for_token_text_display(repeat, wait, fail_silent)
+
+    @staticmethod
+    def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False):
+        """Get the JWT token for SSO login and display corresponding message in text
+        Get the JWT token for SSO login and display corresponding message in text
+        """
+        if not PyGraphistry.api_token():
+            msg_text = '....'
+            if not PyGraphistry.sso_repeat_get_token(repeat, wait):
+                msg_text = f'{msg_text}\nFailed to get token after {repeat*wait} seconds ....'
+                if not fail_silent:
+                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") 
+                else:
+                    msg_text = f'{msg_text}\nGot token'
+                    print(msg_text)
+                    return
+
+            msg_text = f'{msg_text}\nGot token'
+            print(msg_text)
+        else:
+            print('Token is valid, no waiting required.')
+
+
     @staticmethod
     def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False):
         """Get the JWT token for SSO login and display corresponding message in HTML
@@ -2499,20 +2528,23 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
 
 
     @staticmethod
-    def sso_verify_token_html(
+    def sso_verify_token_display(
         repeat: int = 20,
         wait: int = 5,
+        display_mode: str = 'text'
     ) -> bool:
-        from IPython.display import display, HTML, clear_output
+        if display_mode == 'html':
+            from IPython.display import display, HTML, clear_output
+            clear_output()
 
-        clear_output()
         required_login = False
         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>'))
+                if display_mode == 'html':
+                    display(HTML('<br /><strong>Refresh token ....</strong>'))
                 try:
                     PyGraphistry.refresh()
                 except Exception:
@@ -2520,7 +2552,8 @@ def sso_verify_token_html(
 
             else:
                 print("Token is still valid")
-                display(HTML('<br /><strong>Token is still valid ....</strong>'))
+                if display_mode == 'html':
+                    display(HTML('<br /><strong>Token is still valid ....</strong>'))
 
         else:
             required_login = True

From ab2794ca49bfa6df12f32dee6b1d4cfb6688127a Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Thu, 7 Nov 2024 21:58:45 +0800
Subject: [PATCH 13/34] wip (sso): fix syntax error

---
 graphistry/PlotterBase.py  | 4 ++--
 graphistry/__init__.py     | 2 +-
 graphistry/pygraphistry.py | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index 4deb76c6ff..8d4913b03c 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -1407,8 +1407,8 @@ def plot(
             if not PyGraphistry.api_token() and PyGraphistry.sso_state():  # if it is sso login
                 if in_ipython() or in_databricks() or PyGraphistry._config["sso_opt_into_type"] == 'display':
                     PyGraphistry.sso_wait_for_token_text_display()
-                    if not PyGraphistry.sso_verify_token_html():
-                        PyGraphistry.sso_wait_for_token_text_display()
+                    if not PyGraphistry.sso_verify_token_display():
+                        PyGraphistry.sso_wait_for_token_display()
                         PyGraphistry.refresh()
                 else:
                     PyGraphistry.sso_repeat_get_token()
diff --git a/graphistry/__init__.py b/graphistry/__init__.py
index d7af05f064..f5e98be771 100644
--- a/graphistry/__init__.py
+++ b/graphistry/__init__.py
@@ -47,7 +47,7 @@
     PyGraphistry,
     from_igraph,
     from_cugraph,
-    sso_wait_for_token_html_display,
+    sso_wait_for_token_display,
     register_databricks_sso,
 )
 
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 43dbb47ba2..c49c38fcde 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2646,8 +2646,8 @@ def register_databricks_sso(
 switch_org = PyGraphistry.switch_org
 
 # databricks dashboard helper functions
-sso_wait_for_token_html_display = PyGraphistry.sso_wait_for_token_html_display
-sso_verify_token_html = PyGraphistry.sso_verify_token_html
+sso_wait_for_token_display = PyGraphistry.sso_wait_for_token_display
+sso_verify_token_display = PyGraphistry.sso_verify_token_display
 sso_repeat_get_token = PyGraphistry.sso_repeat_get_token
 register_databricks_sso = PyGraphistry.DatabricksHelper.register_databricks_sso
 

From 5b35284328eaeae968245876a857afd4c3403701 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 8 Nov 2024 08:47:28 +0800
Subject: [PATCH 14/34] wip (sso): debugging code

---
 graphistry/pygraphistry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index c49c38fcde..48b8961b6b 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -438,6 +438,7 @@ def verify_token(token=None, fail_silent=False) -> bool:
                 PyGraphistry._is_authenticated = ok
             return ok
         except Exception as e:
+            raise e
             if not fail_silent:
                 util.error("Failed to verify token: %s" % str(e))
             return False

From 8355b0f04b1baf97bba4ca1b2992b14ef86b2a77 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 8 Nov 2024 09:22:16 +0800
Subject: [PATCH 15/34] wip (debug): add debug print

---
 graphistry/pygraphistry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 48b8961b6b..c7434a98e8 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2542,6 +2542,7 @@ def sso_verify_token_display(
         token = PyGraphistry.api_token()
         if token:
             is_valid = PyGraphistry.verify_token()
+            print(f"is_valid : {is_valid")
             if not is_valid:
                 print("***********token not valid, refresh token*****************")
                 if display_mode == 'html':

From 1a2ab33ad49f8c4a6d0c4d148ba1c91c7b7b6999 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 8 Nov 2024 09:24:41 +0800
Subject: [PATCH 16/34] wip (fix): syntax error

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index c7434a98e8..0ba30eba5c 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2542,7 +2542,7 @@ def sso_verify_token_display(
         token = PyGraphistry.api_token()
         if token:
             is_valid = PyGraphistry.verify_token()
-            print(f"is_valid : {is_valid")
+            print(f"is_valid : {is_valid}")
             if not is_valid:
                 print("***********token not valid, refresh token*****************")
                 if display_mode == 'html':

From a53664886f5f20ccc8196d3974c2ac9cdb22e20a Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 8 Nov 2024 10:58:01 +0800
Subject: [PATCH 17/34] wip (debug):debugging print

---
 graphistry/pygraphistry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 0ba30eba5c..cfd580e89b 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -424,6 +424,7 @@ def refresh(token=None, fail_silent=False):
     def verify_token(token=None, fail_silent=False) -> bool:
         """Return True if current or provided token is still valid"""
         using_self_token = token is None
+        print(f"token to verify: {PyGraphistry.api_token() if using_self_token else token}")
         try:
             logger.debug("JWT refresh")
             if using_self_token:

From dd29430693e527b3ac3a17b45e85169c7067bc4c Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Fri, 8 Nov 2024 11:06:42 +0800
Subject: [PATCH 18/34] wip (debug): debugging print

---
 graphistry/arrow_uploader.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index c4a4127643..a0e1bc44b2 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -393,6 +393,7 @@ def verify(self, token=None) -> bool:
             verify=self.certificate_validation,
             json={'token': token})
         log_requests_error(out)
+        return out.text
         return out.status_code == requests.codes.ok
 
     def create_dataset(self, json, validate: bool = True):  # noqa: F811

From 36fcfea604c18774ba3bc895e69fa7b6ce5ae169 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 00:14:44 +0800
Subject: [PATCH 19/34] wip (sso): attempt to display error in iframe

---
 graphistry/PlotterBase.py |  6 +++---
 graphistry/util.py        | 10 +++++++---
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index 8d4913b03c..f242410f55 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -23,7 +23,7 @@
 )
 from .util import (
     error, hash_pdf, in_ipython, in_databricks, make_iframe, random_string, warn,
-    cache_coercion, cache_coercion_helper, WeakValueWrapper
+    cache_coercion, cache_coercion_helper, WeakValueWrapper, make_iframe_srcdoc
 )
 
 from .bolt_util import (
@@ -1408,8 +1408,8 @@ def plot(
                 if in_ipython() or in_databricks() or PyGraphistry._config["sso_opt_into_type"] == 'display':
                     PyGraphistry.sso_wait_for_token_text_display()
                     if not PyGraphistry.sso_verify_token_display():
-                        PyGraphistry.sso_wait_for_token_display()
-                        PyGraphistry.refresh()
+                        msg_html = "<strong>Invalid token due to login timeout</strong>"
+                        extra_html = f"{extra_html} {make_iframe_srcdoc(msg_html)}"
                 else:
                     PyGraphistry.sso_repeat_get_token()
                 
diff --git a/graphistry/util.py b/graphistry/util.py
index c2c47996f1..e4b2b0a3c2 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -10,7 +10,7 @@
 import uuid
 import warnings
 from functools import lru_cache
-from typing import Any
+from typing import Any, Optional
 from collections import UserDict
 
 from .constants import VERBOSE, CACHE_COERCION_SIZE, TRACE
@@ -172,8 +172,12 @@ def check_set_memoize(
         weakref[hashed] = w
     return False
 
+def make_iframe_srcdoc(srcdoc: str):
+    srcdoc = 'srcdoc="{srcdoc}" onload="this.removeAttribute(\'srcdoc\')"'
+    return srcdoc
 
-def make_iframe(url, height, extra_html="", override_html_style=None):
+
+def make_iframe(url, height, extra_html="", override_html_style=None, srcdoc: Optional[str] = None):
     id = uuid.uuid4()
 
     height_str = (
@@ -203,7 +207,7 @@ def make_iframe(url, height, extra_html="", override_html_style=None):
         )
 
     iframe = """
-            <iframe id="%s" src="%s"
+            <iframe id="%s" src="%s" {srcdoc}
                     allowfullscreen="true" webkitallowfullscreen="true" mozallowfullscreen="true"
                     oallowfullscreen="true" msallowfullscreen="true"
                     style="%s"

From 4820620d11cd174f7e75b6d5ad726092c98cccfd Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 00:20:19 +0800
Subject: [PATCH 20/34] wip (databricks): return error msg html in the case of
 plot with invalid token

---
 graphistry/PlotterBase.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index f242410f55..e58286cdcd 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -1409,7 +1409,8 @@ def plot(
                     PyGraphistry.sso_wait_for_token_text_display()
                     if not PyGraphistry.sso_verify_token_display():
                         msg_html = "<strong>Invalid token due to login timeout</strong>"
-                        extra_html = f"{extra_html} {make_iframe_srcdoc(msg_html)}"
+                        # extra_html = f"{extra_html} {make_iframe_srcdoc(msg_html)}"
+                        return HTML(msg_html)
                 else:
                     PyGraphistry.sso_repeat_get_token()
                 

From 2d6ad0bd2096bca86d2a2b0d28dfc9a533dcf786 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 00:57:20 +0800
Subject: [PATCH 21/34] wip (sso): display html message before throw exception

---
 graphistry/pygraphistry.py | 7 +++++--
 graphistry/util.py         | 8 ++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index cfd580e89b..5cbb41d062 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -15,7 +15,7 @@
 from . import util
 from . import bolt_util
 from .plotter import Plotter
-from .util import in_databricks, setup_logger, in_ipython, make_iframe
+from .util import in_databricks, setup_logger, in_ipython, make_iframe, display_message_html
 from .exceptions import SsoRetrieveTokenTimeoutException, TokenExpireException
 
 from .messages import (
@@ -2494,7 +2494,10 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
                 msg_text = f'{msg_text}\nFailed to get token after {repeat*wait} seconds ....'
                 if not fail_silent:
-                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") 
+                    msg = f"Failed to get token after {repeat*wait} seconds. Please re-run the login process"
+                    if in_ipython() or in_databricks or PyGraphistry.set_sso_opt_into_type == "display":
+                        display_message_html("<strong>{msg}</strong>")
+                    raise Exception(msg) 
                 else:
                     msg_text = f'{msg_text}\nGot token'
                     print(msg_text)
diff --git a/graphistry/util.py b/graphistry/util.py
index e4b2b0a3c2..ba47f55eb4 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -172,6 +172,14 @@ def check_set_memoize(
         weakref[hashed] = w
     return False
 
+def display_message_html(message: str, cleared: Optional[bool] = False):
+    from IPython.display import display, HTML, clear_output
+
+    if cleared:
+        clear_output()
+
+    display(HTML(message))
+
 def make_iframe_srcdoc(srcdoc: str):
     srcdoc = 'srcdoc="{srcdoc}" onload="this.removeAttribute(\'srcdoc\')"'
     return srcdoc

From 05d0dff9c14b03ac23addc0e1adfa091a82bff1a Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 01:01:46 +0800
Subject: [PATCH 22/34] wip(sso): fix error

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 5cbb41d062..69586eb757 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2496,7 +2496,7 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
                 if not fail_silent:
                     msg = f"Failed to get token after {repeat*wait} seconds. Please re-run the login process"
                     if in_ipython() or in_databricks or PyGraphistry.set_sso_opt_into_type == "display":
-                        display_message_html("<strong>{msg}</strong>")
+                        display_message_html(f"<strong>{msg}</strong>")
                     raise Exception(msg) 
                 else:
                     msg_text = f'{msg_text}\nGot token'

From fd46826b9d077a4bf44daec32889d6ac52df77e9 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 20:05:35 +0800
Subject: [PATCH 23/34] wip (sso): fix syntax error

---
 graphistry/PlotterBase.py    | 1 +
 graphistry/arrow_uploader.py | 1 -
 graphistry/pygraphistry.py   | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index e58286cdcd..71bebbe453 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -1408,6 +1408,7 @@ def plot(
                 if in_ipython() or in_databricks() or PyGraphistry._config["sso_opt_into_type"] == 'display':
                     PyGraphistry.sso_wait_for_token_text_display()
                     if not PyGraphistry.sso_verify_token_display():
+                        from IPython.core.display import HTML
                         msg_html = "<strong>Invalid token due to login timeout</strong>"
                         # extra_html = f"{extra_html} {make_iframe_srcdoc(msg_html)}"
                         return HTML(msg_html)
diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py
index a0e1bc44b2..c4a4127643 100644
--- a/graphistry/arrow_uploader.py
+++ b/graphistry/arrow_uploader.py
@@ -393,7 +393,6 @@ def verify(self, token=None) -> bool:
             verify=self.certificate_validation,
             json={'token': token})
         log_requests_error(out)
-        return out.text
         return out.status_code == requests.codes.ok
 
     def create_dataset(self, json, validate: bool = True):  # noqa: F811
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 69586eb757..5f9aefe5b8 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2495,7 +2495,7 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
                 msg_text = f'{msg_text}\nFailed to get token after {repeat*wait} seconds ....'
                 if not fail_silent:
                     msg = f"Failed to get token after {repeat*wait} seconds. Please re-run the login process"
-                    if in_ipython() or in_databricks or PyGraphistry.set_sso_opt_into_type == "display":
+                    if in_ipython() or in_databricks() or PyGraphistry.set_sso_opt_into_type == "display":
                         display_message_html(f"<strong>{msg}</strong>")
                     raise Exception(msg) 
                 else:

From 4816fd67dc599f8463186f96fb9f61d4e2daf026 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Sat, 9 Nov 2024 22:23:45 +0800
Subject: [PATCH 24/34] wip (sso): clean up code

---
 graphistry/PlotterBase.py  | 3 +--
 graphistry/pygraphistry.py | 4 +---
 graphistry/util.py         | 4 ----
 3 files changed, 2 insertions(+), 9 deletions(-)

diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py
index 71bebbe453..577fde40b4 100644
--- a/graphistry/PlotterBase.py
+++ b/graphistry/PlotterBase.py
@@ -23,7 +23,7 @@
 )
 from .util import (
     error, hash_pdf, in_ipython, in_databricks, make_iframe, random_string, warn,
-    cache_coercion, cache_coercion_helper, WeakValueWrapper, make_iframe_srcdoc
+    cache_coercion, cache_coercion_helper, WeakValueWrapper
 )
 
 from .bolt_util import (
@@ -1410,7 +1410,6 @@ def plot(
                     if not PyGraphistry.sso_verify_token_display():
                         from IPython.core.display import HTML
                         msg_html = "<strong>Invalid token due to login timeout</strong>"
-                        # extra_html = f"{extra_html} {make_iframe_srcdoc(msg_html)}"
                         return HTML(msg_html)
                 else:
                     PyGraphistry.sso_repeat_get_token()
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 5f9aefe5b8..b622722ec7 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -424,7 +424,6 @@ def refresh(token=None, fail_silent=False):
     def verify_token(token=None, fail_silent=False) -> bool:
         """Return True if current or provided token is still valid"""
         using_self_token = token is None
-        print(f"token to verify: {PyGraphistry.api_token() if using_self_token else token}")
         try:
             logger.debug("JWT refresh")
             if using_self_token:
@@ -439,7 +438,6 @@ def verify_token(token=None, fail_silent=False) -> bool:
                 PyGraphistry._is_authenticated = ok
             return ok
         except Exception as e:
-            raise e
             if not fail_silent:
                 util.error("Failed to verify token: %s" % str(e))
             return False
@@ -722,7 +720,7 @@ def register(
         PyGraphistry.set_bolt_driver(bolt)
         # Reset token creds
         PyGraphistry.__reset_token_creds_in_memory()
-        # Reset sso_state in memory
+        # Reset sso related variables in memory
         PyGraphistry.__reset_sso_variables_in_memory()
 
         if not (username is None) and not (password is None):
diff --git a/graphistry/util.py b/graphistry/util.py
index ba47f55eb4..b4dee6acb9 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -180,10 +180,6 @@ def display_message_html(message: str, cleared: Optional[bool] = False):
 
     display(HTML(message))
 
-def make_iframe_srcdoc(srcdoc: str):
-    srcdoc = 'srcdoc="{srcdoc}" onload="this.removeAttribute(\'srcdoc\')"'
-    return srcdoc
-
 
 def make_iframe(url, height, extra_html="", override_html_style=None, srcdoc: Optional[str] = None):
     id = uuid.uuid4()

From bff245971d4f766fde5b8915dfce8f27cadfa050 Mon Sep 17 00:00:00 2001
From: Simon Abizmil <simon@simonabizmil.com>
Date: Mon, 11 Nov 2024 12:32:48 -0800
Subject: [PATCH 25/34] Update docstr for new/updated methods

---
 graphistry/pygraphistry.py | 141 ++++++++++++++++++++++++++++++++++---
 1 file changed, 130 insertions(+), 11 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index b622722ec7..29568f997a 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2336,7 +2336,22 @@ def layout_settings(
 
     @staticmethod
     def org_name(value=None):
-        """Set or get the org_name when register/login.
+        """Set or get the organization name during registration or login.
+
+        :param value: The organization name to set. If None, the current organization name is returned.
+        :type value: Optional[str]
+        :return: The current organization name if value is None, otherwise None.
+        :rtype: Optional[str]
+
+        **Example: Setting the organization name**
+            ::
+                import graphistry
+                graphistry.org_name("my_org_name")
+
+        **Example: Getting the organization name**
+            ::
+                import graphistry
+                org_name = graphistry.org_name()
         """
 
         if value is None:
@@ -2354,7 +2369,22 @@ def org_name(value=None):
 
     @staticmethod
     def idp_name(value=None):
-        """Set or get the idp_name when register/login.
+        """Set or get the IDP (Identity Provider) name during registration or login.
+
+        :param value: The IDP name to set. If None, the current IDP name is returned.
+        :type value: Optional[str]
+        :return: The current IDP name if value is None, otherwise None.
+        :rtype: Optional[str]
+
+        **Example: Setting the IDP name**
+            ::
+                import graphistry
+                graphistry.idp_name("my_idp_name")
+
+        **Example: Getting the IDP name**
+            ::
+                import graphistry
+                idp_name = graphistry.idp_name()
         """
 
         if value is None:
@@ -2369,7 +2399,22 @@ def idp_name(value=None):
 
     @staticmethod
     def sso_state(value=None):
-        """Set or get the sso_state when register/sso login.
+        """Set or get the SSO state during registration or SSO login.
+
+        :param value: The SSO state to set. If None, the current SSO state is returned.
+        :type value: Optional[str]
+        :return: The current SSO state if value is None, otherwise None.
+        :rtype: Optional[str]
+
+        **Example: Setting the SSO state**
+            ::
+                import graphistry
+                graphistry.sso_state("my_sso_state")
+
+        **Example: Getting the SSO state**
+            ::
+                import graphistry
+                sso_state = graphistry.sso_state()
         """
 
         if value is None:
@@ -2405,7 +2450,22 @@ def scene_settings(
 
     @staticmethod
     def personal_key_id(value: Optional[str] = None):
-        """Set or get the personal_key_id when register.
+        """Set or get the personal_key_id during registration.
+
+        :param value: The personal key ID to set. If None, the current personal key ID is returned.
+        :type value: Optional[str]
+        :return: The current personal key ID if value is None, otherwise None.
+        :rtype: Optional[str]
+
+        **Example: Setting the personal key ID**
+            ::
+                import graphistry
+                graphistry.personal_key_id("my_personal_key_id")
+
+        **Example: Getting the personal key ID**
+            ::
+                import graphistry
+                key_id = graphistry.personal_key_id()
         """
 
         if value is None:
@@ -2419,7 +2479,22 @@ def personal_key_id(value: Optional[str] = None):
 
     @staticmethod
     def personal_key_secret(value: Optional[str] = None):
-        """Set or get the personal_key_secret when register.
+        """Set or get the personal_key_secret during registration.
+
+        :param value: The personal key secret to set. If None, the current personal key secret is returned.
+        :type value: Optional[str]
+        :return: The current personal key secret if value is None, otherwise None.
+        :rtype: Optional[str]
+
+        **Example: Setting the personal key secret**
+            ::
+                import graphistry
+                graphistry.personal_key_secret("my_personal_key_secret")
+
+        **Example: Getting the personal key secret**
+            ::
+                import graphistry
+                secret = graphistry.personal_key_secret()
         """
 
         if value is None:
@@ -2463,8 +2538,24 @@ def _handle_api_response(response):
 
     @staticmethod
     def sso_repeat_get_token(repeat: int = 20, wait: int = 5):
-        """ Function to repeatly call to obtain the jwt token after sso login
-        Function to obtain the JWT token after SSO login
+        """Repeatedly call to obtain the JWT token after SSO login.
+
+        :param repeat: Number of times to attempt obtaining the token, defaults to 20
+        :type repeat: int, optional
+        :param wait: Number of seconds to wait between attempts, defaults to 5
+        :type wait: int, optional
+        :return: The obtained JWT token or None if unsuccessful
+        :rtype: Optional[str]
+
+        **Example:**
+
+        ::
+
+            token = PyGraphistry.sso_repeat_get_token(repeat=10, wait=2)
+            if token:
+                print("Token obtained:", token)
+            else:
+                print("Failed to obtain token")
         """
         
         for _ in range(repeat):
@@ -2484,8 +2575,22 @@ def sso_wait_for_token_display(repeat: int = 20, wait: int = 5, fail_silent: boo
 
     @staticmethod
     def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False):
-        """Get the JWT token for SSO login and display corresponding message in text
-        Get the JWT token for SSO login and display corresponding message in text
+        """Get the JWT token for SSO login and display the corresponding message in text.
+
+        This method attempts to obtain the JWT token for SSO login and displays the result as a text message.
+
+        :param repeat: Number of times to attempt obtaining the token, defaults to 20
+        :type repeat: int, optional
+        :param wait: Number of seconds to wait between attempts, defaults to 5
+        :type wait: int, optional
+        :param fail_silent: Whether to suppress exceptions on failure, defaults to False
+        :type fail_silent: bool, optional
+
+        **Example:**
+
+        ::
+
+            PyGraphistry.sso_wait_for_token_text_display(repeat=10, wait=2, fail_silent=True)
         """
         if not PyGraphistry.api_token():
             msg_text = '....'
@@ -2509,8 +2614,22 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
 
     @staticmethod
     def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent: bool = False):
-        """Get the JWT token for SSO login and display corresponding message in HTML
-        Get the JWT token for SSO login and display corresponding message in HTML
+        """Get the JWT token for SSO login and display the corresponding message in HTML.
+
+        This method attempts to obtain the JWT token for SSO login and displays the result as an HTML message.
+
+        :param repeat: Number of times to attempt obtaining the token, defaults to 20
+        :type repeat: int, optional
+        :param wait: Number of seconds to wait between attempts, defaults to 5
+        :type wait: int, optional
+        :param fail_silent: Whether to suppress exceptions on failure, defaults to False
+        :type fail_silent: bool, optional
+
+        **Example:**
+
+        ::
+
+            PyGraphistry.sso_wait_for_token_html_display(repeat=10, wait=2, fail_silent=True)
         """
         from IPython.display import display, HTML
         if not PyGraphistry.api_token():

From 4b0b129f3f6f56a2ae95e5316fcce0ab111a775a Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 13 Nov 2024 22:08:29 +0800
Subject: [PATCH 26/34] fix (lint): fix lint issue

---
 graphistry/pygraphistry.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 29568f997a..7c41ee3d7f 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2595,9 +2595,9 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
         if not PyGraphistry.api_token():
             msg_text = '....'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
-                msg_text = f'{msg_text}\nFailed to get token after {repeat*wait} seconds ....'
+                msg_text = f'{msg_text}\nFailed to get token after {repeat * wait} seconds ....'
                 if not fail_silent:
-                    msg = f"Failed to get token after {repeat*wait} seconds. Please re-run the login process"
+                    msg = f"Failed to get token after {repeat * wait} seconds. Please re-run the login process"
                     if in_ipython() or in_databricks() or PyGraphistry.set_sso_opt_into_type == "display":
                         display_message_html(f"<strong>{msg}</strong>")
                     raise Exception(msg) 
@@ -2635,9 +2635,9 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
         if not PyGraphistry.api_token():
             msg_html = '<br /><strong> .... </strong>'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
-                msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat*wait} seconds .... </strong>'
+                msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat * wait} seconds .... </strong>'
                 if not fail_silent:
-                    raise Exception(f"Failed to get token after {repeat*wait} seconds. Please re-run the login process") 
+                    raise Exception(f"Failed to get token after {repeat * wait} seconds. Please re-run the login process") 
                 else:
                     msg_html = f'{msg_html}<br /><strong>Got token</strong>'
                     display(HTML(msg_html))
@@ -2683,7 +2683,7 @@ def sso_verify_token_display(
 
         if required_login:
             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 {repeat*wait} seconds, please login within {wait} seconds....</strong><br /><strong>Please close the browser tab and come back to dashboard....</strong>'
+            msg_html = f'<br /><strong>Prepare to sign in ....</strong><br><strong>Please Login with the link appear later. Waiting for success login for {repeat * 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))
 
 

From c048ee2836a7afe9f40bebe6f3d30dee3b900cb0 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Mon, 25 Nov 2024 23:04:03 +0800
Subject: [PATCH 27/34] fix (doc): Add documentations

---
 .../graphistry-notebook-dashboard.ipynb       |  4 +++
 docs/source/server/register.rst               | 14 +++++++++
 graphistry/pygraphistry.py                    | 30 ++++++++++---------
 3 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/demos/demos_databases_apis/databricks_pyspark/graphistry-notebook-dashboard.ipynb b/demos/demos_databases_apis/databricks_pyspark/graphistry-notebook-dashboard.ipynb
index 515a270574..dfdb287bc6 100755
--- a/demos/demos_databases_apis/databricks_pyspark/graphistry-notebook-dashboard.ipynb
+++ b/demos/demos_databases_apis/databricks_pyspark/graphistry-notebook-dashboard.ipynb
@@ -152,6 +152,10 @@
     "                    protocol='https',\n",
     "                    server='hub.graphistry.com')\n",
     "\n",
+    "# Alternatively, use SSO: \n",
+    "# graphistry.register_databricks_sso(api=3, server='hub.graphistry.com', org_name='my-org')\n",
+    "# For more options, see https://github.com/graphistry/pygraphistry#configure\n",
+    "\n",
     "# Alternatively, use username and password: \n",
     "# graphistry.register(api=3, username='...', password='...', protocol='https', server='hub.graphistry.com')\n",
     "# For more options, see https://github.com/graphistry/pygraphistry#configure"
diff --git a/docs/source/server/register.rst b/docs/source/server/register.rst
index 5cc7512d30..9b131a439d 100644
--- a/docs/source/server/register.rst
+++ b/docs/source/server/register.rst
@@ -219,6 +219,20 @@ Register with Custom Browser Routing
         client_protocol_hostname="https://my_ui_server.com"
     )
 
+Register with SSO using helper function(only for databricks, organization with Specific IdP)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+    import graphistry
+
+    graphistry.databricks_register_sso(
+        api=3,
+        org_name="my_org_name",
+        idp_name="my_idp_name",
+        sso_opt_into_type="browser"
+    )
+
 ---
 
 Best Practices
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 7c41ee3d7f..3b3b0f84ab 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -580,6 +580,7 @@ def set_bolt_driver(driver=None):
 
     @staticmethod
     def set_sso_opt_into_type(value: Optional[str]):
+        """Set sso_opt_into_type to memory"""
         PyGraphistry._config["sso_opt_into_type"] = value
 
 
@@ -2655,6 +2656,19 @@ def sso_verify_token_display(
         wait: int = 5,
         display_mode: str = 'text'
     ) -> bool:
+        """Verify the JWT token display the corresponding message either in text or HTML.
+
+        :param repeat: Number of times to attempt obtaining the token, defaults to 20
+        :type repeat: int, optional
+        :param wait: Number of seconds to wait between attempts, defaults to 5
+        :type wait: int, optional
+        :param display_mode: Whether to display message in text or HTML, default is 'text'
+        :type display_mode: str, optional
+        :return: Whether token is still valid
+        :rtype: bool
+
+
+        """        
         if display_mode == 'html':
             from IPython.display import display, HTML, clear_output
             clear_output()
@@ -2663,7 +2677,6 @@ def sso_verify_token_display(
         token = PyGraphistry.api_token()
         if token:
             is_valid = PyGraphistry.verify_token()
-            print(f"is_valid : {is_valid}")
             if not is_valid:
                 print("***********token not valid, refresh token*****************")
                 if display_mode == 'html':
@@ -2692,11 +2705,7 @@ def sso_verify_token_display(
 
     # Databricks Dashboard SSO helper functions
     class DatabricksHelper():
-        """Helper class for databricks.
-
-        **Helper class to improve the sso login flow**
-
-        """
+        """Helper class to improve the sso login flow"""
         @staticmethod
         def register_databricks_sso(
             server: Optional[str] = None,
@@ -2704,16 +2713,9 @@ def register_databricks_sso(
             idp_name: Optional[str] = None,
             **kwargs
         ):
+            """Wrapper function special for databrick run register function"""
             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")
-            
-
-        # @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

From 274126d77e5bc7d536a00911a5d6b7785a396eba Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Tue, 3 Dec 2024 23:01:06 +0800
Subject: [PATCH 28/34] doc (databricks): Add documentation the docString, add
 changelog

---
 CHANGELOG.md               |  7 ++++++
 graphistry/pygraphistry.py | 46 +++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5206cd95cd..17000181d3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,11 +6,18 @@ The changelog format is based on [Keep a Changelog](https://keepachangelog.com/e
 This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) and all PyGraphistry-specific breaking changes are explictly noted here.
 
 ## [Development]
+### Added
+* Added new wrapper function for register to Databricks via SSO, to simplify code to display on dashboard
+
+### Fixes
+* Fix refresh function does not reset some variable in memory that cause weird behavior
 
 ### Docs
 
 * New tutorial on GPU memory management and capacity planning in the GPU section
 
+* Documentation changed for databricks related docs
+
 ## [0.34.17 - 2024-10-20]
 
 ### Added
diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 3b3b0f84ab..5ba1bc36f6 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2666,8 +2666,6 @@ def sso_verify_token_display(
         :type display_mode: str, optional
         :return: Whether token is still valid
         :rtype: bool
-
-
         """        
         if display_mode == 'html':
             from IPython.display import display, HTML, clear_output
@@ -2713,7 +2711,49 @@ def register_databricks_sso(
             idp_name: Optional[str] = None,
             **kwargs
         ):
-            """Wrapper function special for databrick run register function"""
+            """Wrapper function for :func:`graphistry.PyGraphistry.register`, specially for databrick SSO
+        :param server: URL of the visualization server.
+        :type server: Optional[str]
+        :param certificate_validation: Override default-on check for valid TLS certificate by setting to True.
+        :type certificate_validation: Optional[bool]
+        :param bolt: Neo4j bolt information. Optional driver or named constructor arguments for instantiating a new one.
+        :type bolt: Union[dict, Any]
+        :param token_refresh_ms: Ignored for now; JWT token auto-refreshed on plot() calls.
+        :type token_refresh_ms: int
+        :param store_token_creds_in_memory: Store username/password in-memory for JWT token refreshes (Token-originated have a hard limit, so always-on requires creds somewhere)
+        :type store_token_creds_in_memory: Optional[bool]
+        :param client_protocol_hostname: Override protocol and host shown in browser. Defaults to protocol/server or envvar GRAPHISTRY_CLIENT_PROTOCOL_HOSTNAME.
+        :type client_protocol_hostname: Optional[str]
+        :param org_name: Set login organization's name(slug). Defaults to user's personal organization.
+        :type org_name: Optional[str]
+        :param idp_name: Set sso login idp name. Default as None (for site-wide SSO / for the only idp record).
+        :type idp_name: Optional[str]
+        :param sso_opt_into_type: Show the SSO url with display(), webbrowser.open(), or print()
+        :type sso_opt_into_type: Optional[Literal["display", "browser"]]
+        :returns: None.
+        :rtype: None
+
+        **Example: Standard (2.0 api by org_name via SSO configured for site or for organization with only 1 IdP)**
+        ::
+
+            import graphistry
+            graphistry.register(api=3, protocol='http', server='200.1.1.1', org_name="org-name")
+
+        **Example: Standard (2.0 api by org_name via SSO with a specific IdP configured for an organization)**
+        ::
+
+            import graphistry
+            graphistry.register_databricks_sso(server='200.1.1.1', org_name="org-name", idp_name="idp-name")
+
+        **Example: Override SSO url display method to use `display()`, `webbrowser.open()`, or just `print()`**
+        ::
+
+            import graphistry
+            graphistry.register(api=3, protocol='http', server='200.1.1.1', org_name="org-name", sso_opt_into_type="display")
+            graphistry.register(api=3, protocol='http', server='200.1.1.1', org_name="org-name", sso_opt_into_type="browser")
+            graphistry.register(api=3, protocol='http', server='200.1.1.1', org_name="org-name", sso_opt_into_type=None)
+
+        """
             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")
 

From 6ac7cb63c64fbf5e1164974fc0edbd2c9ec2ed5b Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 4 Dec 2024 11:59:37 +0800
Subject: [PATCH 29/34] fix (message): Adjust the  message

---
 graphistry/pygraphistry.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 5ba1bc36f6..aeec595865 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -277,8 +277,8 @@ def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type):
         if in_ipython() or in_databricks() or sso_opt_into_type == 'display':  # If run in notebook, just display the HTML
             # from IPython.core.display import HTML
             from IPython.display import display, HTML
-            display(HTML(f'<a href="{auth_url}" target="_blank">Login SSO</a><br /><span>Please click the above URL to open browser to login</span>'))
-            print("Please click the above URL to open browser to login")
+            display(HTML(f'<a href="{auth_url}" target="_blank" style="display: inline-block; margin-top: 10px; padding: 10px 20px; font-size: 16px; font-weight: bold; color: #fff; background-color: #007bff; text-decoration: none; border-radius: 5px;">Login with SSO</a>'))
+            print('<p style="margin-top: 5px; color: #555;">Please click the button above to open the browser and log in.</p>')
             print(f"If you cannot see the URL, please open browser, browse to this URL: {auth_url}")
             print("Please close browser tab after SSO login to back to notebook")
             # return HTML(make_iframe(auth_url, 20, extra_html=extra_html, override_html_style=override_html_style))
@@ -2596,21 +2596,21 @@ def sso_wait_for_token_text_display(repeat: int = 20, wait: int = 5, fail_silent
         if not PyGraphistry.api_token():
             msg_text = '....'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
-                msg_text = f'{msg_text}\nFailed to get token after {repeat * wait} seconds ....'
+                msg_text = f'{msg_text}\nUnable to retrieve token after {repeat * wait} seconds ....'
                 if not fail_silent:
-                    msg = f"Failed to get token after {repeat * wait} seconds. Please re-run the login process"
+                    msg = f"Unable to retrieve token after {repeat * wait} seconds. Please re-run the login process"
                     if in_ipython() or in_databricks() or PyGraphistry.set_sso_opt_into_type == "display":
                         display_message_html(f"<strong>{msg}</strong>")
                     raise Exception(msg) 
                 else:
-                    msg_text = f'{msg_text}\nGot token'
+                    msg_text = f'{msg_text}\nToken retrieved successfully'
                     print(msg_text)
                     return
 
-            msg_text = f'{msg_text}\nGot token'
+            msg_text = f'{msg_text}\nToken retrieved successfully'
             print(msg_text)
         else:
-            print('Token is valid, no waiting required.')
+            print('Token is valid; no further action needed.')
 
 
     @staticmethod
@@ -2636,18 +2636,18 @@ def sso_wait_for_token_html_display(repeat: int = 20, wait: int = 5, fail_silent
         if not PyGraphistry.api_token():
             msg_html = '<br /><strong> .... </strong>'
             if not PyGraphistry.sso_repeat_get_token(repeat, wait):
-                msg_html = f'{msg_html}<br /><strong>Failed to get token after {repeat * wait} seconds .... </strong>'
+                msg_html = f'{msg_html}<br /><strong>Unable to retrieve token after {repeat * wait} seconds. Please try logging in again.</strong>'
                 if not fail_silent:
-                    raise Exception(f"Failed to get token after {repeat * wait} seconds. Please re-run the login process") 
+                    raise Exception(f"Unable to retrieve token after {repeat * wait} seconds. Please try logging in again.") 
                 else:
-                    msg_html = f'{msg_html}<br /><strong>Got token</strong>'
+                    msg_html = f'{msg_html}<br /><strong>Token retrieved successfully</strong>'
                     display(HTML(msg_html))
                     return
 
-            msg_html = f'{msg_html}<br /><strong>Got token</strong>'
+            msg_html = f'{msg_html}<br /><strong>Token retrieved successfully</strong>'
             display(HTML(msg_html))
         else:
-            display(HTML('<br /><strong>Token is valid, no waiting required.</strong>'))
+            display(HTML('<br /><strong>Token is valid; no further action needed.</strong>'))
 
 
     @staticmethod

From a5779b42c27bbacc11a4b4bb67a78d57e78df0aa Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 4 Dec 2024 12:01:04 +0800
Subject: [PATCH 30/34] fix (message): adjust message

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index aeec595865..15a534396a 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -2678,7 +2678,7 @@ def sso_verify_token_display(
             if not is_valid:
                 print("***********token not valid, refresh token*****************")
                 if display_mode == 'html':
-                    display(HTML('<br /><strong>Refresh token ....</strong>'))
+                    display(HTML('<br /><strong>Refreshing token ....</strong>'))
                 try:
                     PyGraphistry.refresh()
                 except Exception:

From 4df33e2c1c666fb7d50ce8fc40a94b61594b6679 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 25 Dec 2024 21:37:04 +0800
Subject: [PATCH 31/34] fix (util): revert the make_iframe function

---
 graphistry/util.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/graphistry/util.py b/graphistry/util.py
index b4dee6acb9..c581de5c5d 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -10,7 +10,7 @@
 import uuid
 import warnings
 from functools import lru_cache
-from typing import Any, Optional
+from typing import Any
 from collections import UserDict
 
 from .constants import VERBOSE, CACHE_COERCION_SIZE, TRACE
@@ -181,7 +181,7 @@ def display_message_html(message: str, cleared: Optional[bool] = False):
     display(HTML(message))
 
 
-def make_iframe(url, height, extra_html="", override_html_style=None, srcdoc: Optional[str] = None):
+def make_iframe(url, height, extra_html="", override_html_style=None):
     id = uuid.uuid4()
 
     height_str = (
@@ -211,7 +211,7 @@ def make_iframe(url, height, extra_html="", override_html_style=None, srcdoc: Op
         )
 
     iframe = """
-            <iframe id="%s" src="%s" {srcdoc}
+            <iframe id="%s"
                     allowfullscreen="true" webkitallowfullscreen="true" mozallowfullscreen="true"
                     oallowfullscreen="true" msallowfullscreen="true"
                     style="%s"

From e3b1e8c086f40990ade82e60992339062682a15c Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 25 Dec 2024 21:43:44 +0800
Subject: [PATCH 32/34] fix (util): fix make_iframe

---
 graphistry/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/util.py b/graphistry/util.py
index c581de5c5d..55927d7f86 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -211,7 +211,7 @@ def make_iframe(url, height, extra_html="", override_html_style=None):
         )
 
     iframe = """
-            <iframe id="%s"
+            <iframe id="%s" src="%s"
                     allowfullscreen="true" webkitallowfullscreen="true" mozallowfullscreen="true"
                     oallowfullscreen="true" msallowfullscreen="true"
                     style="%s"

From b53051aed2b3e57e4256a35e084efda035689ae1 Mon Sep 17 00:00:00 2001
From: Vaim Dev <dev@v-aim.com>
Date: Wed, 25 Dec 2024 22:01:23 +0800
Subject: [PATCH 33/34] fix (util): Fix syntax error

---
 graphistry/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/util.py b/graphistry/util.py
index 55927d7f86..374566fa8b 100644
--- a/graphistry/util.py
+++ b/graphistry/util.py
@@ -10,7 +10,7 @@
 import uuid
 import warnings
 from functools import lru_cache
-from typing import Any
+from typing import Any, Optional
 from collections import UserDict
 
 from .constants import VERBOSE, CACHE_COERCION_SIZE, TRACE

From baf92c70427906486ebf02e9425003020414fed4 Mon Sep 17 00:00:00 2001
From: vaimdev <dev@v-aim.com>
Date: Fri, 24 Jan 2025 12:17:21 +0800
Subject: [PATCH 34/34] fix (conflict): Fix bug after resolve conflict

---
 graphistry/pygraphistry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py
index 1b52523915..da158da842 100644
--- a/graphistry/pygraphistry.py
+++ b/graphistry/pygraphistry.py
@@ -583,7 +583,7 @@ def set_sso_opt_into_type(value: Optional[str]):
         """Set sso_opt_into_type to memory"""
         PyGraphistry._config["sso_opt_into_type"] = value
 
-    # def set_spanner_config(spanner_config):  
+    @staticmethod 
     def set_spanner_config(spanner_config: Optional[Union[Dict, str]] = None):
         """
         Saves the spanner config to internal Pygraphistry _config