11from .endpoint import Endpoint , api
2- from .exceptions import GraphQLError
3-
2+ from .exceptions import GraphQLError , InvalidGraphQLQuery
43import logging
54import json
65
76logger = logging .getLogger ('tableau.endpoint.metadata' )
87
98
9+ def is_valid_paged_query (parsed_query ):
10+ """Check that the required $first and $afterToken variables are present in the query.
11+ Also check that we are asking for the pageInfo object, so we get the endCursor. There
12+ is no way to do this relilably without writing a GraphQL parser, so simply check that
13+ that the string contains 'hasNextPage' and 'endCursor'"""
14+ return all (k in parsed_query ['variables' ] for k in ('first' , 'afterToken' )) and \
15+ 'hasNextPage' in parsed_query ['query' ] and \
16+ 'endCursor' in parsed_query ['query' ]
17+
18+
19+ def extract_values (obj , key ):
20+ """Pull all values of specified key from nested JSON.
21+ Taken from: https://hackersandslackers.com/extract-data-from-complex-json-python/"""
22+ arr = []
23+
24+ def extract (obj , arr , key ):
25+ """Recursively search for values of key in JSON tree."""
26+ if isinstance (obj , dict ):
27+ for k , v in obj .items ():
28+ if isinstance (v , (dict , list )):
29+ extract (v , arr , key )
30+ elif k == key :
31+ arr .append (v )
32+ elif isinstance (obj , list ):
33+ for item in obj :
34+ extract (item , arr , key )
35+ return arr
36+
37+ results = extract (obj , arr , key )
38+ return results
39+
40+
41+ def get_page_info (result ):
42+ next_page = extract_values (result , 'hasNextPage' ).pop ()
43+ cursor = extract_values (result , 'endCursor' ).pop ()
44+ return next_page , cursor
45+
46+
1047class Metadata (Endpoint ):
1148 @property
1249 def baseurl (self ):
1350 return "{0}/api/metadata/graphql" .format (self .parent_srv .server_address )
1451
15- @api ("3.2 " )
52+ @api ("3.5 " )
1653 def query (self , query , variables = None , abort_on_error = False ):
1754 logger .info ('Querying Metadata API' )
1855 url = self .baseurl
1956
2057 try :
2158 graphql_query = json .dumps ({'query' : query , 'variables' : variables })
22- except Exception :
23- # Place holder for now
24- raise Exception ('Must provide a string' )
59+ except Exception as e :
60+ raise InvalidGraphQLQuery ('Must provide a string' )
2561
2662 # Setting content type because post_reuqest defaults to text/xml
2763 server_response = self .post_request (url , graphql_query , content_type = 'text/json' )
@@ -31,3 +67,55 @@ def query(self, query, variables=None, abort_on_error=False):
3167 raise GraphQLError (results ['errors' ])
3268
3369 return results
70+
71+ @api ("3.5" )
72+ def paginated_query (self , query , variables = None , abort_on_error = False ):
73+ logger .info ('Querying Metadata API using a Paged Query' )
74+ url = self .baseurl
75+
76+ if variables is None :
77+ # default paramaters
78+ variables = {'first' : 100 , 'afterToken' : None }
79+ elif (('first' in variables ) and ('afterToken' not in variables )):
80+ # they passed a page size but not a token, probably because they're starting at `null` token
81+ variables .update ({'afterToken' : None })
82+
83+ graphql_query = json .dumps ({'query' : query , 'variables' : variables })
84+ parsed_query = json .loads (graphql_query )
85+
86+ if not is_valid_paged_query (parsed_query ):
87+ raise InvalidGraphQLQuery ('Paged queries must have a `$first` and `$afterToken` variables as well as '
88+ 'a pageInfo object with `endCursor` and `hasNextPage`' )
89+
90+ results_dict = {'pages' : []}
91+ paginated_results = results_dict ['pages' ]
92+
93+ # get first page
94+ server_response = self .post_request (url , graphql_query , content_type = 'text/json' )
95+ results = server_response .json ()
96+
97+ if abort_on_error and results .get ('errors' , None ):
98+ raise GraphQLError (results ['errors' ])
99+
100+ paginated_results .append (results )
101+
102+ # repeat
103+ has_another_page , cursor = get_page_info (results )
104+
105+ while has_another_page :
106+ # Update the page
107+ variables .update ({'afterToken' : cursor })
108+ # make the call
109+ logger .debug ("Calling Token: " + cursor )
110+ graphql_query = json .dumps ({'query' : query , 'variables' : variables })
111+ server_response = self .post_request (url , graphql_query , content_type = 'text/json' )
112+ results = server_response .json ()
113+ # verify response
114+ if abort_on_error and results .get ('errors' , None ):
115+ raise GraphQLError (results ['errors' ])
116+ # save results and repeat
117+ paginated_results .append (results )
118+ has_another_page , cursor = get_page_info (results )
119+
120+ logger .info ('Sucessfully got all results for paged query' )
121+ return results_dict
0 commit comments