in gremlin/neptune-streams/lambda/neptune-streams-demo/neptune_workload.py [0:0]
def lambda_handler(event, context):
message_network = MessageNetwork()
highest_vertex_id = message_network.get_highest_vertex_id() if 'highestVertexId' not in event else event['highestVertexId']
batch_size = int(os.environ['BATCH_SIZE'])
index = 0
while os.environ['IS_ACTIVE'] == 'True' and index < batch_size:
if (highest_vertex_id < 1000):
highest_vertex_id += 1
message_network.add_user(MessageNetwork.user_id(highest_vertex_id))
else:
if random.randint(1,4) == 1:
highest_vertex_id += 1
message_network.add_user(MessageNetwork.user_id(highest_vertex_id))
else:
from_id = random.randint(1, highest_vertex_id)
to_id = from_id
while to_id == from_id:
to_id = random.randint(1, highest_vertex_id)
message_network.user_messaged_user(MessageNetwork.user_id(from_id), MessageNetwork.user_id(to_id))
index += 1
message_network.set_highest_vertex_id(highest_vertex_id)
message_network.close()
print('highestVertexId: {}'.format(highest_vertex_id))
if os.environ['IS_ACTIVE'] == 'True':
payload = '{{"highestVertexId": {}}}'.format(highest_vertex_id)
client = boto3.client('lambda')
client.invoke(
FunctionName=context.invoked_function_arn,
InvocationType='Event',
Payload=payload.encode()
)