def __enter__(self):
# assign original execute_query()
self.original_execute_query = connection.execute_query
# create a pool with session
self.pool = connection.CONNECTION_POOL_TYPE(pool_size=self.pool_size, with_session=True,
**connection.HOST_PARAMS)
# assign optional binding variables
if self.bindings:
connection.execute_query('g', params=self.bindings, isolate=False, pool=self.pool)
# patch execute_query if we're running non-concurrently
if connection.CONNECTION_TYPE == RexProSyncConnection:
# shadow execute_query with default self.pool
def execute_in_pool(query, params=None, transaction=True, isolate=True,
pool=self.pool, *args, **kwargs):
params = params or {}
return self.original_execute_query(
query, params=params, transaction=transaction, isolate=isolate, pool=pool, *args, **kwargs
)
# patch execute_query to re-use the pool with session
connection.execute_query = execute_in_pool
return self.pool
def __enter__(self):
# assign original execute_query()
self.original_execute_query = connection.execute_query
# create a pool with session
self.pool = connection.CONNECTION_POOL_TYPE(
pool_size=self.pool_size, with_session=True, **connection.HOST_PARAMS
)
# assign optional binding variables
if self.bindings:
connection.execute_query("g", params=self.bindings, isolate=False, pool=self.pool)
# patch execute_query if we're running non-concurrently
if connection.CONNECTION_TYPE == RexProSyncConnection:
# shadow execute_query with default self.pool
def execute_in_pool(query, params=None, transaction=True, isolate=True, pool=self.pool, *args, **kwargs):
params = params or {}
return self.original_execute_query(
query, params=params, transaction=transaction, isolate=isolate, pool=pool, *args, **kwargs
)
# patch execute_query to re-use the pool with session
connection.execute_query = execute_in_pool
return self.pool
def test_blueprints_wrapper(self):
k = 10
wrapper_config = {'class_name': "ReadOnlyGraph", 'bindings': {'k': k}}
with BlueprintsWrapper(**wrapper_config):
gsk = connection.execute_query('"powers of ${k}"')
pysk = "powers of {}".format(k)
self.assertEqual(gsk, pysk)
kk = connection.execute_query("k * k")
self.assertEqual(kk, k * k)
def test_blueprints_wrapper(self):
k = 10
wrapper_config = {
'class_name': "ReadOnlyGraph",
'bindings': {'k': k}
}
with BlueprintsWrapper(**wrapper_config):
gsk = connection.execute_query('"powers of ${k}"')
pysk = "powers of {}".format(k)
self.assertEqual(gsk, pysk)
kk = connection.execute_query("k * k")
self.assertEqual(kk, k * k)
def test_gather_existing_indices(self):
""" Make sure existing indices can be gathered """
v_idx, e_idx = get_existing_indices()
self.assertEqual(len(v_idx), 0)
self.assertEqual(len(e_idx), 0)
# create vertex and edge index
connection.execute_query('g.makeKey(name).dataType(Object.class).indexed(Vertex.class).make(); g.commit()',
params={'name': 'testvertexindex'})
connection.execute_query('g.makeLabel(name).dataType(Object.class).indexed(Edge.class).make(); g.commit()',
params={'name': 'testedgeindex'})
v_idx, e_idx = get_existing_indices()
self.assertEqual(len(v_idx), 1)
self.assertEqual(len(e_idx), 1)
def all(cls, ids, as_dict=False, *args, **kwargs):
"""
Load all edges with the given edge_ids from the graph. By default this will return a list of edges but if
as_dict is True then it will return a dictionary containing edge_ids as keys and edges found as values.
:param ids: A list of titan IDs
:type ids: list
:param as_dict: Toggle whether to return a dictionary or list
:type as_dict: boolean
:rtype: dict | list
"""
if not isinstance(ids, array_types):
raise MogwaiQueryError("ids must be of type list or tuple")
strids = [str(i) for i in ids]
qs = ['ids.collect{g.e(it)}']
results = connection.execute_query('\n'.join(qs), {'ids': strids}, **kwargs)
results = list(filter(None, results))
if len(results) != len(ids):
raise MogwaiQueryError("the number of results don't match the number of edge ids requested")
objects = []
for r in results:
try:
objects += [Element.deserialize(r)]
except KeyError: # pragma: no cover
raise MogwaiQueryError('Edge type "%s" is unknown' % '')
if as_dict: # pragma: no cover
return {e._id: e for e in objects}
return objects
def test_gather_existing_indices(self):
""" Make sure existing indices can be gathered """
v_idx, e_idx = get_existing_indices()
self.assertEqual(len(v_idx), 0)
self.assertEqual(len(e_idx), 0)
# create vertex and edge index
connection.execute_query(
'g.makeKey(name).dataType(Object.class).indexed(Vertex.class).make(); g.commit()',
params={'name': 'testvertexindex'})
connection.execute_query(
'g.makeLabel(name).dataType(Object.class).indexed(Edge.class).make(); g.commit()',
params={'name': 'testedgeindex'})
v_idx, e_idx = get_existing_indices()
self.assertEqual(len(v_idx), 1)
self.assertEqual(len(e_idx), 1)
def index(mc):
counts = mc.get('index_counts')
if not counts:
user_count = User.count()
occupation_count = Occupation.count()
movie_count = Movie.count()
vertex_count = execute_query('g.V.count()')
edge_count = execute_query('g.E.count()')
counts = {'movieCount': movie_count,
'occupationCount': occupation_count,
'userCount': user_count,
'vertexCount': vertex_count,
'edgeCount': edge_count}
mc.set('index_counts', counts)
return counts
return counts
def __enter__(self):
pool = super(BlueprintsWrapper, self).__enter__()
# execute g_assignment
resp = connection.execute_query(self.g_assignment, transaction=False, isolate=False, pool=pool)
# provide a dummy stopTransaction() on non-transactional graphs
connection.execute_query(
"if (!g.metaClass.respondsTo(g, 'stopTransaction')) { g.metaClass.stopTransaction = {null} }",
transaction=False, isolate=False, pool=self.pool
)
# execute wrapper setup
for statement in self.setup:
connection.execute_query(statement, transaction=False, isolate=False, pool=pool)
return pool
def _execute(self, func, deserialize=True, *args, **kwargs):
tmp = "{}.{}()".format(self._get_partial(), func)
self._vars.update({"id": self._vertex._id, "limit": self._limit})
results = connection.execute_query(tmp, self._vars, **kwargs)
if deserialize:
return [Element.deserialize(r) for r in results]
else:
return results
def _execute(self, func, deserialize=True, *args, **kwargs):
tmp = "{}.{}()".format(self._get_partial(), func)
self._vars.update({"id": self._vertex._id, "limit": self._limit})
results = connection.execute_query(tmp, self._vars, **kwargs)
if deserialize:
return [Element.deserialize(r) for r in results]
else:
return results
def all(cls, ids=[], as_dict=False, match_length=True, *args, **kwargs):
"""
Load all vertices with the given ids from the graph. By default this will return a list of vertices but if
as_dict is True then it will return a dictionary containing ids as keys and vertices found as values.
:param ids: A list of titan ids
:type ids: list
:param as_dict: Toggle whether to return a dictionary or list
:type as_dict: boolean
:rtype: dict | list
"""
if not isinstance(ids, array_types):
raise MogwaiQueryError("ids must be of type list or tuple")
if len(ids) == 0:
results = connection.execute_query(
'g.V("element_type","%s").toList()' % cls.get_element_type(),
**kwargs)
else:
strids = [str(i) for i in ids]
results = connection.execute_query('ids.collect{g.v(it)}',
{'ids': strids}, **kwargs)
results = list(filter(None, results))
if len(results) != len(ids) and match_length:
raise MogwaiQueryError(
"the number of results don't match the number of ids requested"
)
objects = []
for r in results:
try:
objects += [Element.deserialize(r)]
except KeyError: # pragma: no cover
raise MogwaiQueryError('Vertex type "%s" is unknown' %
r.get('element_type', ''))
if as_dict: # pragma: no cover
return {v._id: v for v in objects}
return objects
def isolation_test(scope):
wrapper_config = {"bindings": {"scope": scope}, "pool_size": 5}
scope_values = []
with SessionPoolManager(**wrapper_config) as pool:
for i in range(7):
scope_val = connection.execute_query("scope *= 2", isolate=False, pool=pool)
scope_values.append(scope_val)
return scope, scope_values
def __enter__(self):
pool = super(BlueprintsWrapper, self).__enter__()
# execute g_assignment
resp = connection.execute_query(self.g_assignment, transaction=False, isolate=False, pool=pool)
# provide a dummy stopTransaction() on non-transactional graphs
connection.execute_query(
"if (!g.metaClass.respondsTo(g, 'stopTransaction')) { g.metaClass.stopTransaction = {null} }",
transaction=False,
isolate=False,
pool=self.pool,
)
# execute wrapper setup
for statement in self.setup:
connection.execute_query(statement, transaction=False, isolate=False, pool=pool)
return pool
def _simple_traversal(self, operation, *args, **kwargs):
"""
Perform a simple traversal starting from the current edge returning a list of results.
:param operation: The operation to be performed
:type operation: str
:rtype: list
"""
results = connection.execute_query('g.e(id).%s()' % operation, {'id': self.id}, **kwargs)
return [Element.deserialize(r) for r in results]
def isolation_test(scope):
wrapper_config = {'bindings': {'scope': scope}, 'pool_size': 5}
scope_values = []
with SessionPoolManager(**wrapper_config) as pool:
for i in range(7):
scope_val = connection.execute_query("scope *= 2",
isolate=False,
pool=pool)
scope_values.append(scope_val)
return scope, scope_values
def _reload_values(self, *args, **kwargs):
"""
Method for reloading the current vertex by reading its current values from the database.
"""
reloaded_values = {}
results = connection.execute_query('g.v(id)', {'id': self._id}, **kwargs)
#del results['_id']
del results['_type']
reloaded_values['_id'] = results['_id']
for name, value in results.get('_properties', {}).items():
reloaded_values[name] = value
return reloaded_values
def all(cls, ids=[], as_dict=False, match_length=True, *args, **kwargs):
"""
Load all vertices with the given ids from the graph. By default this will return a list of vertices but if
as_dict is True then it will return a dictionary containing ids as keys and vertices found as values.
:param ids: A list of titan ids
:type ids: list
:param as_dict: Toggle whether to return a dictionary or list
:type as_dict: boolean
:rtype: dict | list
"""
if not isinstance(ids, array_types):
raise MogwaiQueryError("ids must be of type list or tuple")
if len(ids) == 0:
results = connection.execute_query('g.V("element_type","%s").toList()' % cls.get_element_type(), **kwargs)
else:
strids = [str(i) for i in ids]
results = connection.execute_query('ids.collect{g.v(it)}', {'ids': strids}, **kwargs)
results = list(filter(None, results))
if len(results) != len(ids) and match_length:
raise MogwaiQueryError("the number of results don't match the number of ids requested")
objects = []
for r in results:
try:
objects += [Element.deserialize(r)]
except KeyError: # pragma: no cover
raise MogwaiQueryError('Vertex type "%s" is unknown' % r.get('element_type', ''))
if as_dict: # pragma: no cover
return {v._id: v for v in objects}
return objects
def _reload_values(self, *args, **kwargs):
"""
Method for reloading the current vertex by reading its current values from the database.
"""
reloaded_values = {}
results = connection.execute_query('g.v(id)', {'id': self._id},
**kwargs)
#del results['_id']
del results['_type']
reloaded_values['_id'] = results['_id']
for name, value in results.get('_properties', {}).items():
reloaded_values[name] = value
return reloaded_values
def _reload_values(self, *args, **kwargs):
""" Re-read the values for this edge from the graph database. """
reloaded_values = {}
results = connection.execute_query('g.e(id)', {'id': self._id}, **kwargs)
if results: # note this won't work if you update a node for titan pre-0.5.x, new id's are created
#del results['_id']
del results['_type']
reloaded_values['_id'] = results['_id']
for name, value in results.get('_properties', {}).items():
reloaded_values[name] = value
if results['_id']:
setattr(self, '_id', results['_id'])
return reloaded_values
else:
return {}
def nested_wrappers(scope):
partition_config = {
'write': "characters",
'bindings': {'scope': scope},
'pool_size': 3
}
with PartitionGraph(**partition_config) as pool:
pile = eventlet.GreenPile()
[pile.spawn(isolation_query, i) for i in range(3)]
scope_val = connection.execute_query('scope', pool=pool)
scope_v = BlueprintsWrapperVertex.create(name=scope_val, pool=pool)
scope_v['nested_values'] = list(pile)
scope_v.save(pool=pool)
return scope, scope_v
def nested_wrappers(scope):
partition_config = {
'write': "characters",
'bindings': {
'scope': scope
},
'pool_size': 3
}
with PartitionGraph(**partition_config) as pool:
pile = eventlet.GreenPile()
[pile.spawn(isolation_query, i) for i in range(3)]
scope_val = connection.execute_query('scope', pool=pool)
scope_v = BlueprintsWrapperVertex.create(name=scope_val, pool=pool)
scope_v['nested_values'] = list(pile)
scope_v.save(pool=pool)
return scope, scope_v
def isolation_query(scope):
wrapper_config = {
'class_name': "ReadOnlyGraph",
'bindings': {'scope': scope},
'pool_size': 5
}
scope_values = []
with BlueprintsWrapper(**wrapper_config) as pool:
for i in range(7):
scope_val = connection.execute_query(
"sleep sleep_length\nreturn scope",
params={'sleep_length': 100 * (i % 2 + 1) - scope*10},
pool=pool
)
scope_values.append(scope_val)
return scope, scope_values
def isolation_query(scope):
wrapper_config = {
'class_name': "ReadOnlyGraph",
'bindings': {
'scope': scope
},
'pool_size': 5
}
scope_values = []
with BlueprintsWrapper(**wrapper_config) as pool:
for i in range(7):
scope_val = connection.execute_query(
"sleep sleep_length\nreturn scope",
params={'sleep_length': 100 * (i % 2 + 1) - scope * 10},
pool=pool)
scope_values.append(scope_val)
return scope, scope_values
from mogwai import connection, properties
from mogwai.models import Vertex, Edge
from mogwai.tools import SessionPoolManager, BlueprintsWrapper, PartitionGraph
from mogwai.exceptions import MogwaiGremlinException
connection.setup('localhost')
##
# Persist a session with SessionPoolManager
##
k = 10
with SessionPoolManager(bindings={'k': k}):
gsk = connection.execute_query('"powers of ${k}"')
pysk = "powers of {}".format(k)
assert gsk == pysk
kk = connection.execute_query("k * k")
assert kk == k * k
##
# Wrap the graph with a Blueprints Implementation
##
class BlueprintsWrapperVertex(Vertex):
element_type = 'blueprints_wrapper_vertex'
name = properties.String(required=True, max_length=128)
def __exit__(self, exc, message, traceback):
# execute g = g.baseGraph
connection.execute_query("g = g.baseGraph", transaction=False, isolate=False, pool=self.pool)
return super(BlueprintsWrapper, self).__exit__(exc, message, traceback)
def get_existing_indices():
""" Find all Vertex and Edge types available in the database """
vertex_indices = connection.execute_query('g.getIndexedKeys(Vertex.class)')
edge_indices = connection.execute_query('g.getIndexedKeys(Edge.class)')
return vertex_indices, edge_indices
def __exit__(self, exc, message, traceback):
# execute g = g.baseGraph
connection.execute_query("g = g.baseGraph", transaction=False, isolate=False, pool=self.pool)
return super(BlueprintsWrapper, self).__exit__(exc, message, traceback)
def test_wrapper_isolation(self):
connection.execute_query("k")
def test_wrapper_isolation(self):
connection.execute_query("k")
def get_existing_indices():
""" Find all Vertex and Edge types available in the database """
vertex_indices = connection.execute_query('g.getIndexedKeys(Vertex.class)')
edge_indices = connection.execute_query('g.getIndexedKeys(Edge.class)')
return vertex_indices, edge_indices
def __call__(self, instance, *args, **kwargs):
"""
Intercept attempts to call the GremlinMethod attribute and perform a gremlin query returning the results.
:param instance: The class instance the method was called on
:param pool: The RexPro connection pool to execute the query with (optional)
:type instance: object
"""
self._setup()
# pop the optional execute query arguments from kwargs
query_kwargs = connection.pop_execute_query_kwargs(kwargs)
query_kwargs['transaction'] = query_kwargs.get('transaction') or self.transaction
args = list(args)
if not self.classmethod:
args = [instance._id] + args
params = self.defaults.copy()
if len(args + list(kwargs.values())) > len(self.arg_list): # pragma: no cover
raise TypeError('%s() takes %s args, %s given' % (self.attr_name, len(self.arg_list), len(args)))
#check for and calculate callable defaults
for k, v in params.items():
if callable(v):
params[k] = v()
arglist = self.arg_list[:]
for arg in args:
params[arglist.pop(0)] = arg
for k, v in kwargs.items():
if k not in arglist:
an = self.attr_name
if k in params: # pragma: no cover
raise TypeError("%s() got multiple values for keyword argument '%s'" % (an, k))
else: # pragma: no cover
raise TypeError("%s() got an unexpected keyword argument '%s'" % (an, k))
arglist.pop(arglist.index(k))
params[k] = v
params = self.transform_params_to_database(params)
import_list = []
for imp in self.imports + self.extra_imports:
if imp is not None:
for import_string in imp.import_list:
import_list.append(import_string)
import_string = '\n'.join(import_list)
script = '\n'.join([import_string, self.function_body])
try:
if hasattr(instance, 'get_element_type'):
context = "vertices.{}".format(instance.get_element_type())
elif hasattr(instance, 'get_label'):
context = "edges.{}".format(instance.get_label())
else:
context = "other"
context = "{}.{}".format(context, self.method_name)
tmp = connection.execute_query(script, params, context=context, **query_kwargs)
except MogwaiQueryError as pqe: # pragma: no cover
import pprint
msg = "Error while executing Gremlin method\n\n"
msg += "[Method]\n{}\n\n".format(self.method_name)
msg += "[Params]\n{}\n\n".format(pprint.pformat(params))
msg += "[Function Body]\n{}\n".format(self.function_body)
msg += "[Imports]\n{}\n".format(import_string)
msg += "\n[Error]\n{}\n".format(pqe)
if hasattr(pqe, 'raw_response'):
msg += "\n[Raw Response]\n{}\n".format(pqe.raw_response)
raise MogwaiGremlinException(msg)
return tmp
from mogwai import connection, properties
from mogwai.models import Vertex, Edge
from mogwai.tools import SessionPoolManager, BlueprintsWrapper, PartitionGraph
from mogwai.exceptions import MogwaiGremlinException
connection.setup('localhost')
##
# Persist a session with SessionPoolManager
##
k = 10
with SessionPoolManager(bindings={'k': k}):
gsk = connection.execute_query('"powers of ${k}"')
pysk = "powers of {}".format(k)
assert gsk == pysk
kk = connection.execute_query("k * k")
assert kk == k * k
##
# Wrap the graph with a Blueprints Implementation
##
class BlueprintsWrapperVertex(Vertex):
element_type = 'blueprints_wrapper_vertex'
name = properties.String(required=True, max_length=128)
