From 5ce2f3dad84f32e3990a396a2a3b2a031acc4a56 Mon Sep 17 00:00:00 2001 From: Teo Date: Tue, 12 Sep 2023 16:47:26 +0800 Subject: [PATCH 01/10] feat (pygraphistry) set sso login to switch org api --- graphistry/arrow_uploader.py | 8 ++-- graphistry/pygraphistry.py | 83 +++++++++++++++++++++++++++++------- 2 files changed, 72 insertions(+), 19 deletions(-) diff --git a/graphistry/arrow_uploader.py b/graphistry/arrow_uploader.py index 765fe3067e..145382c71b 100644 --- a/graphistry/arrow_uploader.py +++ b/graphistry/arrow_uploader.py @@ -320,10 +320,8 @@ def sso_get_token(self, state): # from .pygraphistry import PyGraphistry base_path = self.server_base_path - out = requests.get( - f'{base_path}/api/v2/o/sso/oidc/jwt/{state}/', - verify=self.certificate_validation - ) + url = f'{base_path}/api/v2/o/sso/oidc/jwt/{state}/' + out = requests.get(url,verify=self.certificate_validation) json_response = None try: json_response = out.json() @@ -338,6 +336,8 @@ def sso_get_token(self, state): if 'active_organization' in json_response['data']: logger.debug("@ArrowUploader.sso_get_token, org_name: {}".format(json_response['data']['active_organization']['slug'])) self.org_name = json_response['data']['active_organization']['slug'] + if 'state' in json_response['data']: + self.state = json_response['data']['state'] except Exception as e: logger.error('Unexpected SSO authentication error: %s', out, exc_info=True) diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index 9b4430d9d8..5d677fc984 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -219,9 +219,7 @@ def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SEC """ if PyGraphistry._config['store_token_creds_in_memory']: - PyGraphistry.relogin = lambda: PyGraphistry.sso_login( - org_name, idp_name, sso_timeout - ) + PyGraphistry.relogin = lambda: PyGraphistry.sso_login(org_name, idp_name, sso_timeout) PyGraphistry._is_authenticated = False arrow_uploader = ArrowUploader( @@ -248,12 +246,12 @@ def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SEC @staticmethod - def _handle_auth_url(auth_url, sso_timeout): + def _handle_auth_url(auth_url, sso_timeout, relogin=False): """Internal function to handle what to do with the auth_url based on the client mode python/ipython console or notebook. :param auth_url: SSO auth url retrieved via API - :type auth_url: str + :type auth_url: str or list in list([[name1,url1], [name2,url2], [name3,url3]) :param sso_timeout: Set sso login getting token timeout in seconds (blocking mode), set to None if non-blocking mode. Default as SSO_GET_TOKEN_ELAPSE_SECONDS. :type sso_timeout: Optional[int] :returns: None. @@ -266,17 +264,48 @@ def _handle_auth_url(auth_url, sso_timeout): if in_ipython() or in_databricks(): # If run in notebook, just display the HTML # from IPython.core.display import HTML from IPython.display import display, HTML - display(HTML(f'Login SSO')) - print("Please click the above link to open browser to login") + if isinstance(auth_url ,list): + for auth_url_each in auth_url: + # pop up window + # display(HTML(f'{auth_url_each[0]}')) + display(HTML(f'{auth_url_each[0]}')) + else: + # display(HTML(f'Login SSO')) + display(HTML(f'Login SSO')) + if relogin: + print("Please click the above link to open browser to access SSO organization") + else: + print("Please click the above link to open browser to login") 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)) else: - print("Please minimize browser after SSO login to back to pygraphistry") - import webbrowser - input("Press Enter to open browser ...") - # open browser to auth_url - webbrowser.open(auth_url) + print("Please minimize browser after SSO login to back to pygraphistry") + if isinstance(auth_url ,list): + if len(auth_url) == 1: + print(f"idp name: {auth_url[0][0]}") + input("Press Enter to open browser ...") + webbrowser.open(auth_url[0][1]) + else: + while True: + url_dict = {} + for index, (name, url) in enumerate(auth_url, start=1): + url_dict[str(index)] = url + print(f"{index}: {name}") + input_key = input("Enter a number above to open the browser or 'quit' to exit: ") + + if input_key in url_dict: + selected_url = url_dict[input_key] + webbrowser.open(selected_url) + break + elif input_key.strip().lower() == 'quit': + break + else: + print("Invalid key. No URL found.") + else: + input("Press Enter to open browser ...") + # open browser to auth_url + webbrowser.open(auth_url) if sso_timeout is not None: time.sleep(1) @@ -305,7 +334,7 @@ def _handle_auth_url(auth_url, sso_timeout): # set org_name to sso org PyGraphistry._config['org_name'] = org_name - print("Successfully get a token") + print(f"Successfully get a token, org_name is {org_name}") return PyGraphistry.api_token() else: return None @@ -338,6 +367,7 @@ def _sso_get_token(): try: token = arrow_uploader.token org_name = arrow_uploader.org_name + PyGraphistry.sso_state(arrow_uploader.state) except Exception: pass logger.debug("jwt token :{}".format(token)) @@ -2386,10 +2416,33 @@ def switch_org(value): result = PyGraphistry._handle_api_response(response) if result is True: - PyGraphistry._config['org_name'] = value.strip() - logger.info("Switched to organization: {}".format(value.strip())) + PyGraphistry._api_response_switch_org(response) else: # print the error message raise Exception(result) + + @staticmethod + def _api_response_switch_org(response): + try: + json_response = response.json() + message = json_response.get('message', '') + data = json_response.get('data', '') + if message.startswith('Switch to organization'): + PyGraphistry._config['org_name'] = data['organization_slug'] + logger.info("Switched to organization: {}".format(data['organization_slug'])) + elif message.startswith('Login to SSO for switch to organization') or message.startswith('Choose SSO for switch to organization'): + idp_name_with_url_list = [] + idp_state_list = [] + for idp_name in data['idp']: + idp_name_with_url_list.append([idp_name, data['idp'][idp_name]['auth_url']]) + idp_state_list.append(data['idp'][idp_name]['state']) + multiple_idp_state = 'SEPARATE'.join(idp_state_list) + PyGraphistry.sso_state(multiple_idp_state) + PyGraphistry._handle_auth_url(idp_name_with_url_list, sso_timeout=SSO_GET_TOKEN_ELAPSE_SECONDS, relogin=True) + else: + return message + except: + logger.error('Error: %s', response, exc_info=True) + raise Exception("Unknown Error") @staticmethod def _handle_api_response(response): From 760ddfa52b47e7e4a0c9076fd5366ab4a7d905f8 Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Wed, 17 Jan 2024 10:45:07 +0800 Subject: [PATCH 02/10] fix(sso): :bug: fix require sso_opt_into_type positional argument --- graphistry/pygraphistry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index 2c94ac4aaf..e5a782a142 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -250,7 +250,7 @@ def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SEC return auth_url @staticmethod - def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type, relogin=False): + def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type='browser', relogin=False): """Internal function to handle what to do with the auth_url based on the client mode python/ipython console or notebook. From d233704dd8240af3a28aef7dfdd104266bf9f8cc Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Wed, 17 Jan 2024 10:56:28 +0800 Subject: [PATCH 03/10] revert(sso): :rewind: revert some code to master version --- graphistry/pygraphistry.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index e5a782a142..affe5da2dd 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -287,8 +287,9 @@ def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type='browser', relogin 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)) elif sso_opt_into_type == 'browser': - import webbrowser print("Please minimize browser after your SSO login and go back to pygraphistry") + + import webbrowser if isinstance(auth_url ,list): if len(auth_url) == 1: print(f"idp name: {auth_url[0][0]}") From 4f23642d2d909ebc87f6ed99a3f7482d0900ac54 Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Mon, 22 Jan 2024 12:48:44 +0800 Subject: [PATCH 04/10] fix(sso): :fire: remove redundance code --- graphistry/pygraphistry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index affe5da2dd..95a1b712c9 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -379,7 +379,6 @@ def _sso_get_token(): try: token = arrow_uploader.token org_name = arrow_uploader.org_name - PyGraphistry.sso_state(arrow_uploader.state) except Exception: pass logger.debug("jwt token :{}".format(token)) From 6d97bdfbe3d92793467cfb9ae3da801dc9c5196a Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Mon, 22 Jan 2024 15:05:56 +0800 Subject: [PATCH 05/10] fix(sso): :bug: change SEPARATE to | --- graphistry/pygraphistry.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index 95a1b712c9..94cdde87c5 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -32,6 +32,7 @@ ############################################################################### SSO_GET_TOKEN_ELAPSE_SECONDS = 50 +SSO_STATE_SPLIT = "|" EnvVarNames = { "api_key": "GRAPHISTRY_API_KEY", @@ -2457,7 +2458,7 @@ def _api_response_switch_org(response): for idp_name in data['idp']: idp_name_with_url_list.append([idp_name, data['idp'][idp_name]['auth_url']]) idp_state_list.append(data['idp'][idp_name]['state']) - multiple_idp_state = 'SEPARATE'.join(idp_state_list) + multiple_idp_state = SSO_STATE_SPLIT.join(idp_state_list) PyGraphistry.sso_state(multiple_idp_state) PyGraphistry._handle_auth_url(idp_name_with_url_list, sso_timeout=SSO_GET_TOKEN_ELAPSE_SECONDS, relogin=True) else: From 6bc1b58467a6280b81a881b53d972f8462cc7d6c Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Tue, 23 Jan 2024 10:32:13 +0800 Subject: [PATCH 06/10] fix(sso): :bug: fix fail test-full-ai --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 8b048e6abc..c6cc3e9403 100755 --- a/setup.py +++ b/setup.py @@ -42,7 +42,7 @@ def unique_flatten_dict(d): } base_extras_heavy = { - 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn>=1.0'], + 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn==1.2.2'], } # https://github.com/facebookresearch/faiss/issues/1589 for faiss-cpu 1.6.1, #'setuptools==67.4.0' removed base_extras_heavy['ai'] = base_extras_heavy['umap-learn'] + ['scipy', 'dgl', 'torch<2', 'sentence-transformers', 'faiss-cpu', 'joblib'] From 3de5d74d2cb74b0271dc49e51140a24ef4a060a8 Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Tue, 23 Jan 2024 10:56:16 +0800 Subject: [PATCH 07/10] fix(sso): :bug: upgrate scikit-learn to 1.3.2 In the last passing test, scikit-learn is version 1.3.2 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index c6cc3e9403..2dd943e6ff 100755 --- a/setup.py +++ b/setup.py @@ -42,7 +42,7 @@ def unique_flatten_dict(d): } base_extras_heavy = { - 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn==1.2.2'], + 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn==1.3.2'], } # https://github.com/facebookresearch/faiss/issues/1589 for faiss-cpu 1.6.1, #'setuptools==67.4.0' removed base_extras_heavy['ai'] = base_extras_heavy['umap-learn'] + ['scipy', 'dgl', 'torch<2', 'sentence-transformers', 'faiss-cpu', 'joblib'] From 172387ba16d39759d3612086adfa3e44e1feb85a Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Wed, 24 Jan 2024 10:22:33 +0800 Subject: [PATCH 08/10] fix(sso): :bug: testai cfg to skip the problematic recent scikits --- setup.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 2dd943e6ff..6d7b733cf3 100755 --- a/setup.py +++ b/setup.py @@ -27,7 +27,8 @@ def unique_flatten_dict(d): 'docs': ['sphinx==3.4.3', 'docutils==0.16', 'sphinx_autodoc_typehints==1.11.1', 'sphinx-rtd-theme==0.5.1', 'Jinja2<3.1'], 'test': ['flake8>=5.0', 'mock', 'mypy', 'pytest'] + stubs, 'testai': [ - 'numba>=0.57.1' # https://github.com/numba/numba/issues/8615 + 'numba>=0.57.1', # https://github.com/numba/numba/issues/8615 + 'scikit-learn<=1.3.2' ], 'build': ['build'] } @@ -42,7 +43,7 @@ def unique_flatten_dict(d): } base_extras_heavy = { - 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn==1.3.2'], + 'umap-learn': ['umap-learn', 'dirty-cat==0.2.0', 'scikit-learn>=1.0'], } # https://github.com/facebookresearch/faiss/issues/1589 for faiss-cpu 1.6.1, #'setuptools==67.4.0' removed base_extras_heavy['ai'] = base_extras_heavy['umap-learn'] + ['scipy', 'dgl', 'torch<2', 'sentence-transformers', 'faiss-cpu', 'joblib'] From b866eafdfb8a4d12a38c52e05becd501913ab3b9 Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Wed, 24 Jan 2024 12:44:22 +0800 Subject: [PATCH 09/10] fix(sso): :bug: exclude scikit-learn version that cause problem --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 6d7b733cf3..7ceb64b2e3 100755 --- a/setup.py +++ b/setup.py @@ -28,7 +28,7 @@ def unique_flatten_dict(d): 'test': ['flake8>=5.0', 'mock', 'mypy', 'pytest'] + stubs, 'testai': [ 'numba>=0.57.1', # https://github.com/numba/numba/issues/8615 - 'scikit-learn<=1.3.2' + 'scikit-learn!=1.4.0,!=1.4.0rc1' ], 'build': ['build'] } From dcf47a04ee18beaaf017ae68b3b7f58a098e9749 Mon Sep 17 00:00:00 2001 From: Vaim Dev Date: Mon, 26 Feb 2024 15:26:24 +0800 Subject: [PATCH 10/10] fix(sso): :bug: use master code for setup.py --- setup.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/setup.py b/setup.py index faef74c716..c81db1b09c 100755 --- a/setup.py +++ b/setup.py @@ -29,8 +29,7 @@ def unique_flatten_dict(d): 'docs': ['sphinx==3.4.3', 'docutils==0.16', 'sphinx_autodoc_typehints==1.11.1', 'sphinx-rtd-theme==0.5.1', 'Jinja2<3.1'], 'test': ['flake8>=5.0', 'mock', 'mypy', 'pytest'] + stubs + test_workarounds, 'testai': [ - 'numba>=0.57.1', # https://github.com/numba/numba/issues/8615 - 'scikit-learn!=1.4.0,!=1.4.0rc1' + 'numba>=0.57.1' # https://github.com/numba/numba/issues/8615 ], 'build': ['build'] }