Personal key (graphistry#420)

* feat: Add org_name for logging in and corresponding Dataset upload with org_name intact

* fix: org_name global value retrieving bug found by test script

* refactor : clean up debugging code

* unittest: Add test scenarios for login with org_name

* fix: send org_name when create_file, privacy setting with mode_action

* fix: Remove debugging code, fix backward incompatibility

* fix: request json error

* fix: only pass org_name in json if org_name is not None, which failed in old server

* fix: arrow_uploader login handling for org_name

* fix (security): Adjust default mode_action to READ when share as public, set to WRITE only when share as private

* fix (test): Fix test scripts and add more test cases

* documetation on org-name parameter

* documentation: Update README.md

* documentation: Add detail

* fix: rearrange the sequence of exception thrown when new pygraphistry login with org on old server

* doc: Add documentation for sharing tutorial notebook on sharing within organization, minor correction for other documentation

* docs(readme.md): orgs

* docs(readme.md): org mode privacy

* feat (sso login): Add initial code to call API to get SSO login page and use state to retrieve token

* fix (sso): Login with SSO to obtain jwt token

* fix (cleanup): Clean up code for SSO login

* feat : Use timeout=None to replace flag for blocking mode

* refactor (sso): Handle cases for ipython console or jupyter notebook

* doc (register function for SSO): Added docstring doc for register function for SSO login

* doc (other SSO related function): Add docstring documentation

* fix (sso login): to allow blocking mode for running in notebook

* fix (register function): do not throw exception if register function does not pass in username/pwd or org_name

* fix (typecheck): type check failure fix

* fix (org_name): json object conversion problem for org_name if using property instead of calling function directly. Add some debug code

* refactor (print -> logger.debug): change print to logger.debug

* fix (org_name should only pass if org_name has value)

* fix (type imcompatible) test failure

* fix (org_name to cope with old and new server)

* fix (mypy): fix mypy issue

* fix (login): allow login with token

* wip (login with personal key): added personal key id and personal key for the login/register

* feat (personal key): Add personal key login capability initial code

* wip (sso login): SSO login fix for site wide

* Update README.md

* wip (site wide sso login): when no org slug and idp_name passed in register, try with site wide SSO

* fix (sso): Display SSO timeout exception in better and understandable

* fix (register): adjust register logic to check the missing username, password, missing personal key id & personal key. Add pytest for testing the scenario

* feat (sso enhancement): Add is_sso_login parameter to handle whether to do sso login when register

* feat (test scripts): add test script for the register function

* fix (arrow_uploader): bug in register with sso_login

* feat (organization): Add switch_org function to allow switching of organization

* fix (switch org): Fix switch org API

* fix (typecheck and lint)

* feat (add test script): add test script for switch org

* refactor (personal key to personal key secret): refactor variable and fix test scripts

* fix: docstring/typecheck Optional fix

* fix: typecheck docstring

* fix (raise exception instead of print): use a messages.py to keep the message as constant

* docs(readme): login

* fix : raise Exception instead of printing

* fix (debug info): remove debugging info

* fix(docs): personal key

* wip (switch org): fix after switch org with org_name('xxx'), plotting does not take the updated org_name. Fix done, pending to remove debugging code

* fix (clean up): Clean up debugging code for switch org

* fix (mypy newer version issues): Optional[xxx] = None, instead of xxx = None

* fix (mypy): fix the default value of layer to None

* fix (test scripts): add ipython in dev_extra(stubs), fix unauthenticated issue in test_ipython

* fix (personal key): fix personal key login does not switch org_name

* fix(logging): print to logger

* fix(types)

Co-authored-by: lmeyerov <leo@graphistry.com>

CHANGELOG.md

@@ -7,6 +7,15 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Development]
 
+### Added
+
+* Personal keys: `register(personal_key_id=..., personal_key_secret=...)`
+* SSO: `register()` (no user/pass), `register(idp_name=...)` (org-specific IDP)
+
+### Fixed
+
+* Type errors
+
 ## [0.28.4 - 2022-10-22]
 
 ### Added

README.md

