def lambda_handler()

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()
        )