def healthCheck(url, role, secret, channel, retry=False, httpCheck=False):
'''Perform a health check'''
credentials = createCredentials(role, secret)
if httpCheck:
healthCheckHttp(url)
healthCheckPubSub(url, credentials, channel, retry)
healthCheckKVStore(url, credentials)
def test_publish(runner):
port = runner.port
url = getDefaultHealthCheckUrl(None, port)
role = getDefaultRoleForApp('health')
secret = getDefaultSecretForApp('health')
creds = createCredentials(role, secret)
connection = Connection(url, creds)
asyncio.get_event_loop().run_until_complete(clientCoroutine(connection))
def admin(endpoint, appkey, rolename, rolesecret, action, connection_id):
'''Execute admin operations on the server
\b
cobra admin --action disconnect --connection_id 3919dc67
'''
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
asyncio.get_event_loop().run_until_complete(
adminCoroutine(url, credentials, action, connection_id))
def test_monitor(runner):
'''Starts a server, then run a health check'''
port = runner.port
url = getDefaultMonitorUrl(None, port)
role = getDefaultRoleForApp('stats')
secret = getDefaultSecretForApp('stats')
creds = createCredentials(role, secret)
connection = Connection(url, creds)
asyncio.get_event_loop().run_until_complete(clientCoroutine(connection))
monitor(connection)
def test_admin_disconnect_one(runner):
'''Starts a server, then run a health check'''
port = runner.port
url = getDefaultHealthCheckUrl(None, port)
role = getDefaultRoleForApp('health')
secret = getDefaultSecretForApp('health')
creds = createCredentials(role, secret)
connection = Connection(url, creds)
connectionToBeClosed = Connection(url, creds)
asyncio.get_event_loop().run_until_complete(
clientCoroutine(connection, connectionToBeClosed))
def monitor(
endpoint,
appkey,
rolename,
rolesecret,
raw,
role_filter,
channel_filter,
metric_filter,
hide_nodes,
hide_roles,
hide_channels,
hide_summary,
subscribers,
system,
once,
tidy,
unsafe,
):
'''Monitor cobra
'''
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
if tidy:
hide_nodes = True
hide_roles = True
hide_channels = True
retry = not unsafe
runMonitor(
url,
credentials,
raw,
role_filter,
channel_filter,
metric_filter,
not hide_nodes,
not hide_roles,
not hide_channels,
not hide_summary,
subscribers,
system,
once,
retry,
)
def test_monitor_redis_down(redisDownRunner):
'''Starts a server, then run a health check'''
port = redisDownRunner.port
url = getDefaultMonitorUrl(None, port)
role = getDefaultRoleForApp('stats')
secret = getDefaultSecretForApp('stats')
creds = createCredentials(role, secret)
connection = Connection(url, creds)
# FIXME: bring this back
asyncio.get_event_loop().run_until_complete(
clientCoroutineRedisDown(connection))
with pytest.raises(ActionException):
monitor(connection)
def test_save_position(runner):
'''Starts a server, then run a health check'''
port = runner.port
url = getDefaultHealthCheckUrl(None, port)
role = getDefaultRoleForApp('health')
secret = getDefaultSecretForApp('health')
creds = createCredentials(role, secret)
connection = Connection(url, creds)
_ = Connection(url, creds)
uniqueId = uuid.uuid4().hex[:8]
channel = 'test_save_position_channel::' + uniqueId
resumeFromLastPositionId = 'last_position_id::' + uniqueId
asyncio.get_event_loop().run_until_complete(
clientCoroutine(connection, channel, url, creds,
resumeFromLastPositionId))
def read(endpoint, appkey, rolename, rolesecret, channel, position):
'''Read from the cobra key value store
'''
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
async def handler(url, credentials, channel, position):
connection = Connection(url, credentials)
await connection.connect()
try:
data = await connection.read(channel, position)
except ActionException as e:
logging.error(f'Action error: {e}')
return
await connection.close()
print()
print(f'handler received message {data}')
asyncio.get_event_loop().run_until_complete(
handler(url, credentials, channel, position))
def subscribe(
endpoint,
appkey,
rolename,
rolesecret,
channel,
position,
stream_sql,
resume_from_last_position,
batch_size,
disable_debug_memory,
):
'''Subscribe to a channel
'''
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
resumeFromLastPositionId = ''
if resume_from_last_position:
resumeFromLastPositionId = f'{channel}::{stream_sql}'
asyncio.get_event_loop().run_until_complete(
subscribeClient(
url,
credentials,
channel,
position,
stream_sql,
MessageHandlerClass,
{
'resume_from_last_position': resume_from_last_position,
'disable_debug_memory': disable_debug_memory,
},
resumeFromLastPosition=resume_from_last_position,
resumeFromLastPositionId=resumeFromLastPositionId,
batchSize=batch_size,
)
)
def publish(
endpoint,
appkey,
channel,
path,
rolename,
rolesecret,
batch,
batch_events_path,
limit,
repeat,
delay,
summary,
):
'''Publish to a channel
'''
if batch:
path = batch_events_path
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
run(url, channel, path, credentials, repeat, delay, limit, summary)
def write(endpoint, appkey, rolename, rolesecret, channel, data, repeat):
'''Write to the cobra key value store
'''
url = makeUrl(endpoint, appkey)
credentials = createCredentials(rolename, rolesecret)
async def handler(url, credentials, channel, data, repeat):
connection = Connection(url, credentials)
await connection.connect()
try:
while True:
await connection.write(channel, data)
if not repeat:
break
except ActionException as e:
logging.error(f'Action error: {e}')
return
await connection.close()
asyncio.get_event_loop().run_until_complete(
handler(url, credentials, channel, data, repeat))
def main(url, role, secret, channel, position, username, password):
credentials = createCredentials(role, secret)
asyncio.get_event_loop().run_until_complete(
runSubscriber(url, credentials, channel, position, password))
async def runClient(
url,
role,
secret,
channel,
position,
stream_sql,
verbose,
username,
password,
loop,
inputs,
stop,
):
credentials = createCredentials(role, secret)
q: asyncio.Queue[str] = asyncio.Queue(loop=loop)
args = {'verbose': verbose, 'queue': q}
task = asyncio.ensure_future(
subscribeClient(url, credentials, channel, position, stream_sql,
MessageHandlerClass, args))
addTaskCleanup(task)
try:
while True:
incoming: asyncio.Future[Any] = asyncio.ensure_future(q.get())
outgoing: asyncio.Future[Any] = asyncio.ensure_future(inputs.get())
done: Set[asyncio.Future[Any]]
pending: Set[asyncio.Future[Any]]
done, pending = await asyncio.wait(
[incoming, outgoing, stop],
return_when=asyncio.FIRST_COMPLETED)
# Cancel pending tasks to avoid leaking them.
if incoming in pending:
incoming.cancel()
if outgoing in pending:
outgoing.cancel()
if incoming in done:
try:
(message, position) = incoming.result()
except websockets.exceptions.ConnectionClosed:
break
else:
data = message.get('data', {})
user = data.get('user', 'unknown user')
text = data.get('text', '<invalid message>')
encrypted = data.get('encrypted', False)
if encrypted:
text = decrypt(text, password)
messageId = message.get('id')
# Use redis position to get a datetime
timestamp = position.split('-')[0]
dt = datetime.datetime.fromtimestamp(int(timestamp) / 1000)
dtFormatted = dt.strftime('[%H:%M:%S]')
maxUserNameLength = 12
padding = (maxUserNameLength - len(user)) * ' '
user = colorize(user)
print_during_input(
f'{dtFormatted} {padding} {user}: {text}')
if outgoing in done:
text = outgoing.result()
messageId = uuid.uuid4().hex # FIXME needed ?
encrypted = False
if password is not None:
text = encrypt(text, password)
encrypted = True
message = {
'data': {
'encrypted': encrypted,
'user': username,
'text': text
},
'id': messageId,
}
await args['connection'].publish(channel, message)
if stop in done:
break
except Exception as e:
logging.error(f'Caught exception: {e}')
finally:
connection = args.get('connection')
if connection is not None:
closeStatus = await args['connection'].close()
print_over_input(f"Connection closed: {closeStatus}.")
task.cancel()
await task
exit_from_event_loop_thread(loop, stop)
