Graphistry 2.0 REST API Tutorial: Login, create dataset, and upload nodes/edges as json, csv, parquet, arrow, and more

uploader.ipynb

uploader.py

import graphistry, io, json, logging, pandas as pd, pyarrow as pa, requests

logger = logging.getLogger('Uploader')

class Uploader:
    
    @property
    def token(self) -> str:
        if self.__token is None:
            raise Exception("Not logged in")
        return self.__token

    @property
    def dataset_id(self) -> str:
        if self.__dataset_id is None:
            raise Exception("Must first create a dataset")
        return self.__dataset_id

    @property
    def base_path(self) -> str:
        return self.__base_path

    @property
    def view_base_path(self) -> str:
        return self.__view_base_path

    @property
    def url_params(self) -> dict:
        if self.__url_params is None:
            return {}
        else:
            return self.__url_params
        
    def settings(self, url_params=None):
        if not (url_params is None):
            self.__url_params = url_params
        return self
    
    def __init__(self, base_path='http://nginx', view_base_path='http://localhost'):
        self.__base_path = base_path
        self.__view_base_path = view_base_path
        self.__token = None
        self.__dataset_id = None
        self.__url_params = None
    
    def login(self, username, password):   
        base_path = self.base_path
        out = requests.post(
            f'{base_path}/api-token-auth/',
            json={'username': username, 'password': password})
        json_response = None
        try:
            json_response = out.json()
            if not ('token' in json_response):
                raise Exception(out.text)
        except Exception as e:
            logger.error('Error: %s', out)
            raise Exception(out.text)
            
        self.__token = out.json()['token']        
        return self
    
    def create_dataset(self, json):
        tok = self.token 
        
        out = requests.post(
            self.base_path + '/api/v2/upload/datasets/',
             headers={'Authorization': f'Bearer {tok}'},
             json=json).json()
        
        if not out['success']:
            raise Exception(out)
        
        self.__dataset_id = out['data']['dataset_id']

        return out
        
    #PyArrow's table.getvalues().to_pybytes() fails to hydrate some reason, 
    #  so work around by consolidate into a virtual file and sending that
    def arrow_to_buffer(self, table: pa.Table):
        b = io.BytesIO()
        writer = pa.RecordBatchFileWriter(b, table.schema)
        writer.write_table(table)
        writer.close()
        return b.getvalue()
    
    def post_g(self, g, name=None):
        
        def maybe_bindings(g, bindings):
            out = {}
            for old_field_name, new_field_name in bindings:
                try:
                    val = getattr(g, old_field_name)
                    if val is None:
                        continue
                    else:
                        out[new_field_name] = val
                except AttributeError:
                    continue
            logger.debug('bindings: %s', out)
            return out
        
        self.__url_params = g._url_params if not (g._url_params is None) else {}
        
        node_encodings = maybe_bindings(
                g,
                [
                    ['_node', 'node'],
                    ['_point_color', 'node_color'],
                    ['_point_label', 'node_label'],
                    ['_point_opacity', 'node_opacity'],
                    ['_point_size', 'node_size'],                    
                    ['_point_title', 'node_title'],
                    ['_point_weight', 'node_weight']
                ])
        if not (g._nodes is None):
            if 'x' in g._nodes:
                node_encodings['x'] = 'x'
            if 'y' in g._nodes:
                node_encodings['y'] = 'y'
        
        
        self.create_dataset({
            "node_encodings": {"bindings": node_encodings},
            "edge_encodings": {"bindings": maybe_bindings(
                g,
                [
                    ['_source', 'source'],
                    ['_destination', 'destination'],
                    ['_edge_color', 'edge_color'],
                    ['_edge_label', 'edge_label'],
                    ['_edge_opacity', 'edge_opacity'],
                    ['_edge_size', 'edge_size'],
                    ['_edge_title', 'edge_title'],
                    ['_edge_weight', 'edge_weight']
                ])
            },
            "metadata": {},
            "name": ("mytestviz" if name is None else name)
        })
        
        self.g_post_edges(g)
        
        if not (g._nodes is None):
            self.g_post_nodes(g)
        
        return self
    
    def to_url(self, view_base_path=None):
        path = view_base_path if not (view_base_path is None) else self.view_base_path
        dataset_id = self.dataset_id
        params = [ str(k) + '=' + str(v) for k, v in self.url_params.items() ]
        url_params = ('&' + '&'.join(params)) if len(params) > 0 else '' 
            
        return f'{path}/graph/graph.html?dataset={dataset_id}{url_params}'
    
    def plot(self, render=True):
        if render:
            try:
                from IPython.core.display import display, HTML
                url = self.to_url()
                logger.debug('url: %s', url)
                return display(HTML(f'<iframe src="{url}" width="100%" height="600"/>'))
            except Exception as e:
                logger.debug(e)
        return self.to_url()
            
    def g_post_edges(self, g):
        
        arr = pa.Table.from_pandas(g._edges, preserve_index=False).replace_schema_metadata({})
        buf = self.arrow_to_buffer(arr)
        
        dataset_id = self.dataset_id
        tok = self.token
        base_path = self.base_path

        out = requests.post(
            f'{base_path}/api/v2/upload/datasets/{dataset_id}/edges/arrow',
            headers={'Authorization': f'Bearer {tok}'},
            data=buf).json()
        
        if not out['success']:
            raise Exception(out)
        
        return out
    
    def g_post_nodes(self, g):

        arr = pa.Table.from_pandas(g._nodes, preserve_index=False).replace_schema_metadata({})
        buf = self.arrow_to_buffer(arr)

        dataset_id = self.dataset_id
        tok = self.token
        base_path = self.base_path

        out = requests.post(
            f'{base_path}/api/v2/upload/datasets/{dataset_id}/nodes/arrow',
            headers={'Authorization': f'Bearer {tok}'},
            data=buf).json()
        
        if not out['success']:
            raise Exception(out)

        return out
    
    def post_edges_arrow(self, arr, opts=''):
        return self.post_arrow(arr, 'edges', opts) 

    def post_nodes_arrow(self, arr, opts=''):
        return self.post_arrow(arr, 'nodes', opts) 

    def post_arrow(self, arr, graph_type, opts=''):
        buf = self.arrow_to_buffer(arr)

        dataset_id = self.dataset_id
        tok = self.token
        base_path = self.base_path

        url = f'{base_path}/api/v2/upload/datasets/{dataset_id}/{graph_type}/arrow'
        if len(opts) > 0:
            url = f'{url}?{opts}'
        out = requests.post(
            url,
            headers={'Authorization': f'Bearer {tok}'},
            data=buf).json()
        
        if not out['success']:
            raise Exception(out)
            
        return out

    def post_edges_file(self, file_path, file_type='csv'):
        return self.post_file(file_path, 'edges', file_type)

    def post_nodes_file(self, file_path, file_type='csv'):
        return self.post_file(file_path, 'nodes', file_type)

    def post_file(self, file_path, graph_type='edges', file_type='csv'):

        dataset_id = self.dataset_id
        tok = self.token
        base_path = self.base_path

        with open(file_path, 'rb') as file:        
            out = requests.post(
                f'{base_path}/api/v2/upload/datasets/{dataset_id}/{graph_type}/{file_type}',
                headers={'Authorization': f'Bearer {tok}'},
                data=file.read()).json()
            if not out['success']:
                raise Exception(out)
            
            return out