| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import json |
| from warnings import warn |
| |
| from cassandra import ConsistencyLevel |
| from cassandra.query import Statement, SimpleStatement |
| from cassandra.datastax.graph.types import Vertex, Edge, Path, VertexProperty |
| from cassandra.datastax.graph.graphson import GraphSON2Reader, GraphSON3Reader |
| |
| |
| __all__ = [ |
| 'GraphProtocol', 'GraphOptions', 'GraphStatement', 'SimpleGraphStatement', |
| 'single_object_row_factory', 'graph_result_row_factory', 'graph_object_row_factory', |
| 'graph_graphson2_row_factory', 'Result', 'graph_graphson3_row_factory' |
| ] |
| |
| # (attr, description, server option) |
| _graph_options = ( |
| ('graph_name', 'name of the targeted graph.', 'graph-name'), |
| ('graph_source', 'choose the graph traversal source, configured on the server side.', 'graph-source'), |
| ('graph_language', 'the language used in the queries (default "gremlin-groovy")', 'graph-language'), |
| ('graph_protocol', 'the graph protocol that the server should use for query results (default "graphson-1-0")', 'graph-results'), |
| ('graph_read_consistency_level', '''read `cassandra.ConsistencyLevel <http://docs.datastax.com/en/developer/python-driver/latest/api/cassandra/#cassandra.ConsistencyLevel>`_ for graph queries (if distinct from session default). |
| Setting this overrides the native `Statement.consistency_level <http://docs.datastax.com/en/developer/python-driver/latest/api/cassandra/query/#cassandra.query.Statement.consistency_level>`_ for read operations from Cassandra persistence''', 'graph-read-consistency'), |
| ('graph_write_consistency_level', '''write `cassandra.ConsistencyLevel <http://docs.datastax.com/en/developer/python-driver/latest/api/cassandra/#cassandra.ConsistencyLevel>`_ for graph queries (if distinct from session default). |
| Setting this overrides the native `Statement.consistency_level <http://docs.datastax.com/en/developer/python-driver/latest/api/cassandra/query/#cassandra.query.Statement.consistency_level>`_ for write operations to Cassandra persistence.''', 'graph-write-consistency') |
| ) |
| _graph_option_names = tuple(option[0] for option in _graph_options) |
| |
| # this is defined by the execution profile attribute, not in graph options |
| _request_timeout_key = 'request-timeout' |
| |
| |
| class GraphProtocol(object): |
| |
| GRAPHSON_1_0 = b'graphson-1.0' |
| """ |
| GraphSON1 |
| """ |
| |
| GRAPHSON_2_0 = b'graphson-2.0' |
| """ |
| GraphSON2 |
| """ |
| |
| GRAPHSON_3_0 = b'graphson-3.0' |
| """ |
| GraphSON3 |
| """ |
| |
| |
| class GraphOptions(object): |
| """ |
| Options for DSE Graph Query handler. |
| """ |
| # See _graph_options map above for notes on valid options |
| |
| DEFAULT_GRAPH_PROTOCOL = GraphProtocol.GRAPHSON_1_0 |
| DEFAULT_GRAPH_LANGUAGE = b'gremlin-groovy' |
| |
| def __init__(self, **kwargs): |
| self._graph_options = {} |
| kwargs.setdefault('graph_source', 'g') |
| kwargs.setdefault('graph_language', GraphOptions.DEFAULT_GRAPH_LANGUAGE) |
| for attr, value in kwargs.items(): |
| if attr not in _graph_option_names: |
| warn("Unknown keyword argument received for GraphOptions: {0}".format(attr)) |
| setattr(self, attr, value) |
| |
| def copy(self): |
| new_options = GraphOptions() |
| new_options._graph_options = self._graph_options.copy() |
| return new_options |
| |
| def update(self, options): |
| self._graph_options.update(options._graph_options) |
| |
| def get_options_map(self, other_options=None): |
| """ |
| Returns a map for these options updated with other options, |
| and mapped to graph payload types. |
| """ |
| options = self._graph_options.copy() |
| if other_options: |
| options.update(other_options._graph_options) |
| |
| # cls are special-cased so they can be enums in the API, and names in the protocol |
| for cl in ('graph-write-consistency', 'graph-read-consistency'): |
| cl_enum = options.get(cl) |
| if cl_enum is not None: |
| options[cl] = ConsistencyLevel.value_to_name[cl_enum].encode() |
| return options |
| |
| def set_source_default(self): |
| """ |
| Sets ``graph_source`` to the server-defined default traversal source ('default') |
| """ |
| self.graph_source = 'default' |
| |
| def set_source_analytics(self): |
| """ |
| Sets ``graph_source`` to the server-defined analytic traversal source ('a') |
| """ |
| self.graph_source = 'a' |
| |
| def set_source_graph(self): |
| """ |
| Sets ``graph_source`` to the server-defined graph traversal source ('g') |
| """ |
| self.graph_source = 'g' |
| |
| def set_graph_protocol(self, protocol): |
| """ |
| Sets ``graph_protocol`` as server graph results format (See :class:`cassandra.datastax.graph.GraphProtocol`) |
| """ |
| self.graph_protocol = protocol |
| |
| @property |
| def is_default_source(self): |
| return self.graph_source in (b'default', None) |
| |
| @property |
| def is_analytics_source(self): |
| """ |
| True if ``graph_source`` is set to the server-defined analytics traversal source ('a') |
| """ |
| return self.graph_source == b'a' |
| |
| @property |
| def is_graph_source(self): |
| """ |
| True if ``graph_source`` is set to the server-defined graph traversal source ('g') |
| """ |
| return self.graph_source == b'g' |
| |
| |
| for opt in _graph_options: |
| |
| def get(self, key=opt[2]): |
| return self._graph_options.get(key) |
| |
| def set(self, value, key=opt[2]): |
| if value is not None: |
| # normalize text here so it doesn't have to be done every time we get options map |
| if isinstance(value, str): |
| value = value.encode() |
| self._graph_options[key] = value |
| else: |
| self._graph_options.pop(key, None) |
| |
| def delete(self, key=opt[2]): |
| self._graph_options.pop(key, None) |
| |
| setattr(GraphOptions, opt[0], property(get, set, delete, opt[1])) |
| |
| |
| class GraphStatement(Statement): |
| """ An abstract class representing a graph query.""" |
| |
| @property |
| def query(self): |
| raise NotImplementedError() |
| |
| def __str__(self): |
| return u'<GraphStatement query="{0}">'.format(self.query) |
| __repr__ = __str__ |
| |
| |
| class SimpleGraphStatement(GraphStatement, SimpleStatement): |
| """ |
| Simple graph statement for :meth:`.Session.execute_graph`. |
| Takes the same parameters as :class:`.SimpleStatement`. |
| """ |
| @property |
| def query(self): |
| return self._query_string |
| |
| |
| def single_object_row_factory(column_names, rows): |
| """ |
| returns the JSON string value of graph results |
| """ |
| return [row[0] for row in rows] |
| |
| |
| def graph_result_row_factory(column_names, rows): |
| """ |
| Returns a :class:`Result <cassandra.datastax.graph.Result>` object that can load graph results and produce specific types. |
| The Result JSON is deserialized and unpacked from the top-level 'result' dict. |
| """ |
| return [Result(json.loads(row[0])['result']) for row in rows] |
| |
| |
| def graph_object_row_factory(column_names, rows): |
| """ |
| Like :func:`~.graph_result_row_factory`, except known element types (:class:`~.Vertex`, :class:`~.Edge`) are |
| converted to their simplified objects. Some low-level metadata is shed in this conversion. Unknown result types are |
| still returned as :class:`Result <cassandra.datastax.graph.Result>`. |
| """ |
| return _graph_object_sequence(json.loads(row[0])['result'] for row in rows) |
| |
| |
| def _graph_object_sequence(objects): |
| for o in objects: |
| res = Result(o) |
| if isinstance(o, dict): |
| typ = res.value.get('type') |
| if typ == 'vertex': |
| res = res.as_vertex() |
| elif typ == 'edge': |
| res = res.as_edge() |
| yield res |
| |
| |
| class _GraphSONContextRowFactory(object): |
| graphson_reader_class = None |
| graphson_reader_kwargs = None |
| |
| def __init__(self, cluster): |
| context = {'cluster': cluster} |
| kwargs = self.graphson_reader_kwargs or {} |
| self.graphson_reader = self.graphson_reader_class(context, **kwargs) |
| |
| def __call__(self, column_names, rows): |
| return [self.graphson_reader.read(row[0])['result'] for row in rows] |
| |
| |
| class _GraphSON2RowFactory(_GraphSONContextRowFactory): |
| """Row factory to deserialize GraphSON2 results.""" |
| graphson_reader_class = GraphSON2Reader |
| |
| |
| class _GraphSON3RowFactory(_GraphSONContextRowFactory): |
| """Row factory to deserialize GraphSON3 results.""" |
| graphson_reader_class = GraphSON3Reader |
| |
| |
| graph_graphson2_row_factory = _GraphSON2RowFactory |
| graph_graphson3_row_factory = _GraphSON3RowFactory |
| |
| |
| class Result(object): |
| """ |
| Represents deserialized graph results. |
| Property and item getters are provided for convenience. |
| """ |
| |
| value = None |
| """ |
| Deserialized value from the result |
| """ |
| |
| def __init__(self, value): |
| self.value = value |
| |
| def __getattr__(self, attr): |
| if not isinstance(self.value, dict): |
| raise ValueError("Value cannot be accessed as a dict") |
| |
| if attr in self.value: |
| return self.value[attr] |
| |
| raise AttributeError("Result has no top-level attribute %r" % (attr,)) |
| |
| def __getitem__(self, item): |
| if isinstance(self.value, dict) and isinstance(item, str): |
| return self.value[item] |
| elif isinstance(self.value, list) and isinstance(item, int): |
| return self.value[item] |
| else: |
| raise ValueError("Result cannot be indexed by %r" % (item,)) |
| |
| def __str__(self): |
| return str(self.value) |
| |
| def __repr__(self): |
| return "%s(%r)" % (Result.__name__, self.value) |
| |
| def __eq__(self, other): |
| return self.value == other.value |
| |
| def as_vertex(self): |
| """ |
| Return a :class:`Vertex` parsed from this result |
| |
| Raises TypeError if parsing fails (i.e. the result structure is not valid). |
| """ |
| try: |
| return Vertex(self.id, self.label, self.type, self.value.get('properties', {})) |
| except (AttributeError, ValueError, TypeError): |
| raise TypeError("Could not create Vertex from %r" % (self,)) |
| |
| def as_edge(self): |
| """ |
| Return a :class:`Edge` parsed from this result |
| |
| Raises TypeError if parsing fails (i.e. the result structure is not valid). |
| """ |
| try: |
| return Edge(self.id, self.label, self.type, self.value.get('properties', {}), |
| self.inV, self.inVLabel, self.outV, self.outVLabel) |
| except (AttributeError, ValueError, TypeError): |
| raise TypeError("Could not create Edge from %r" % (self,)) |
| |
| def as_path(self): |
| """ |
| Return a :class:`Path` parsed from this result |
| |
| Raises TypeError if parsing fails (i.e. the result structure is not valid). |
| """ |
| try: |
| return Path(self.labels, self.objects) |
| except (AttributeError, ValueError, TypeError): |
| raise TypeError("Could not create Path from %r" % (self,)) |
| |
| def as_vertex_property(self): |
| return VertexProperty(self.value.get('label'), self.value.get('value'), self.value.get('properties', {})) |