async def request_events(url: str,
resource: str,
namespace: str = 'default') -> None:
"""Get k8s events for some resource.
:param url: The location of the server
:param resource: The name of the resource you are asking about.
:param namespace: The namespace of the resource you are asking about.
"""
async with websockets.connect(url) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'EVENTS',
APIField.REQUEST: {
'resource': resource,
'namespace': namespace
}
}))
resp = json.loads(await websocket.recv())
if resp[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(resp)
return
if resp[APIField.STATUS] == APIStatus.OK:
rows = [Event(**r) for r in resp[APIField.RESPONSE]]
print_table(rows)
async def request_logs(ws: str,
resource: str,
namespace: str,
container: str,
follow: bool,
lines: Optional[int] = None) -> None:
"""Make the websocket request to get the logs.
:param ws: The websocket host
:param resource: The thing to get logs from
:param namespace: The namespace the pod is in
:param container: The container to get the logs from, only needed
when there are multiple containers in a pod
:param follow: Should you get all the logs available right now or stream
them as they come in?
"""
work = {'resource': resource, 'namespace': namespace, 'follow': follow}
if container is not None:
work['container'] = container
if lines is not None:
work['lines'] = lines
async with websockets.connect(ws) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'LOGS',
APIField.REQUEST: work
}))
line = json.loads(await websocket.recv())
while line[APIField.STATUS] != APIStatus.END:
if line[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(line)
break
LOGGER.info(line[APIField.RESPONSE])
line = json.loads(await websocket.recv())
async def request_status(ws: str,
work: str,
columns: Set[str],
all_cols: bool = False) -> None:
"""Request the status of an odin job over web-sockets
:param ws: The web socket
:param work: the job name
:param columns: A set of columns to include in the output
:param all_cols: Should we just show all columns, If true then columns in ignored
"""
async with websockets.connect(ws) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'STATUS',
APIField.REQUEST: work
}))
results = json.loads(await websocket.recv())
if results[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(results)
return
if results[APIField.STATUS] == APIStatus.OK:
results = results[APIField.RESPONSE]
for result in results:
rows = [Row(**r) for r in result['task_statuses']]
show_status(Pipeline(**result['pipeline_status']), rows,
columns, all_cols)
async def ping(uri: str, message: str) -> None:
"""Ping odin at uri and send message.
:param uri: The location of the server
:param message: The message you expect to see back
:raises RuntimeError: If the server returns an error
"""
async with websockets.connect(uri) as websocket:
await websocket.send(json.dumps({APIField.COMMAND: 'PING', APIField.REQUEST: message}))
resp = json.loads(await websocket.recv())
if resp[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(resp)
raise RuntimeError(resp)
LOGGER.info(resp[APIField.RESPONSE])
async def request_cleanup(ws: str, work: str, purge_db: bool = False, purge_fs: bool = False):
"""Request the work is cleaned up by the server."""
async with websockets.connect(ws) as websocket:
args = {'work': work, 'purge_db': purge_db, 'purge_fs': purge_fs}
await websocket.send(json.dumps({APIField.COMMAND: 'CLEANUP', APIField.REQUEST: args}))
results = json.loads(await websocket.recv())
if results[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(results)
return
if results[APIField.STATUS] == APIStatus.OK:
cleaned = results[APIField.RESPONSE]
print("Results of this request:")
print_table([Cleaned(**c) for c in cleaned])
async def request_generate_config(ws, config):
"""Use async to open a connection to serve.py and generate a config."""
async with websockets.connect(ws) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'GENERATE',
APIField.REQUEST: config
}))
result = json.loads(await websocket.recv())
if result[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(result)
return
if result[APIField.STATUS] == APIStatus.OK:
LOGGER.info('Generated pipeline is called %s',
result[APIField.RESPONSE])
def create_user_http(url: str, jwt_token: str, username: str, password: str,
firstname: str, lastname: str) -> None:
"""Create or update a user over HTTP
:param url: the base URL
:param jwt_token: The JWT token representing this authentication
:param username: The user ID
:param password: The updated password
:param firstname: The firstname
:param lastname: The lastname
"""
user = {"username": username, "password": password}
if firstname:
user['firstname'] = firstname
if lastname:
user['lastname'] = lastname
headers = {'Authorization': f'Bearer {jwt_token}'}
try:
response = requests.get(f'{url}/v1/users/{username}')
if response.status_code == 401:
raise ValueError("Invalid login")
if response.status_code != 200:
# No such user exists so do a POST
response = requests.post(f'{url}/v1/users',
headers=headers,
json={"user": user})
if response.status_code != 200:
raise Exception(f"Failed to create user: {username}")
results = response.json()
LOGGER.info("Created new user")
LOGGER.info(json.dumps(results))
return
results = response.json()
LOGGER.info("Found existing user")
LOGGER.info(json.dumps(results))
except Exception as ex:
LOGGER.error(ex)
return
response = requests.put(f'{url}/v1/users/{username}',
json=user,
headers=headers)
results = response.json()
LOGGER.info(json.dumps(results))
async def request_pipeline_definitions(ws: str, pipeline: str) -> None:
"""Use async to open a connection to serve.py and get a pipeline definition."""
async with websockets.connect(ws) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'SHOW',
APIField.REQUEST: pipeline
}))
result = json.loads(await websocket.recv())
if result[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(result)
return
if result[APIField.STATUS] == APIStatus.OK:
for file_name, file_contents in result[APIField.RESPONSE].items():
LOGGER.info(file_name)
LOGGER.info("=" * 100)
LOGGER.info(file_contents)
LOGGER.info("")
def _authenticate(url, username, passwd):
response = None
url = f'{url}/v1/auth'
try:
response = requests.post(url, data={'username': username, 'password': passwd})
results = response.json()
return results['message']
except Exception as ex:
try:
response = requests.post(url, json={'username': username, 'password': passwd})
results = response.json()
return results['message']
except Exception as ex:
LOGGER.error(url)
if response:
LOGGER.error(response.status_code)
raise ex
async def request_data(url: str, resource: str) -> None:
"""Get k8s data for some resource.
:param url: The location of the server
:param resource: The name of the resource you are asking about.
:param namespace: The namespace of the resource you are asking about.
"""
async with websockets.connect(url) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'DATA',
APIField.REQUEST: {
'resource': resource
}
}))
resp = json.loads(await websocket.recv())
if resp[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(resp)
return
if resp[APIField.STATUS] == APIStatus.OK:
print(json.dumps(resp[APIField.RESPONSE]))
async def schedule_pipeline(ws, work) -> None:
"""Use async to open a connection to serve.py and launch work
Blocks until the job completes (and websocket stays open)
"""
async with websockets.connect(ws) as websocket:
await websocket.send(
json.dumps({
APIField.COMMAND: 'START',
APIField.REQUEST: work
}))
result = json.loads(await websocket.recv())
while result[APIField.STATUS] != APIStatus.END:
if result[APIField.STATUS] == APIStatus.ERROR:
LOGGER.error(result)
return
if result[APIField.RESPONSE].startswith('PIPE_ID'):
pipe_id = result.split(' ')[-1]
LOGGER.info('Started %s', pipe_id)
else:
LOGGER.info(result[APIField.RESPONSE])
result = json.loads(await websocket.recv())