def generic_subscribe(client, handle_channel_data, channels,
query=None, delivery=None):
logger.info(
'Subscribing to %s, press C-c to stop',
channels[0] if len(channels) == 1 else '{0} channels'.format(
len(channels)))
class SubscriptionObserver(object):
def on_enter_subscribed(self):
logger.info('Subscription became active')
def on_enter_failed(self, reason):
logger.error('Subscription failed because:\n%s', reason)
sys.exit(1)
def on_subscription_data(self, data):
handle_channel_data(data)
so = SubscriptionObserver()
delivery = delivery or SubscriptionMode.ADVANCED
for channel in channels:
args = {}
if query:
args.update({'filter': query})
client.subscribe(channel, delivery, so, args)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
sys.exit(0)
def on_subscription_data(data):
data['timestamp'] = time.time()
output = json.dumps(data)
print(output, file=output_stream)
sys.stdout.flush()
if count_limit['count_limit'] is not None:
count_limit['count_limit'] -= len(data['messages'])
if count_limit['count_limit'] <= 0:
logger.info('Message count limit reached')
stop_main_thread()
if size_limit['size_limit'] is not None:
size_limit['size_limit'] -= len(output)
if size_limit['size_limit'] <= 0:
logger.info('Log size limit reached')
stop_main_thread()
def publish(client, channel, enable_acks):
print('Sending input to {0}, press C-d or C-c to stop'.format(channel))
try:
counter = Counter()
if enable_acks:
def callback(reply):
if reply['action'] != 'rtm/publish/ok':
print('Publish failed: ', file=sys.stderr)
sys.exit(1)
counter.decrement()
else:
callback = None
while True:
line = sys.stdin.readline()
if not line:
break
line = line.rstrip()
try:
message = json.loads(line)
except:
message = line
if enable_acks:
counter.increment()
client.publish(channel, message, callback=callback)
except KeyboardInterrupt:
pass
if not enable_acks:
return
if counter.value() > 0:
sleep = 0.1
while sleep < 1.0:
time.sleep(sleep)
sleep += sleep
if not counter.value():
break
else:
logger.info('%s publishes remain unacked', counter.value())
def stop_recording():
logger.info('Time limit reached')
stop_main_thread()
def on_enter_subscribed(self):
logger.info('Subscription became active')
def replay(client, override_channel=None, fast=False, input_file=None, enable_acks=True):
try:
counter = Counter()
if enable_acks:
def callback(reply):
if reply['action'] != 'rtm/publish/ok':
print('Publish failed: ', file=sys.stderr)
sys.exit(1)
counter.decrement()
else:
callback = None
first_message_send_date = None
first_message_recv_date = None
if input_file:
input_stream = open(input_file)
else:
input_stream = sys.stdin
while True:
line = input_stream.readline()
if not line:
break
line = line.rstrip()
try:
data = json.loads(line)
current_message_recv_date = data['timestamp']
channel = override_channel or data['subscription_id']
messages = data['messages']
now = time.time()
if first_message_send_date is not None and not fast:
sleep_amount = (first_message_send_date
+ (current_message_recv_date - first_message_recv_date)
- now)
else:
sleep_amount = 0
if sleep_amount > 0:
time.sleep(sleep_amount)
if first_message_send_date is None:
first_message_send_date = time.time()
first_message_recv_date = current_message_recv_date
for message in messages:
try:
if enable_acks:
counter.increment()
client.publish(channel, message, callback=callback)
except queue.Full as e:
logger.error('Publish queue is full')
except ValueError:
logger.error('Bad line: %s', line)
except Exception as e:
logger.error('Exception: %s', e)
sys.exit(2)
except KeyboardInterrupt:
pass
if not enable_acks:
return
if counter.value() > 0:
sleep = 0.1
while sleep < 3.0:
time.sleep(sleep)
sleep += sleep
if not counter.value():
break
else:
logger.info('%s publishes remain unacked', counter.value())