@@ -51,16 +51,18 @@ You can use PyGraphistry with traditional Python data sources like CSVs, SQL, Ne
    ```python
   # pip install --user graphistry              # minimal
   # pip install --user graphistry[bolt,gremlin,nodexl,igraph,networkx]  # data plugins
-  # Requires Python 3.8+ (for scikit-learn 1.0+):
-  # pip install --user graphistry[umap-learn]  # UMAP autoML (without text support)
-  # pip install --user graphistry[ai]          # Full UMAP + GNN autoML, including sentence transformers (1GB+)
+  # AI modules: Python 3.8+ with scikit-learn 1.0+:
+  # pip install --user graphistry[umap-learn]  # Lightweight: UMAP autoML (without text support); scikit-learn 1.0+
+  # pip install --user graphistry[ai]          # Heavy: Full UMAP + GNN autoML, including sentence transformers (1GB+)
 
   import graphistry
   graphistry.register(api=3, username='abc', password='xyz')  # Free: hub.graphistry.com
-
-  #graphistry.register(..., org_name='my-org') # Upload into an organization account
-  #graphistry.register(..., protocol='http', server='my.site.ngo')  # Use with a self-hosted server
-
+  #graphistry.register(..., personal_key_id='pkey_id', personal_key_secret='pkey_secret') # Key instead of username+password+org_name
+  #graphistry.register(..., is_sso_login=True)  # SSO instead of password
+  #graphistry.register(..., org_name='my-org') # Upload into an organization account vs personal
+  #graphistry.register(..., protocol='https', server='my.site.ngo')  # Use with a self-hosted server
+  # ... and if client (browser) URLs are different than python server<> graphistry server uploads
+  #graphistry.register(..., client_protocol_hostname='https://public.acme.co')
   ```
 
 * **Notebook-friendly:** PyGraphistry plays well with interactive notebooks like [Jupyter](http://ipython.org), [Zeppelin](https://zeppelin.incubator.apache.org/), and [Databricks](http://databricks.com). Process, visualize, and drill into with graphs directly within your notebooks:

graphistry/ArrowFileUploader.py

@@ -1,6 +1,6 @@
 import pyarrow as pa, requests, sys
 from functools import lru_cache
-from typing import Any, Tuple
+from typing import Any, Tuple, Optional
 from weakref import WeakKeyDictionary
 from .util import setup_logger
 logger = setup_logger(__name__)
@@ -120,7 +120,7 @@ def post_arrow(self, arr: pa.Table, file_id: str, url_opts: str = 'erase=true')
     ###
 
     def create_and_post_file(
-        self, arr: pa.Table, file_id: str = None, file_opts: dict = {}, upload_url_opts: str = 'erase=true', memoize: bool = True
+        self, arr: pa.Table, file_id: Optional[str] = None, file_opts: dict = {}, upload_url_opts: str = 'erase=true', memoize: bool = True
     ) -> Tuple[str, dict]:
         """
             Create file and upload data for it.

graphistry/PlotterBase.py

@@ -1368,6 +1368,8 @@ def plot(
                     .plot(es)
 
         """
+        from .pygraphistry import PyGraphistry
+        logger.debug("1. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
 
         if graph is None:
             if self._edges is None:
@@ -1381,15 +1383,19 @@ def plot(
 
         self._check_mandatory_bindings(not isinstance(n, type(None)))
 
-        from .pygraphistry import PyGraphistry
+        # from .pygraphistry import PyGraphistry
         api_version = PyGraphistry.api_version()
+        logger.debug("2. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
         if api_version == 1:
             dataset = self._plot_dispatch(g, n, name, description, 'json', self._style, memoize)
             if skip_upload:
                 return dataset
             info = PyGraphistry._etl1(dataset)
         elif api_version == 3:
+            logger.debug("3. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
             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)
             if skip_upload:
                 return dataset
@@ -1903,7 +1909,6 @@ def _make_dataset(self, edges, nodes, name, description, mode, metadata=None, me
                 warn('Graph has no edges, may have rendering issues')
         except:
             1
-
         #compatibility checks
         if mode == 'json':
             if not (metadata is None):
@@ -1958,7 +1963,6 @@ def flatten_categorical(df):
     def _make_arrow_dataset(self, edges: pa.Table, nodes: pa.Table, name: str, description: str, metadata) -> ArrowUploader:
 
         from .pygraphistry import PyGraphistry
-
         au : ArrowUploader = ArrowUploader(
             server_base_path=PyGraphistry.protocol() + '://' + PyGraphistry.server(),
             edges=edges, nodes=nodes,
@@ -1971,8 +1975,7 @@ def _make_arrow_dataset(self, edges: pa.Table, nodes: pa.Table, name: str, descr
                 'agentversion': sys.modules['graphistry'].__version__,  # type: ignore
                 **(metadata or {})
             },
-            certificate_validation=PyGraphistry.certificate_validation(),
-            org_name=PyGraphistry.org_name())
+            certificate_validation=PyGraphistry.certificate_validation())
 
         au.edge_encodings = au.g_to_edge_encodings(self)
         au.node_encodings = au.g_to_node_encodings(self)

graphistry/__init__.py

@@ -4,6 +4,7 @@
     protocol,
     server,
     register,
+    sso_get_token,
     privacy,
     login,
     refresh,

graphistry/arrow_uploader.py

@@ -1,6 +1,7 @@
 from typing import List, Optional
 
 import io, pyarrow as pa, requests, sys
+
 from .ArrowFileUploader import ArrowFileUploader
 from .util import setup_logger
 logger = setup_logger(__name__)
@@ -18,12 +19,12 @@ def token(self, token: str):
         self.__token = token
 
     @property
-    def org_name(self) -> str:
+    def org_name(self) -> Optional[str]:
         return self.__org_name
 
     @org_name.setter
-    def org_name(self, org_name: str):
-        self.__org_name = org_name
+    def org_name(self, org_name: str) -> None:
+        self.__org_name: Optional[str] = org_name
 
     @property
     def dataset_id(self) -> str:
@@ -136,6 +137,19 @@ def certificate_validation(self, certificate_validation):
 
     ########################################################################3
 
+    # @property
+    # def sso_state(self) -> str:
+    #     return getattr(self, '__sso_state', "")
+
+    ########################################################################3
+
+    # @property
+    # def sso_auth_url(self) -> str:
+    #     return getattr(self, '__sso_auth_url')
+
+    ########################################################################3
+
+
     def __init__(self, 
             server_base_path='http://nginx', view_base_path='http://localhost',
             name = None,
@@ -146,6 +160,7 @@ def __init__(self,
             metadata = None,
             certificate_validation = True, 
             org_name: Optional[str] = None):
+
         self.__name = name
         self.__description = description
         self.__server_base_path = server_base_path
@@ -158,30 +173,65 @@ def __init__(self,
         self.__edge_encodings = edge_encodings
         self.__metadata = metadata
         self.__certificate_validation = certificate_validation
-        if org_name is not None:
+        self.__org_name = org_name if org_name else None
+
+        if org_name:
             self.__org_name = org_name
-
+        else:
+            # 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']))
+                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))
+
+
     def login(self, username, password, org_name=None):
-        from .pygraphistry import PyGraphistry
+        # base_path = self.server_base_path
+
+        json_data = {'username': username, 'password': password}
+        if org_name:
+            json_data.update({"org_name": org_name})
 
-        base_path = self.server_base_path
         out = requests.post(
-            f'{base_path}/api-token-auth/',
+            f'{self.server_base_path}/api-token-auth/',
+            verify=self.certificate_validation,
+            json=json_data)
+
+        return self._handle_login_response(out, org_name)
+
+    def pkey_login(self, personal_key_id, personal_key_secret, org_name=None):
+        # json_data = {'personal_key_id': personal_key_id, 'personal_key_secret': personal_key}
+        json_data = {}
+        if org_name:
+            json_data.update({"org_name": org_name})
+
+        headers = {"Authorization": f'PersonalKey {personal_key_id}:{personal_key_secret}'}
+
+        url = f'{self.server_base_path}/api/v2/auth/pkey/jwt/'
+
+        out = requests.get(
+            url,
             verify=self.certificate_validation,
-            json={'username': username, 'password': password, "org_name": org_name})
+            json=json_data, headers=headers)
+        return self._handle_login_response(out, org_name)
+
+    def _handle_login_response(self, out, org_name):
+        from .pygraphistry import PyGraphistry
         json_response = None
         try:
             json_response = out.json()
-
             if not ('token' in json_response):
                 raise Exception(out.text)
 
             org = json_response.get('active_organization',{})
             logged_in_org_name = org.get('slug', None)
-
             if org_name:  # caller pass in org_name
                 if not logged_in_org_name:  # no active_organization in JWT payload
-                    raise Exception("Server does not support organization, please omit org_name")
+                    raise Exception("You are not authorized to the organization '{}', or server does not support organization, please omit org_name parameter".format(org_name))
                 else:
                     # if JWT response with org_name different than the pass in org_name
                     # => org_name not found and return default organization (currently is personal org)
@@ -195,9 +245,16 @@ def login(self, username, password, org_name=None):
                     raise Exception("Organization {} is not found".format(org_name))
 
                 if not is_member: 
-                    raise Exception("You are not a member of {}".format(org_name))
+                    raise Exception("You are not authorized or not a member of {}".format(org_name))
 
-            PyGraphistry.org_name(logged_in_org_name)
+            if logged_in_org_name is None and org_name is None:
+                if 'org_name' in PyGraphistry._config:
+                    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']))
+                PyGraphistry._config['org_name'] = logged_in_org_name 
+                # PyGraphistry.org_name(logged_in_org_name)
         except Exception:
             logger.error('Error: %s', out, exc_info=True)
             raise
@@ -206,6 +263,87 @@ def login(self, username, password, org_name=None):
 
         return self
 
+    def sso_login(self, org_name=None, idp_name=None):
+        """
+        Koa, 04 May 2022    Get SSO login auth_url or token
+        """
+        # from .pygraphistry import PyGraphistry
+        base_path = self.server_base_path
+
+        if org_name is None and idp_name is None:
+            print("Login to site wide SSO")
+            url = f'{base_path}/api/v2/g/sso/oidc/login/'
+        elif org_name is not None and idp_name is None:
+            print("Login to {} organization level SSO".format(org_name))
+            url = f'{base_path}/api/v2/o/{org_name}/sso/oidc/login/'
+        elif org_name is not None and idp_name is not None:
+            print("Login to {} idp {} SSO".format(org_name, idp_name))
+            url = f'{base_path}/api/v2/o/{org_name}/sso/oidc/login/{idp_name}/'
+
+        # print("url : {}".format(url))
+        out = requests.post(
+            url, data={'client-type': 'pygraphistry'},
+            verify=self.certificate_validation
+        )
+        # print(out.text)
+        json_response = None
+        try:
+            json_response = out.json()
+            logger.debug("@ArrowUploader.sso_login, json_response: {}".format(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']))
+                    if 'state' in json_response['data']:
+                        self.sso_state = json_response['data']['state']
+                        self.sso_auth_url = json_response['data']['auth_url']
+                    else:
+                        self.token = json_response['data']['token']
+                elif json_response['status'] == 'ERR':
+                    raise Exception(json_response['message'])
+
+        except Exception:
+            logger.error('Error: %s', out, exc_info=True)
+            raise
+
+        return self
+
+    def sso_get_token(self, state):
+        """
+        Koa, 04 May 2022    Use state to get token
+        """
+
+        # 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
+        )
+        json_response = None
+        try:
+            json_response = out.json()
+            # print("get_jwt : {}".format(json_response))
+            self.token = None
+            if not ('status' in json_response):
+                raise Exception(out.text)
+            else:
+                if json_response['status'] == 'OK':
+                    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']))
+                        self.org_name = json_response['data']['active_organization']['slug']
+
+        except Exception:
+            logger.error('Error: %s', out, exc_info=True)
+            # raise
+
+        return self
+
+
     def refresh(self, token=None):
         if token is None:
             token = self.token
@@ -240,10 +378,9 @@ def verify(self, token=None) -> bool:
 
     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))
         res = requests.post(
             self.server_base_path + '/api/v2/upload/datasets/',
             verify=self.certificate_validation,
@@ -351,6 +488,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))
         if as_files:
 
             file_uploader = ArrowFileUploader(self)

graphistry/bolt_util.py

@@ -82,7 +82,7 @@ def flatten_spatial_col(df : pd.DataFrame, col : str) -> pd.DataFrame:  # noqa:
     for prop in ['x', 'y', 'z', 'srid', 'longtitude', 'latitude', 'height']:
         try:
             # v4.x + v5.x
-            s = df[col].apply(lambda v: getattr(v, prop, None))
+            s = df[col].apply(lambda v: getattr(v, prop, None))  # type: ignore
             if len(s.dropna()) > 0:
                 out_df[f'{col}_{prop}'] = s
         except: