Skip to content
This repository was archived by the owner on Sep 22, 2023. It is now read-only.

Commit 5480412

Browse files
achimnolachimnol
authored
fix: Support new standard-complient GQL endpoint (#168)
* refactor: Reuse Admin._query() for all other GQL invocations
1 parent a62badc commit 5480412

File tree

13 files changed

+140
-420
lines changed

13 files changed

+140
-420
lines changed

‎changes/168.fix

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Support the new standard-compliant GQL endpoint in the manager with the API version v6.20210815

‎src/ai/backend/client/config.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ class Undefined(enum.Enum):
3737
_config = None
3838
_undefined = Undefined.token
3939

40-
API_VERSION = (6, '20200815')
40+
API_VERSION = (6, '20210815')
4141

4242
DEFAULT_CHUNK_SIZE = 16 * (2**20) # 16 MiB
4343
MAX_INFLIGHT_CHUNKS = 4

‎src/ai/backend/client/func/admin.py

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
from typing import Any, Mapping, Optional
22

33
from .base import api_function, BaseFunction
4+
from ..exceptions import BackendAPIError
45
from ..request import Request
6+
from ..session import api_session
57

68
__all__ = (
79
'Admin',
@@ -22,7 +24,8 @@ class Admin(BaseFunction):
2224
@api_function
2325
@classmethod
2426
async def query(
25-
cls, query: str,
27+
cls,
28+
query: str,
2629
variables: Optional[Mapping[str, Any]] = None,
2730
) -> Any:
2831
"""
@@ -35,11 +38,34 @@ async def query(
3538
3639
:returns: The object parsed from the response JSON string.
3740
"""
41+
return await cls._query(query, variables)
42+
43+
@classmethod
44+
async def _query(
45+
cls,
46+
query: str,
47+
variables: Optional[Mapping[str, Any]] = None,
48+
) -> Any:
49+
"""
50+
Internal async implementation of the query() method,
51+
which may be reused by other functional APIs to make GQL requests.
52+
"""
3853
gql_query = {
3954
'query': query,
4055
'variables': variables if variables else {},
4156
}
42-
rqst = Request('POST', '/admin/graphql')
43-
rqst.set_json(gql_query)
44-
async with rqst.fetch() as resp:
45-
return await resp.json()
57+
if api_session.get().api_version >= (6, '20210815'):
58+
rqst = Request('POST', '/admin/gql')
59+
rqst.set_json(gql_query)
60+
async with rqst.fetch() as resp:
61+
response = await resp.json()
62+
errors = response.get("errors", [])
63+
if errors:
64+
raise BackendAPIError(400, reason="GraphQL-generated error", data=errors)
65+
else:
66+
return response["data"]
67+
else:
68+
rqst = Request('POST', '/admin/graphql')
69+
rqst.set_json(gql_query)
70+
async with rqst.fetch() as resp:
71+
return await resp.json()

‎src/ai/backend/client/func/agent.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from .base import api_function, BaseFunction
88
from ..request import Request
9+
from ..session import api_session
910
from ..pagination import generate_paginated_results
1011

1112
__all__ = (
@@ -86,14 +87,8 @@ async def detail(
8687
""")
8788
query = query.replace('$fields', ' '.join(fields))
8889
variables = {'agent_id': agent_id}
89-
rqst = Request('POST', '/admin/graphql')
90-
rqst.set_json({
91-
'query': query,
92-
'variables': variables,
93-
})
94-
async with rqst.fetch() as resp:
95-
data = await resp.json()
96-
return data['agent']
90+
data = await api_session.get().Admin._query(query, variables)
91+
return data['agent']
9792

9893

9994
class AgentWatcher(BaseFunction):

‎src/ai/backend/client/func/domain.py

Lines changed: 13 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from typing import Iterable, Sequence
33

44
from .base import api_function, BaseFunction
5-
from ..request import Request
5+
from ..session import api_session
66

77
__all__ = (
88
'Domain',
@@ -39,13 +39,8 @@ async def list(cls, fields: Iterable[str] = None) -> Sequence[dict]:
3939
}
4040
""")
4141
query = query.replace('$fields', ' '.join(fields))
42-
rqst = Request('POST', '/admin/graphql')
43-
rqst.set_json({
44-
'query': query,
45-
})
46-
async with rqst.fetch() as resp:
47-
data = await resp.json()
48-
return data['domains']
42+
data = await api_session.get().Admin._query(query)
43+
return data['domains']
4944

5045
@api_function
5146
@classmethod
@@ -67,14 +62,8 @@ async def detail(cls, name: str, fields: Iterable[str] = None) -> Sequence[dict]
6762
""")
6863
query = query.replace('$fields', ' '.join(fields))
6964
variables = {'name': name}
70-
rqst = Request('POST', '/admin/graphql')
71-
rqst.set_json({
72-
'query': query,
73-
'variables': variables,
74-
})
75-
async with rqst.fetch() as resp:
76-
data = await resp.json()
77-
return data['domain']
65+
data = await api_session.get().Admin._query(query, variables)
66+
return data['domain']
7867

7968
@api_function
8069
@classmethod
@@ -109,14 +98,8 @@ async def create(cls, name: str, description: str = '', is_active: bool = True,
10998
'integration_id': integration_id,
11099
},
111100
}
112-
rqst = Request('POST', '/admin/graphql')
113-
rqst.set_json({
114-
'query': query,
115-
'variables': variables,
116-
})
117-
async with rqst.fetch() as resp:
118-
data = await resp.json()
119-
return data['create_domain']
101+
data = await api_session.get().Admin._query(query, variables)
102+
return data['create_domain']
120103

121104
@api_function
122105
@classmethod
@@ -149,14 +132,8 @@ async def update(cls, name: str, new_name: str = None, description: str = None,
149132
'integration_id': integration_id,
150133
},
151134
}
152-
rqst = Request('POST', '/admin/graphql')
153-
rqst.set_json({
154-
'query': query,
155-
'variables': variables,
156-
})
157-
async with rqst.fetch() as resp:
158-
data = await resp.json()
159-
return data['modify_domain']
135+
data = await api_session.get().Admin._query(query, variables)
136+
return data['modify_domain']
160137

161138
@api_function
162139
@classmethod
@@ -172,14 +149,8 @@ async def delete(cls, name: str):
172149
}
173150
""")
174151
variables = {'name': name}
175-
rqst = Request('POST', '/admin/graphql')
176-
rqst.set_json({
177-
'query': query,
178-
'variables': variables,
179-
})
180-
async with rqst.fetch() as resp:
181-
data = await resp.json()
182-
return data['delete_domain']
152+
data = await api_session.get().Admin._query(query, variables)
153+
return data['delete_domain']
183154

184155
@api_function
185156
@classmethod
@@ -195,11 +166,5 @@ async def purge(cls, name: str):
195166
}
196167
""")
197168
variables = {'name': name}
198-
rqst = Request('POST', '/admin/graphql')
199-
rqst.set_json({
200-
'query': query,
201-
'variables': variables,
202-
})
203-
async with rqst.fetch() as resp:
204-
data = await resp.json()
205-
return data['purge_domain']
169+
data = await api_session.get().Admin._query(query, variables)
170+
return data['purge_domain']

‎src/ai/backend/client/func/group.py

Lines changed: 9 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from typing import Iterable, Sequence
33

44
from .base import api_function, BaseFunction
5-
from ..request import Request
5+
from ..session import api_session
66

77
__all__ = (
88
'Group',
@@ -42,13 +42,7 @@ async def list(cls, domain_name: str,
4242
""")
4343
query = query.replace('$fields', ' '.join(fields))
4444
variables = {'domain_name': domain_name}
45-
rqst = Request('POST', '/admin/graphql')
46-
rqst.set_json({
47-
'query': query,
48-
'variables': variables,
49-
})
50-
async with rqst.fetch() as resp:
51-
data = await resp.json()
45+
data = await api_session.get().Admin._query(query, variables)
5246
return data['groups']
5347

