async def manage_runners(runners_uri, runners, session: aiohttp.ClientSession):
# To find and enable/disable gitlab-runners
# https://docs.gitlab.com/ee/api/runners.html#list-projects-runners
# https://docs.gitlab.com/ee/api/runners.html#enable-a-runner-in-project
# https://docs.gitlab.com/ee/api/runners.html#disable-a-runner-from-project
for runner in runners:
if not isinstance(runner, dict):
continue
if runner.get("is_shared"):
continue
runner_id = runner.get("id")
if runner_id == DISABLE_RUNNER_ID:
print(DRY_RUN + "disable: " + runner["name"])
if not DRY_RUN:
async with session.delete(runners_uri + f"/{runner_id}") as resp:
print(resp.status)
print()
elif runner_id == ENABLE_RUNNER_ID:
print(DRY_RUN + "enable: " + runner["name"])
if not DRY_RUN:
async with session.post(
runners_uri, params={"runner_id": runner_id}
) as resp:
print(resp.status)
print()
else:
print(runner["description"])
class HTTP:
def __init__(self):
self.sess = ClientSession(headers={'Authorization': f'Bot {TOKEN}'})
async def init(self):
if self.sess.closed:
self.sess = ClientSession(
headers={'Authorization': f'Bot {TOKEN}'})
async def pong(self, url: str):
async with self.sess.get(url) as resp:
return str(resp.status).startswith('2')
async def send_message(self,
channel: str,
message: str,
allowed_mentions: dict = {}):
async with self.sess.post(BASE + f"/channels/{channel}/messages",
json={
"content": message,
**allowed_mentions
}) as resp:
return await resp.json()
async def add_role(self, guild, user, role):
async with self.sess.put(
BASE + f"/guilds/{guild}/members/{user}/roles/{role}") as resp:
return resp.status == 204
async def del_role(self, guild, user, role):
async with self.sess.delete(
BASE + f"/guilds/{guild}/members/{user}/roles/{role}") as resp:
return resp.status == 204
async def abort_multipart_upload(self, session: aiohttp.ClientSession) -> None:
"""Method to abort a multipart upload
Args:
session: The current aiohttp session
Returns:
None
"""
if not self.multipart_upload_id:
raise ValueError("An upload id is required to abort a multipart upload")
try_count = 0
error_status = None
error_msg = None
while try_count < 3:
async with session.delete(f"{self.service_root}/{self.object_id}/multipart/{self.multipart_upload_id}",
headers=self.object_service_headers) as response:
if response.status != 204:
# An error occurred
logger.warning(f"Failed to abort {self.object_details.dataset_path}. Try {try_count + 1} of 3")
error_msg = await response.text()
error_status = response.status
await asyncio.sleep(try_count ** 2)
try_count += 1
else:
# All good.
logger.info(f"Aborted multipart upload {self.multipart_upload_id} for "
f"{self.object_details.dataset_path} at {self.object_details.revision[0:8]}.")
return
raise IOError(f"Failed to abort multipart upload for {self.object_details.dataset_path}."
f" Status: {error_status}. Response: {error_msg}")
async def unsubscribe(session: aiohttp.ClientSession, sub_id: str) -> None:
headers = {
'X-MPBX-API-AUTH-TOKEN': API_TOKEN,
}
url = ('https://cloudpbx.beeline.ru/apis/portal/subscription?'
f'subscriptionId={sub_id}')
async with session.delete(url, headers=headers):
pass
async def delete_invite(
self,
id_: str,
*,
session: aiohttp.ClientSession,
) -> None:
async with session.delete(
self._admin_v1_endpoint("/invites/{}".format(id_)), ) as resp:
self._raise_error_from_response(resp)
async def _delete(cls,
session: aiohttp.ClientSession,
url_fragment: str,
data: Union[dict, list] = None,
**kwargs: Any) -> Any:
url = f"{Settings.rest_url}/{url_fragment}"
async with session.delete(url, json=data, **kwargs) as resp:
return resp.status, await resp.json()
async def delete_user_by_localpart(
self,
localpart: str,
*,
session: aiohttp.ClientSession,
) -> None:
async with session.delete(
self._admin_v1_endpoint("/users/{}".format(localpart)),
) as resp:
self._raise_error_from_response(resp)
async def remove_group_member(
self,
id_: str,
localpart: str,
*,
session: aiohttp.ClientSession,
) -> None:
async with session.delete(
self._admin_v1_endpoint("/groups/{}/members/{}".format(
id_, localpart)), ) as resp:
self._raise_error_from_response(resp)
class Client:
_client = None
def __init__(self, loop, url=None):
self._client = ClientSession(loop=loop)
self._url = url
@property
def cli(self):
return self._client
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_value, traceback):
pass
def handler_url(self, url):
if url.startswith("http"):
return url
if self._url:
return "{}{}".format(self._url, url)
return url
def request(self, method, url, *args, **kwargs):
return self._client.request(method, self.handler_url(url), *args,
**kwargs)
def get(self, url, allow_redirects=True, **kwargs):
return self._client.get(self.handler_url(url),
allow_redirects=True,
**kwargs)
def post(self, url, data=None, **kwargs):
return self._client.post(self.handler_url(url), data=data, **kwargs)
def put(self, url, data=None, **kwargs):
return self._client.put(self.handler_url(url), data=data, **kwargs)
def delete(self, url, **kwargs):
return self._client.delete(self.handler_url(url), **kwargs)
def head(self, url, allow_redirects=False, **kwargs):
return self._client.head(self.handler_url(url),
allow_redirects=allow_redirects,
**kwargs)
def options(self, url, allow_redirects=True, **kwargs):
return self._client.options(self.handler_url(url),
allow_redirects=allow_redirects,
**kwargs)
def close(self):
self._client.close()
async def run_delete(session_to_use: ClientSession):
async with session_to_use.delete(
(base_url if len(records_to_delete) > 1 else base_url +
f"/{records_to_delete[0]}"),
params=({
"records": records_to_delete
} if len(records_to_delete) > 1 else None),
headers=self.auth_header,
) as r:
if r.status != 200:
log.warning(f"Failed to delete IDs: {records_to_delete}")
raise AirTableError(r.url, await r.json())
async def cancel_query_task(
self, session: aiohttp.ClientSession, query_task_id: str
):
""" Cancels a query task.
Args:
query_task_id: ID for the query task to cancel.
"""
logger.debug(f"Cancelling query task: {query_task_id}")
url = utils.compose_url(self.api_url, path=["running_queries", query_task_id])
async with session.delete(url=url) as response:
await response.read()
class HTTP:
def __init__(self):
self.sess = ClientSession(headers={'Authorization': f'Bot {TOKEN}'})
async def init(self):
if self.sess.closed:
self.sess = ClientSession(
headers={'Authorization': f'Bot {TOKEN}'})
async def request(self, method, url, json={}):
async with self.sess.request(method, BASE + url, json=json) as resp:
return await resp.json()
async def send_message(self, channel: str, message: str):
return await self.request("POST",
f"/channels/{channel}/messages",
json={
"content": message,
**AM
})
async def report(self, message: str):
return await self.send_message(CHANNEL, message)
async def get_user(self, server: str, userid: str):
async with self.sess.get(MCPBASE + f"/{server}/{userid}",
headers={"auth": MC_TOKEN}) as resp:
if resp.status > 299:
return False
return await resp.json()
async def grant_perms(self, server: str, group: str, userid: str,
mcname: str):
async with self.sess.post(MCPBASE +
f"/{server}/{group}/{userid}/{mcname}",
headers={"auth": MC_TOKEN}) as resp:
if resp.status > 299:
return False
return await resp.json()
async def revoke_perms(
self,
server: str,
userid: str,
):
async with self.sess.delete(MCPBASE + f"/{server}/{userid}",
headers={"auth": MC_TOKEN}) as resp:
if resp.status > 299:
return False
return await resp.json()
async def delete_file(session: ClientSession, file_id: str, location_id: str,
user_id: UserID) -> None:
if (not isinstance(file_id, str) or not isinstance(location_id, str)
or not isinstance(user_id, int)):
raise exceptions.StorageInvalidCall(
f"invalid call: user_id '{user_id}', location_id '{location_id}', file_id '{file_id}' are invalid",
)
if file_id is None or location_id is None or user_id is None:
raise exceptions.StorageInvalidCall(
f"invalid call: user_id '{user_id}', location_id '{location_id}', file_id '{file_id}' are not allowed to be empty",
)
async with session.delete(
f"{_base_url()}/locations/{location_id}/files/{quote(file_id, safe='')}",
params={"user_id": f"{user_id}"},
) as response:
response.raise_for_status()
class JupyterClient:
log: BoundLoggerLazyProxy
user: User
session: ClientSession
headers: Dict[str, str]
xsrftoken: str
jupyter_url: str
def __init__(self, user: User, log: BoundLoggerLazyProxy,
options: Dict[str, Any]):
self.user = user
self.log = log
self.jupyter_base = options.get("nb_url", "/nb/")
self.jupyter_url = Configuration.environment_url + self.jupyter_base
self.xsrftoken = "".join(
random.choices(string.ascii_uppercase + string.digits, k=16))
self.jupyter_options_form = options.get("jupyter_options_form", {})
self.headers = {
"Authorization": "Bearer " + user.token,
"x-xsrftoken": self.xsrftoken,
}
self.session = ClientSession(headers=self.headers)
self.session.cookie_jar.update_cookies(
BaseCookie({"_xsrf": self.xsrftoken}))
__ansi_reg_exp = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
@classmethod
def _ansi_escape(cls, line: str) -> str:
return cls.__ansi_reg_exp.sub("", line)
async def hub_login(self) -> None:
async with self.session.get(self.jupyter_url + "hub/login") as r:
if r.status != 200:
await self._raise_error("Error logging into hub", r)
async def ensure_lab(self) -> None:
self.log.info("Ensure lab")
running = await self.is_lab_running()
if running:
await self.lab_login()
else:
await self.spawn_lab()
async def lab_login(self) -> None:
self.log.info("Logging into lab")
lab_url = self.jupyter_url + f"user/{self.user.username}/lab"
async with self.session.get(lab_url) as r:
if r.status != 200:
await self._raise_error("Error logging into lab", r)
async def is_lab_running(self) -> bool:
self.log.info("Is lab running?")
hub_url = self.jupyter_url + "hub"
async with self.session.get(hub_url) as r:
if r.status != 200:
self.log.error(f"Error {r.status} from {r.url}")
spawn_url = self.jupyter_url + "hub/spawn"
self.log.info(f"Going to {hub_url} redirected to {r.url}")
if str(r.url) == spawn_url:
return False
return True
async def spawn_lab(self) -> None:
spawn_url = self.jupyter_url + "hub/spawn"
pending_url = (self.jupyter_url +
f"hub/spawn-pending/{self.user.username}")
lab_url = self.jupyter_url + f"user/{self.user.username}/lab"
# DM-23864: Do a get on the spawn URL even if I don't have to.
async with self.session.get(spawn_url) as r:
await r.text()
async with self.session.post(spawn_url,
data=self.jupyter_options_form,
allow_redirects=False) as r:
if r.status != 302:
await self._raise_error("Spawn did not redirect", r)
redirect_url = (self.jupyter_base +
f"hub/spawn-pending/{self.user.username}")
if r.headers["Location"] != redirect_url:
await self._raise_error("Spawn didn't redirect to pending", r)
# Jupyterlab will give up a spawn after 900 seconds, so we shouldn't
# wait longer than that.
max_poll_secs = 900
poll_interval = 15
retries = max_poll_secs / poll_interval
while retries > 0:
async with self.session.get(pending_url) as r:
if str(r.url) == lab_url:
self.log.info(f"Lab spawned, redirected to {r.url}")
return
if not r.ok:
await self._raise_error("Error spawning", r)
self.log.info(f"Still waiting for lab to spawn {r}")
retries -= 1
await asyncio.sleep(poll_interval)
raise Exception("Giving up waiting for lab to spawn!")
async def delete_lab(self) -> None:
headers = {"Referer": self.jupyter_url + "hub/home"}
server_url = (self.jupyter_url +
f"hub/api/users/{self.user.username}/server")
self.log.info(f"Deleting lab for {self.user.username} at {server_url}")
async with self.session.delete(server_url, headers=headers) as r:
if r.status not in [200, 202, 204]:
await self._raise_error("Error deleting lab", r)
async def create_kernel(self, kernel_name: str = "LSST") -> str:
kernel_url = (self.jupyter_url +
f"user/{self.user.username}/api/kernels")
body = {"name": kernel_name}
async with self.session.post(kernel_url, json=body) as r:
if r.status != 201:
await self._raise_error("Error creating kernel", r)
response = await r.json()
return response["id"]
async def run_python(self, kernel_id: str, code: str) -> str:
kernel_url = (
self.jupyter_url +
f"user/{self.user.username}/api/kernels/{kernel_id}/channels")
msg_id = uuid4().hex
msg = {
"header": {
"username": "",
"version": "5.0",
"session": "",
"msg_id": msg_id,
"msg_type": "execute_request",
},
"parent_header": {},
"channel": "shell",
"content": {
"code": code,
"silent": False,
"store_history": False,
"user_expressions": {},
"allow_stdin": False,
},
"metadata": {},
"buffers": {},
}
async with self.session.ws_connect(kernel_url) as ws:
await ws.send_json(msg)
while True:
r = await ws.receive_json()
self.log.debug(f"Recieved kernel message: {r}")
msg_type = r["msg_type"]
if msg_type == "error":
error_message = "".join(r["content"]["traceback"])
raise NotebookException(self._ansi_escape(error_message))
elif (msg_type == "stream"
and msg_id == r["parent_header"]["msg_id"]):
return r["content"]["text"]
elif msg_type == "execute_reply":
status = r["content"]["status"]
if status == "ok":
return ""
else:
raise NotebookException(
f"Error content status is {status}")
def dump(self) -> dict:
return {
"cookies": [str(cookie) for cookie in self.session.cookie_jar],
}
async def _raise_error(self, msg: str, r: ClientResponse) -> None:
raise Exception(f"{msg}: {r.status} {r.url}: {r.headers}")
async def delete_all(session: aiohttp.ClientSession):
async with session.delete('http://0.0.0.0:8080/') as resp:
assert resp.status == 200
class AioHttpClient(HttpClient):
def __init__(self,
*,
connector=None,
loop=None,
cookies=None,
headers=None,
skip_auto_headers=None,
auth=None,
json_serialize=json.dumps,
request_class=ClientRequest,
response_class=ClientResponse,
ws_response_class=ClientWebSocketResponse,
version=http.HttpVersion11,
cookie_jar=None,
connector_owner=True,
raise_for_status=False,
read_timeout=sentinel,
conn_timeout=None,
auto_decompress=True,
trust_env=False,
**kwargs):
"""
The class packaging a class ClientSession to perform HTTP request and manager that these HTTP connection.
For details of the params: http://aiohttp.readthedocs.io/en/stable/client_advanced.html#client-session
"""
super(AioHttpClient, self).__init__(**kwargs)
self.client = ClientSession(connector=connector,
loop=loop,
cookies=cookies,
headers=headers,
skip_auto_headers=skip_auto_headers,
auth=auth,
json_serialize=json_serialize,
request_class=request_class,
response_class=response_class,
ws_response_class=ws_response_class,
version=version,
cookie_jar=cookie_jar,
connector_owner=connector_owner,
raise_for_status=raise_for_status,
read_timeout=read_timeout,
conn_timeout=conn_timeout,
auto_decompress=auto_decompress,
trust_env=trust_env)
def request(self, method, url, *args, **kwargs):
return self.client.request(method=method, url=url, **kwargs)
def get(self, url, *args, **kwargs):
return self.client.get(url=url, **kwargs)
def post(self, url, *args, data=None, **kwargs):
return self.client.post(url=url, data=data, **kwargs)
def put(self, url, *args, data=None, **kwargs):
return self.client.put(url=url, data=data, **kwargs)
def delete(self, url, *args, **kwargs):
return self.client.delete(url=url, **kwargs)
def options(self, url, *args, **kwargs):
return self.client.options(url=url, **kwargs)
def head(self, url, *args, **kwargs):
return self.client.head(url=url, **kwargs)
def patch(self, url, *args, data=None, **kwargs):
return self.client.patch(url=url, data=data, **kwargs)
async def close(self):
await self.client.close()
async def get_response(self, response):
text = await response.text()
return Response(url=response.url,
status=response.status,
charset=response.charset,
content_type=response.content_type,
content_length=response.content_length,
reason=response.reason,
headers=response.headers,
text=text,
selector=etree.HTML(text))
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.close()
class ApiInstance:
def __init__(self,
base_url: str,
timeout: int = 5000,
headers: dict = {},
*,
logger: logging.Logger):
self.base_url = base_url
self.headers = headers
self.timeout = timeout
self.session = ClientSession(timeout=ClientTimeout(total=timeout))
self.logger = logger
async def __aenter__(self) -> "ApiInstance":
return self
async def __aexit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> None:
await self.close()
@staticmethod
def check_result(res: Any) -> Any:
if res == "":
return None
return json_loads_attrs(res)
async def call(self,
method: str,
prefix: str,
data: Any,
headers: dict = {}):
new_headers = CIMultiDict()
new_headers.update(self.headers)
new_headers.update(headers)
data_binary = json.dumps(data).encode("utf-8")
# TODO: this is too much code duplication but I cannot think of
# a way outside macros that could abstract async with block
# and sadly there are no macro in python
if method == "get":
async with self.session.get(self.base_url + "/" + prefix,
headers=new_headers) as resp:
await check_response(resp, self.logger)
res = await resp.text()
return self.check_result(res)
elif method == "post":
async with self.session.post(self.base_url + "/" + prefix,
data=data_binary,
headers=new_headers) as resp:
await check_response(resp, self.logger)
res = await resp.text()
return self.check_result(res)
elif method == "put":
async with self.session.put(self.base_url + "/" + prefix,
data=data_binary,
headers=new_headers) as resp:
await check_response(resp, self.logger)
res = await resp.text()
return self.check_result(res)
elif method == "patch":
async with self.session.patch(self.base_url + "/" + prefix,
data=data_binary,
headers=new_headers) as resp:
await check_response(resp, self.logger)
res = await resp.text()
return self.check_result(res)
elif method == "delete":
async with self.session.delete(self.base_url + "/" + prefix,
headers=new_headers) as resp:
await check_response(resp, self.logger)
res = await resp.text()
return self.check_result(res)
async def make_request(self, method: str, prefix: str, data: Any,
headers: dict):
try:
return await self.call(method, prefix, data, headers)
except (ConflictingEntityException, EntityNotFoundException,
PermissionDeniedException, ProcedureInvocationException,
UnauthorizedException, ValidationException,
BadRequestException, PapieaServerException, ApiException):
raise
except:
self.logger.debug("RENEWING SESSION")
await self.renew_session()
return await self.call(method, prefix, data, headers)
async def post(self, prefix: str, data: Any, headers: dict = {}) -> Any:
return await self.make_request("post", prefix, data, headers)
async def put(self, prefix: str, data: Any, headers: dict = {}) -> Any:
return await self.make_request("put", prefix, data, headers)
async def patch(self, prefix: str, data: Any, headers: dict = {}) -> Any:
return await self.make_request("patch", prefix, data, headers)
async def get(self, prefix: str, headers: dict = {}) -> Any:
return await self.make_request("get", prefix, {}, headers)
async def delete(self, prefix: str, headers: dict = {}) -> Any:
return await self.make_request("delete", prefix, {}, headers)
async def close(self):
await self.session.close()
async def renew_session(self):
await self.close()
self.session = ClientSession(timeout=ClientTimeout(total=self.timeout))
class AsyncUTMTestSession:
"""
Requests Asyncio client session that provides additional functionality for running DSS concurrency tests:
* Adds a prefix to URLs that start with a '/'.
* Automatically applies authorization according to adapter, when present
"""
def __init__(self, prefix_url: str, auth_adapter: Optional[AuthAdapter] = None):
self._client = None
loop = asyncio.get_event_loop()
loop.run_until_complete(self.build_session())
self._prefix_url = prefix_url[0:-1] if prefix_url[-1] == '/' else prefix_url
self.auth_adapter = auth_adapter
self.default_scopes = None
async def build_session(self):
self._client = ClientSession()
def close(self):
loop = asyncio.get_event_loop()
loop.run_until_complete(self._client.close())
def adjust_request_kwargs(self, url, method, kwargs):
if self.auth_adapter:
scopes = None
if 'scopes' in kwargs:
scopes = kwargs['scopes']
del kwargs['scopes']
if 'scope' in kwargs:
scopes = [kwargs['scope']]
del kwargs['scope']
if scopes is None:
scopes = self.default_scopes
if not scopes:
raise ValueError('All tests must specify auth scope for all session requests. Either specify as an argument for each individual HTTP call, or decorate the test with @default_scope.')
headers = {}
for k, v in self.auth_adapter.get_headers(url, scopes).items():
headers[k] = v
kwargs['headers'] = headers
if method == 'PUT' and kwargs.get('data'):
kwargs['json'] = kwargs['data']
del kwargs['data']
return kwargs
async def put(self, url, **kwargs):
url = self._prefix_url + url
if 'auth' not in kwargs:
kwargs = self.adjust_request_kwargs(url, 'PUT', kwargs)
async with self._client.put(url, **kwargs) as response:
return response.status, await response.json()
async def get(self, url, **kwargs):
url = self._prefix_url + url
if 'auth' not in kwargs:
kwargs = self.adjust_request_kwargs(url, 'GET', kwargs)
async with self._client.get(url, **kwargs) as response:
return response.status, await response.json()
async def post(self, url, **kwargs):
url = self._prefix_url + url
if 'auth' not in kwargs:
kwargs = self.adjust_request_kwargs(url, 'POST', kwargs)
async with self._client.post(url, **kwargs) as response:
return response.status, await response.json()
async def delete(self, url, **kwargs):
url = self._prefix_url + url
if 'auth' not in kwargs:
kwargs = self.adjust_request_kwargs(url, 'DELETE', kwargs)
async with self._client.delete(url, **kwargs) as response:
return response.status, await response.json()
