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())
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") loop.run_until_complete(gc.close())
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={}): @asyncio.coroutine def stream(gc): result_data = [] resp = yield from gc.submit(gremlin_query_str, bindings=bindings) 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 all_edges(self): return self.execute('g.E()') def all_nodes(self): return self.execute('g.V()') def add_edge(self, n1, d1, n2, d2, d, label): """ Adds edge with label and dict d between nodes with names n1 and n2, if it does not exist. If the nodes do not exist they are created with dicts d1 and d2. """ r1 = self.add_node(n1, d1) r2 = self.add_node(n2, d2) id1 = r1[0]['id'] id2 = r2[0]['id'] properties_str = ', '.join( map( lambda x: '\'{0}\',\'{1}\''.format(x[0], x[1]) if x[0] != 'label' else "", d.items())) r = self.execute( 'g.V({0}).next().addEdge(\'{2}\', g.V({1}).next(), {3})'.format( id1, id2, label, properties_str)) return r def add_node(self, n, d): """ Adds node with name n to the DB if it does not already exist. A node with name n and dictionary d exists in the DB if c1) there is already a node with name n in the db, or c2) there is already a node with dictionary d in the db """ d['ident'] = n c1 = self.execute('g.V().has(\'ident\', \'{}\')'.format(n)) properties_str = ', '.join( map( lambda x: '__().has(\'{0}\',\'{1}\')'.format( x[0], escape(x[1])), d.items())) c2 = self.execute('g.V().and({0})'.format(properties_str)) if not (c1 or c2): properties_str = ', '.join( map(lambda x: '\'{0}\',\'{1}\''.format(x[0], escape(x[1])), d.items())) # Hack: for now, label all new nodes as segment nodes. c1 = self.execute( 'g.addV(label,\'Segment\',{})'.format(properties_str)) assert 'ident' in d, d logger.debug('add_node: Added node with properties {}'.format(d)) else: if c1: logger.debug( 'add_node: Node with name {} already exists'.format(n)) if c2: logger.debug( 'add_node: Node with dict {} already exists'.format(d)) return c1 def load_from_document_graph(self, dg): for n1, n2 in dg.g.edges(): d1 = dg.g.node[n1] d2 = dg.g.node[n2] d = dg.g.edge[n1][n2] label = d['label'] self.add_edge(n1, d1, n2, d2, d, label) def load_segments_from_document_graph(self, dg): """ Given a document graph dg, add edges (and any missing nodes) generated by segmentation """ for n1, n2 in dg.g.edges(): d1 = dg.g.node[n1] d2 = dg.g.node[n2] d = dg.g.edge[n1][n2] label = d['label'] if label.startswith('segment:'): self.add_edge(n1, d1, n2, d2, d, label) def read_into_document_graph(self): doc = Document() node_id2name_map = {} nodes = self.all_nodes() if not nodes: return DocumentGraph(doc) for v in nodes: d = v['properties'] assert 'label' in v, v assert 'ident' in d, d resource_id = d['ident'][0]['value'] logger.info('%9d %s' % (int(v['id']), resource_id)) node_id2name_map[v['id']] = resource_id resource_type = v['label'] att_val_list = [(str(k), str(val[0]['value'])) for (k, val) in d.items() if k not in ['ident', 'label']] r = ResourceFactory.create(resource_type, resource_id, att_val_list) doc.expression_list.append(r) pprint.pprint(node_id2name_map) edges = self.all_edges() for e in edges: event_type = e['label'] ev = EventFactory.create(event_type, node_id2name_map.get(e['outV'], None), node_id2name_map.get(e['inV'], None)) doc.expression_list.append(ev) return DocumentGraph(doc) def drop_db(self): r = self.execute('g.V().drop().iterate()') assert r is None
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())