5448
@api_function
@@ -70,13 +64,7 @@ async def detail(cls, gid: str, fields: Iterable[str] = None) -> Sequence[dict]:
7064
""")
7165
query = query.replace('$fields', ' '.join(fields))
7266
variables = {'gid': gid}
73-
rqst = Request('POST', '/admin/graphql')
74-
rqst.set_json({
75-
'query': query,
76-
'variables': variables,
77-
})
78-
async with rqst.fetch() as resp:
79-
data = await resp.json()
67+
data = await api_session.get().Admin._query(query, variables)
8068
return data['group']
8169

8270
@api_function
@@ -111,13 +99,7 @@ async def create(cls, domain_name: str, name: str, description: str = '',
11199
'integration_id': integration_id,
112100
},
113101
}
114-
rqst = Request('POST', '/admin/graphql')
115-
rqst.set_json({
116-
'query': query,
117-
'variables': variables,
118-
})
119-
async with rqst.fetch() as resp:
120-
data = await resp.json()
102+
data = await api_session.get().Admin._query(query, variables)
121103
return data['create_group']
122104

123105
@api_function
@@ -149,13 +131,7 @@ async def update(cls, gid: str, name: str = None, description: str = None,
149131
'integration_id': integration_id,
150132
},
151133
}
152-
rqst = Request('POST', '/admin/graphql')
153-
rqst.set_json({
154-
'query': query,
155-
'variables': variables,
156-
})
157-
async with rqst.fetch() as resp:
158-
data = await resp.json()
134+
data = await api_session.get().Admin._query(query, variables)
159135
return data['modify_group']
160136

161137
@api_function
@@ -172,13 +148,7 @@ async def delete(cls, gid: str):
172148
}
173149
""")
174150
variables = {'gid': gid}
175-
rqst = Request('POST', '/admin/graphql')
176-
rqst.set_json({
177-
'query': query,
178-
'variables': variables,
179-
})
180-
async with rqst.fetch() as resp:
181-
data = await resp.json()
151+
data = await api_session.get().Admin._query(query, variables)
182152
return data['delete_group']
183153

