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")
def __init__(self, broker='ws://localhost:8182/'):
self.broker = broker
self.loop = asyncio.get_event_loop()
self.gc = GremlinClient(url=broker, loop=self.loop)
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
def __init__(self):
self.loop = asyncio.get_event_loop()
self.gc = GremlinClient(loop=self.loop)
def setUp(self):
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(None)
self.gc = GremlinClient(url="ws://localhost:8182/", loop=self.loop)
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())
