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