184154
@api_function
@@ -195,13 +165,7 @@ async def purge(cls, gid: str):
195165
}
196166
""")
197167
variables = {'gid': gid}
198-
rqst = Request('POST', '/admin/graphql')
199-
rqst.set_json({
200-
'query': query,
201-
'variables': variables,
202-
})
203-
async with rqst.fetch() as resp:
204-
data = await resp.json()
168+
data = await api_session.get().Admin._query(query, variables)
205169
return data['purge_group']
206170

207171
@api_function
@@ -226,13 +190,7 @@ async def add_users(cls, gid: str, user_uuids: Iterable[str],
226190
'user_uuids': user_uuids,
227191
},
228192
}
229-
rqst = Request('POST', '/admin/graphql')
230-
rqst.set_json({
231-
'query': query,
232-
'variables': variables,
233-
})
234-
async with rqst.fetch() as resp:
235-
data = await resp.json()
193+
data = await api_session.get().Admin._query(query, variables)
236194
return data['modify_group']
237195

238196
@api_function
@@ -257,11 +215,5 @@ async def remove_users(cls, gid: str, user_uuids: Iterable[str],
257215
'user_uuids': user_uuids,
258216
},
259217
}
260-
rqst = Request('POST', '/admin/graphql')
261-
rqst.set_json({
262-
'query': query,
263-
'variables': variables,
264-
})
265-
async with rqst.fetch() as resp:
266-
data = await resp.json()
218+
data = await api_session.get().Admin._query(query, variables)
267219
return data['modify_group']

0 commit comments

Comments
 (0)