77
88import requests
99from retrying import retry
10+ from warnings import warn
1011
1112from jupiterone .errors import (
1213 JupiterOneClientError ,
2223 DELETE_ENTITY ,
2324 UPDATE_ENTITY ,
2425 CREATE_RELATIONSHIP ,
25- DELETE_RELATIONSHIP
26+ DELETE_RELATIONSHIP ,
27+ CURSOR_QUERY_V1
2628)
2729
30+
2831def retry_on_429 (exc ):
2932 """ Used to trigger retry on rate limit """
3033 return isinstance (exc , JupiterOneApiRetryError )
@@ -68,12 +71,12 @@ def account(self, value: str):
6871
6972 @property
7073 def token (self ):
71- """ Your JupiteOne access token """
74+ """ Your JupiterOne access token """
7275 return self ._token
7376
7477 @token .setter
7578 def token (self , value : str ):
76- """ Your JupiteOne access token """
79+ """ Your JupiterOne access token """
7780 if not value :
7881 raise JupiterOneClientError ('token is required' )
7982 self ._token = value
@@ -89,11 +92,11 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict:
8992 if variables :
9093 data .update (variables = variables )
9194
92- response = requests .post (self .query_endpoint , headers = self .headers , json = data )
95+ response = requests .post (self .query_endpoint , headers = self .headers , json = data , timeout = 60 )
9396
9497 # It is still unclear if all responses will have a status
9598 # code of 200 or if 429 will eventually be used to
96- # indicate rate limitting . J1 devs are aware.
99+ # indicate rate limits being hit . J1 devs are aware.
97100 if response .status_code == 200 :
98101 if response ._content :
99102 content = json .loads (response ._content )
@@ -108,29 +111,59 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict:
108111 elif response .status_code == 401 :
109112 raise JupiterOneApiError ('401: Unauthorized. Please supply a valid account id and API token.' )
110113
111- elif response .status_code in [429 , 500 ]:
114+ elif response .status_code in [429 , 503 ]:
112115 raise JupiterOneApiRetryError ('JupiterOne API rate limit exceeded' )
113116
114- else :
115- try :
116- content = json .loads (response ._content )
117- raise JupiterOneApiError ('{}: {}' .format (response .status_code , content .get ('error' ) or 'Unknown Error' ))
118- except ValueError as e :
119- raise JupiterOneApiError ('{}: {}' .format (response .status_code , 'Unknown Error' ));
120-
117+ elif response .status_code in [500 ]:
118+ raise JupiterOneApiError ('JupiterOne API internal server error.' )
121119
122- def query_v1 (self , query : str , ** kwargs ) -> Dict :
123- """ Performs a V1 graph query
120+ else :
121+ content = response ._content
122+ if isinstance (content , (bytes , bytearray )):
123+ content = content .decode ("utf-8" )
124+ if 'application/json' in response .headers .get ('Content-Type' , 'text/plain' ):
125+ data = json .loads (content )
126+ content = data .get ('error' , data .get ('errors' , content ))
127+ raise JupiterOneApiError ('{}:{}' .format (response .status_code , content ))
128+
129+ def _cursor_query (self , query : str , cursor : str = None , include_deleted : bool = False ) -> Dict :
130+ """ Performs a V1 graph query using cursor pagination
124131 args:
125132 query (str): Query text
126- skip (int): Skip entity count
127- limit (int): Limit entity count
133+ cursor (str): A pagination cursor for the initial query
128134 include_deleted (bool): Include recently deleted entities in query/search
129135 """
130- skip : int = kwargs .pop ('skip' , J1QL_SKIP_COUNT )
131- limit : int = kwargs .pop ('limit' , J1QL_LIMIT_COUNT )
132- include_deleted : bool = kwargs .pop ('include_deleted' , False )
133136
137+ results : List = []
138+ while True :
139+ variables = {
140+ 'query' : query ,
141+ 'includeDeleted' : include_deleted
142+ }
143+
144+ if cursor is not None :
145+ variables ['cursor' ] = cursor
146+
147+ response = self ._execute_query (query = CURSOR_QUERY_V1 , variables = variables )
148+ data = response ['data' ]['queryV1' ]['data' ]
149+
150+ if 'vertices' in data and 'edges' in data :
151+ return data
152+
153+ results .extend (data )
154+
155+ if 'cursor' in response ['data' ]['queryV1' ] and response ['data' ]['queryV1' ]['cursor' ] is not None :
156+ cursor = response ['data' ]['queryV1' ]['cursor' ]
157+ else :
158+ break
159+
160+ return {'data' : results }
161+
162+ def _limit_and_skip_query (self ,
163+ query : str ,
164+ skip : int = J1QL_SKIP_COUNT ,
165+ limit : int = J1QL_LIMIT_COUNT ,
166+ include_deleted : bool = False ) -> Dict :
134167 results : List = []
135168 page : int = 0
136169
@@ -159,6 +192,39 @@ def query_v1(self, query: str, **kwargs) -> Dict:
159192
160193 return {'data' : results }
161194
195+ def query_v1 (self , query : str , ** kwargs ) -> Dict :
196+ """ Performs a V1 graph query
197+ args:
198+ query (str): Query text
199+ skip (int): Skip entity count
200+ limit (int): Limit entity count
201+ cursor (str): A pagination cursor for the initial query
202+ include_deleted (bool): Include recently deleted entities in query/search
203+ """
204+ uses_limit_and_skip : bool = 'skip' in kwargs .keys () or 'limit' in kwargs .keys ()
205+ skip : int = kwargs .pop ('skip' , J1QL_SKIP_COUNT )
206+ limit : int = kwargs .pop ('limit' , J1QL_LIMIT_COUNT )
207+ include_deleted : bool = kwargs .pop ('include_deleted' , False )
208+ cursor : str = kwargs .pop ('cursor' , None )
209+
210+ if uses_limit_and_skip :
211+ warn ('limit and skip pagination is no longer a recommended method for pagination. '
212+ 'To read more about using cursors checkout the JupiterOne documentation: '
213+ 'https://docs.jupiterone.io/features/admin/parameters#query-parameterlist' ,
214+ DeprecationWarning , stacklevel = 2 )
215+ return self ._limit_and_skip_query (
216+ query = query ,
217+ skip = skip ,
218+ limit = limit ,
219+ include_deleted = include_deleted
220+ )
221+ else :
222+ return self ._cursor_query (
223+ query = query ,
224+ cursor = cursor ,
225+ include_deleted = include_deleted
226+ )
227+
162228 def create_entity (self , ** kwargs ) -> Dict :
163229 """ Creates an entity in graph. It will also update an existing entity.
164230
@@ -206,7 +272,7 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict:
206272 Update an existing entity.
207273
208274 args:
209- entity_id (str): The _id of the entity to udate
275+ entity_id (str): The _id of the entity to update
210276 properties (dict): Dictionary of key/value entity properties
211277 """
212278 variables = {
@@ -218,7 +284,7 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict:
218284
219285 def create_relationship (self , ** kwargs ) -> Dict :
220286 """
221- Create a relationship (edge) between two entities (veritces ).
287+ Create a relationship (edge) between two entities (vertices ).
222288
223289 args:
224290 relationship_key (str): Unique key for the relationship
0 commit comments