class Runner:
def __init__(self):
self.loop = asyncio.get_event_loop()
self.gc = GremlinClient(loop=self.loop)
def fetch(self, query):
return self.loop.run_until_complete(self.gc.execute(query))
def fetch_data(self, query, bindings = {}):
result = self.loop.run_until_complete(self.gc.execute(query, bindings=bindings))
data = []
for r in result:
if r.data:
data = data + r.data
return data
def close(self):
self.loop.run_until_complete(self.gc.close())
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
class GremlinQueryRunner:
def __init__(self):
self.loop = asyncio.get_event_loop()
self.gc = GremlinClient(loop=self.loop)
def fetch(self, query):
return self.loop.run_until_complete(self.gc.execute(query))
def close(self):
self.loop.run_until_complete(self.gc.close())
#! /usr/bin/env python3
import sys
import asyncio
from aiogremlin import GremlinClient, GremlinServerError
QSuccess = "graph.addVertex('Entity')" # should insert
QFail = "graph.addVertex('badLabel')" # should fail to insert
if __name__ == '__main__':
loop = asyncio.get_event_loop()
gc = GremlinClient(loop=loop)
try:
e = gc.execute(QSuccess)
r = loop.run_until_complete(e)
print("Successfully inserted vertex of type Entity as expected")
except:
print(
"Should be able to insert vertex of type 'Entity', schema creation failed: ",
sys.exc_info()[0])
try:
e = gc.execute(QFail)
r = loop.run_until_complete(e)
print(
"Schema creation failed, successful insertion of vertex 'badLabel'"
)
except GremlinServerError:
print("Could not insert vertex of type 'badLabel' as expected")
class TitanClient:
def __init__(self, broker='ws://localhost:8182/'):
self.broker = broker
self.loop = asyncio.get_event_loop()
self.gc = GremlinClient(url=broker, loop=self.loop)
def close(self):
self.loop.run_until_complete(self.gc.close())
self.loop.close()
def execute(self, gremlin_query_str, bindings={}):
'''Run a single gremlin query and collect complete
results, optionally with parameters'''
@asyncio.coroutine
def stream(gc):
result_data = []
try:
resp = yield from gc.submit(gremlin_query_str,
bindings=bindings)
except aiohttp.errors.ClientOSError as e:
sys.stdout.write("Cannot connect to " + self.broker + "\n")
sys.exit()
else:
while True:
result = yield from resp.stream.read()
if result is None:
break
assert result.status_code in [206, 200,
204], result.status_code
if not result.status_code == 204:
result_data += result.data
return result_data
result = self.loop.run_until_complete(stream(self.gc))
return result
def execute_many(self, numprocs, queries):
'''
Execute many queries simultaneously, with up to numprocs
concurrent requests.
'''
sem = asyncio.Semaphore(numprocs)
@asyncio.coroutine
def fetch(name, query):
with (yield from sem):
result = yield from self.gc.execute(query)
return (name, query, result)
jobs = [fetch(name, query) for (name, query) in queries]
results = self.loop.run_until_complete(asyncio.gather(*jobs))
return results
def execute_many_dbg(self, numprocs, queries):
'''
Debugging version of execute_many. Prints start, end and timing
information for each job.
'''
sem = asyncio.Semaphore(numprocs)
@asyncio.coroutine
def fetch(name, query):
with (yield from sem):
print("Starting: %s" % name)
t1 = time.time()
result = yield from self.gc.execute(query)
t2 = time.time()
print("Finished: %s in %fs" % (name, t2 - t1))
return (name, query, result)
jobs = [fetch(name, query) for (name, query) in queries]
results = self.loop.run_until_complete(asyncio.gather(*jobs))
return results
def execute_many_params(self, numprocs, query, params):
'''
Execute many variants of the same query simultaneously,
with up to numprocs concurrent requests, using the parameter bindings
given in the list params.
'''
sem = asyncio.Semaphore(numprocs)
@asyncio.coroutine
def fetch(bindings):
with (yield from sem):
result = yield from self.gc.execute(query, bindings=bindings)
return (bindings, result)
jobs = [fetch(bindings) for (bindings) in params]
results = self.loop.run_until_complete(asyncio.gather(*jobs))
return results
def execute_many_params_dbg(self, numprocs, query, params):
'''
Debugging version of execute_many_params. Prints job start and
end messages with times.
'''
sem = asyncio.Semaphore(numprocs)
@asyncio.coroutine
def fetch(bindings):
with (yield from sem):
print("Starting: %a" % (bindings))
t1 = time.time()
result = yield from self.gc.execute(query, bindings=bindings)
t2 = time.time()
print("Finished: %a in %fs" % (bindings, t2 - t1))
return (bindings, result)
print("For query: %s" % (query))
jobs = [fetch(bindings) for (bindings) in params]
results = self.loop.run_until_complete(asyncio.gather(*jobs))
return results
def drop_db(self):
r = self.execute('g.V().drop().iterate()')
count = self.execute('g.V().count()')[0]
assert count is 0
class GremlinClientTest(unittest.TestCase):
def setUp(self):
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(None)
self.gc = GremlinClient(url="ws://localhost:8182/", loop=self.loop)
def tearDown(self):
self.loop.run_until_complete(self.gc.close())
self.loop.close()
def test_connection(self):
@asyncio.coroutine
def go():
ws = yield from self.gc._connector.ws_connect(self.gc.url)
self.assertFalse(ws.closed)
yield from ws.close()
self.loop.run_until_complete(go())
def test_execute(self):
@asyncio.coroutine
def go():
resp = yield from self.gc.execute("x + x", bindings={"x": 4})
return resp
results = self.loop.run_until_complete(go())
self.assertEqual(results[0].data[0], 8)
def test_sub_waitfor(self):
sub1 = self.gc.execute("x + x", bindings={"x": 1})
sub2 = self.gc.execute("x + x", bindings={"x": 2})
sub3 = self.gc.execute("x + x", bindings={"x": 4})
coro = asyncio.gather(*[asyncio.async(sub1, loop=self.loop),
asyncio.async(sub2, loop=self.loop),
asyncio.async(sub3, loop=self.loop)],
loop=self.loop)
# Here I am looking for resource warnings.
results = self.loop.run_until_complete(coro)
self.assertIsNotNone(results)
def test_resp_stream(self):
@asyncio.coroutine
def stream_coro():
results = []
resp = yield from self.gc.submit("x + x", bindings={"x": 4})
while True:
f = yield from resp.stream.read()
if f is None:
break
results.append(f)
self.assertEqual(results[0].data[0], 8)
self.loop.run_until_complete(stream_coro())
def test_execute_error(self):
execute = self.gc.execute("x + x g.asdfas", bindings={"x": 4})
try:
self.loop.run_until_complete(execute)
error = False
except:
error = True
self.assertTrue(error)
def test_rebinding(self):
execute = self.gc.execute("graph2.addVertex()")
try:
self.loop.run_until_complete(execute)
error = False
except:
error = True
self.assertTrue(error)
@asyncio.coroutine
def go():
result = yield from self.gc.execute(
"graph2.addVertex()", rebindings={"graph2": "graph"})
self.assertEqual(len(result), 1)
self.loop.run_until_complete(go())
bad_ls_expected_result = 70960
# Q1="g.V().has('type','file').count()"
# Q2="g.V().has('type','netflow').count()"
# Q3="g.V().has('type','memory').count()"
# Q4="g.V().has('type','resource').count()"
# Q5="g.V().has('type','subject').count()"
# Q6="g.V().has('type','host').count()"
# Q7="g.V().has('type','agent').count()"
# QUERIES = [Q1,Q2,Q3,Q4,Q5,Q6,Q7]
QUERIES = []
if __name__ == '__main__':
chan = KafkaConsumer('pe')
msg = next(chan)
print("Received a signal: ", msg)
loop = asyncio.get_event_loop()
gc = GremlinClient(loop=loop)
for q in QUERIES:
e = gc.execute(q)
r = loop.run_until_complete(e)
print(q, "\n\t", r[0].data[0])
execute = gc.execute(QUERY)
result = loop.run_until_complete(execute)
print("total nodes: ", result[0].data[0])
print("success rate: ", result[0].data[0] / bad_ls_expected_result)
loop.run_until_complete(gc.close())
