Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix (sso): In databricks, the HTML is not displayable, print out the … #512

Merged
merged 14 commits into from
Oct 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm

## [Development]

## [0.29.6 - 2023-10-23]

### Docs

* Memgraph: Add tutorial (https://github.com/graphistry/pygraphistry/pull/507 by https://github.com/karmenrabar)
Expand All @@ -15,6 +17,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm

* Guard against potential `requests`` null dereference in uploader error handling

### Security

* Add control `register(..., sso_opt_into_type='browser' | 'display' | None)`
* Fix display of SSO URL

## [0.29.5 - 2023-08-23]

### Fixed
Expand Down
22 changes: 12 additions & 10 deletions graphistry/arrow_uploader.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,12 +183,12 @@ def __init__(self,
# check current org_name
from .pygraphistry import PyGraphistry
if 'org_name' in PyGraphistry._config:
logger.debug("@ArrowUploader.__init__: There is an org_name : {}".format(PyGraphistry._config['org_name']))
logger.debug("@ArrowUploader.__init__: There is an org_name : %s", PyGraphistry._config['org_name'])
self.__org_name = PyGraphistry._config['org_name']
else:
self.__org_name = None

logger.debug("2. @ArrowUploader.__init__: After set self.org_name: {}, self.__org_name : {}".format(self.org_name, self.__org_name))
logger.debug("2. @ArrowUploader.__init__: After set self.org_name: %s, self.__org_name : %s", self.org_name, self.__org_name)


def login(self, username, password, org_name=None):
Expand Down Expand Up @@ -254,7 +254,7 @@ def _handle_login_response(self, out, org_name):
del PyGraphistry._config['org_name']
else:
if org_name in PyGraphistry._config:
logger.debug("@ArrowUploder, handle login reponse, org_name: {}".format(PyGraphistry._config['org_name']))
logger.debug("@ArrowUploder, handle login reponse, org_name: %s", PyGraphistry._config['org_name'])
PyGraphistry._config['org_name'] = logged_in_org_name
# PyGraphistry.org_name(logged_in_org_name)
except Exception:
Expand Down Expand Up @@ -287,17 +287,18 @@ def sso_login(self, org_name=None, idp_name=None):
url, data={'client-type': 'pygraphistry'},
verify=self.certificate_validation
)
# print(out.text)

json_response = None
try:
logger.debug("@ArrowUploader.sso_login, out.text: %s", out.text)
json_response = out.json()
logger.debug("@ArrowUploader.sso_login, json_response: {}".format(json_response))
logger.debug("@ArrowUploader.sso_login, json_response: %s", json_response)
self.token = None
if not ('status' in json_response):
raise Exception(out.text)
else:
if json_response['status'] == 'OK':
logger.debug("@ArrowUploader.sso_login, json_data : {}".format(json_response['data']))
logger.debug("@ArrowUploader.sso_login, json_data : %s", json_response['data'])
if 'state' in json_response['data']:
self.sso_state = json_response['data']['state']
self.sso_auth_url = json_response['data']['auth_url']
Expand All @@ -308,6 +309,7 @@ def sso_login(self, org_name=None, idp_name=None):

except Exception:
logger.error('Error: %s', out, exc_info=True)
print("\nThere is error with the SSO login, please check your SSO and IDP configuration")
raise

return self
Expand Down Expand Up @@ -336,7 +338,7 @@ def sso_get_token(self, state):
if 'token' in json_response['data']:
self.token = json_response['data']['token']
if 'active_organization' in json_response['data']:
logger.debug("@ArrowUploader.sso_get_token, org_name: {}".format(json_response['data']['active_organization']['slug']))
logger.debug("@ArrowUploader.sso_get_token, org_name: %s", json_response['data']['active_organization']['slug'])
self.org_name = json_response['data']['active_organization']['slug']

except Exception as e:
Expand Down Expand Up @@ -382,7 +384,7 @@ def create_dataset(self, json): # noqa: F811
tok = self.token
if self.org_name:
json['org_name'] = self.org_name
logger.debug("@ArrowUploder create_dataset json: {}".format(json))
logger.debug("@ArrowUploder create_dataset json: %s", json)
res = requests.post(
self.server_base_path + '/api/v2/upload/datasets/',
verify=self.certificate_validation,
Expand Down Expand Up @@ -490,7 +492,7 @@ def post(self, as_files: bool = True, memoize: bool = True):
"""
Note: likely want to pair with self.maybe_post_share_link(g)
"""
logger.debug("@ArrowUploader.post, self.org_name : {}".format(self.org_name))
logger.debug("@ArrowUploader.post, self.org_name : %s", self.org_name)
if as_files:

file_uploader = ArrowFileUploader(self)
Expand Down Expand Up @@ -656,7 +658,7 @@ def post_arrow(self, arr: pa.Table, graph_type: str, opts: str = ''):
raise Exception('No success indicator in server response')
return out
except requests.exceptions.HTTPError as e:
logger.error('Failed to post arrow to %s (%s)', sub_path, e.request.url if e.request is not None else "(No request)", exc_info=True)
logger.error('Failed to post arrow to %s (%s)', sub_path, "{}/{}{}".format(self.server_base_path, sub_path, f"?{opts}" if len(opts) > 0 else ""), exc_info=True)
logger.error('%s', e)
logger.error('%s', e.response.text)
raise e
Expand Down
51 changes: 35 additions & 16 deletions graphistry/pygraphistry.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ def pkey_login(personal_key_id, personal_key_secret, org_name=None, fail_silent=
return PyGraphistry.api_token()

@staticmethod
def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SECONDS):
def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SECONDS, sso_opt_into_type=None):
"""Authenticate with SSO and set token for reuse (api=3).

:param org_name: Set login organization's name(slug). Defaults to user's personal organization.
Expand All @@ -211,16 +211,18 @@ def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SEC
:type idp_name: Optional[str]
: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.
:rtype: None
: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: token or auth_url
:rtype: Optional[str]

SSO Login logic.

"""

if PyGraphistry._config['store_token_creds_in_memory']:
PyGraphistry.relogin = lambda: PyGraphistry.sso_login(
org_name, idp_name, sso_timeout
org_name, idp_name, sso_timeout, sso_opt_into_type
)

PyGraphistry._is_authenticated = False
Expand All @@ -244,39 +246,45 @@ def sso_login(org_name=None, idp_name=None, sso_timeout=SSO_GET_TOKEN_ELAPSE_SEC
auth_url = arrow_uploader.sso_auth_url
# print("auth_url : {}".format(auth_url))
if auth_url and not PyGraphistry.api_token():
PyGraphistry._handle_auth_url(auth_url, sso_timeout)

PyGraphistry._handle_auth_url(auth_url, sso_timeout, sso_opt_into_type)
return auth_url

@staticmethod
def _handle_auth_url(auth_url, sso_timeout):
def _handle_auth_url(auth_url, sso_timeout, sso_opt_into_type):
"""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
: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.
:rtype: None
: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: token
:rtype: token: Optional[str]

SSO Login logic.

"""

if in_ipython() or in_databricks(): # If run in notebook, just display the HTML
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>'))
print("Please click the above link to open browser to login")
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")
# 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")
elif sso_opt_into_type == 'browser':
print("Please minimize browser after your SSO login and go back to pygraphistry")

import webbrowser
input("Press Enter to open browser ...")
# open browser to auth_url
webbrowser.open(auth_url)
else:
print(f"Please open a browser, browse to this URL, and sign in: {auth_url}")
print("After, if you get timeout error, run graphistry.sso_get_token() to complete the authentication")

if sso_timeout is not None:
time.sleep(1)
Expand Down Expand Up @@ -305,7 +313,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("Successfully logged in")
return PyGraphistry.api_token()
else:
return None
Expand Down Expand Up @@ -562,7 +570,8 @@ def register(
org_name: Optional[str] = None,
idp_name: Optional[str] = None,
is_sso_login: Optional[bool] = False,
sso_timeout: Optional[int] = SSO_GET_TOKEN_ELAPSE_SECONDS
sso_timeout: Optional[int] = SSO_GET_TOKEN_ELAPSE_SECONDS,
sso_opt_into_type: Optional[Literal["display", "browser"]] = None
):
"""API key registration and server selection

Expand Down Expand Up @@ -606,6 +615,8 @@ def register(
:type idp_name: Optional[str]
: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]
: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

Expand All @@ -621,6 +632,14 @@ def register(
import graphistry
graphistry.register(api=3, protocol='http', server='200.1.1.1', org_name="org-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)

**Example: Standard (2.0 api by username/password with org_name)**
::

Expand Down Expand Up @@ -690,7 +709,7 @@ def register(
PyGraphistry.api_token(token or PyGraphistry._config['api_token'])
elif not (org_name is None) or is_sso_login:
print(MSG_REGISTER_ENTER_SSO_LOGIN)
PyGraphistry.sso_login(org_name, idp_name, sso_timeout=sso_timeout)
PyGraphistry.sso_login(org_name, idp_name, sso_timeout=sso_timeout, sso_opt_into_type=sso_opt_into_type)

@staticmethod
def __check_login_type_to_reset_token_creds(
Expand Down