gremlinpython 3.6.8__py3-none-any.whl

gremlin_python/driver/driver_remote_connection.py

@@ -0,0 +1,183 @@
+#
+# 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 logging
+from concurrent.futures import Future
+import warnings
+
+from gremlin_python.driver import client, serializer
+from gremlin_python.driver.remote_connection import (
+    RemoteConnection, RemoteTraversal)
+from gremlin_python.process.strategies import OptionsStrategy
+from gremlin_python.process.traversal import Bytecode
+import uuid
+
+log = logging.getLogger("gremlinpython")
+
+__author__ = 'David M. Brown (davebshow@gmail.com), Lyndon Bauto (lyndonb@bitquilltech.com)'
+
+
+class DriverRemoteConnection(RemoteConnection):
+
+    def __init__(self, url, traversal_source="g", protocol_factory=None,
+                 transport_factory=None, pool_size=None, max_workers=None,
+                 username="", password="", kerberized_service='',
+                 message_serializer=None, graphson_reader=None,
+                 graphson_writer=None, headers=None, session=None,
+                 enable_user_agent_on_connect=True, enable_compression=False, **transport_kwargs):
+        log.info("Creating DriverRemoteConnection with url '%s'", str(url))
+        self.__url = url
+        self.__traversal_source = traversal_source
+        self.__protocol_factory = protocol_factory
+        self.__transport_factory = transport_factory
+        self.__pool_size = pool_size
+        self.__max_workers = max_workers
+        self.__username = username
+        self.__password = password
+        self.__kerberized_service = kerberized_service
+        self.__message_serializer = message_serializer
+        self.__graphson_reader = graphson_reader
+        self.__graphson_writer = graphson_writer
+        self.__headers = headers
+        self.__session = session
+        self.__enable_user_agent_on_connect = enable_user_agent_on_connect
+        self.__enable_compression = enable_compression
+        self.__transport_kwargs = transport_kwargs
+
+        # keeps a list of sessions that have been spawned from this DriverRemoteConnection
+        # so that they can be closed if this parent session is closed.
+        self.__spawned_sessions = []
+
+        if message_serializer is None and graphson_reader is not None and graphson_writer is not None:
+            message_serializer = serializer.GraphSONMessageSerializer(
+                reader=graphson_reader,
+                writer=graphson_writer)
+        self._client = client.Client(url, traversal_source,
+                                     protocol_factory=protocol_factory,
+                                     transport_factory=transport_factory,
+                                     pool_size=pool_size,
+                                     max_workers=max_workers,
+                                     message_serializer=message_serializer,
+                                     username=username,
+                                     password=password,
+                                     kerberized_service=kerberized_service,
+                                     headers=headers,
+                                     session=session,
+                                     enable_user_agent_on_connect=enable_user_agent_on_connect,
+                                     enable_compression=enable_compression,
+                                     **transport_kwargs)
+        self._url = self._client._url
+        self._traversal_source = self._client._traversal_source
+
+    def close(self):
+        # close this client and any DriverRemoteConnection instances spawned from this one
+        # for a session
+        if len(self.__spawned_sessions) > 0:
+            log.info("closing spawned sessions from DriverRemoteConnection with url '%s'", str(self._url))
+            for spawned_session in self.__spawned_sessions:
+                spawned_session.close()
+            self.__spawned_sessions.clear()
+
+        if self.__session:
+            log.info("closing DriverRemoteConnection with url '%s' with session '%s'",
+                         str(self._url), str(self.__session))
+        else:
+            log.info("closing DriverRemoteConnection with url '%s'", str(self._url))
+
+        self._client.close()
+
+    def submit(self, bytecode):
+        log.debug("submit with bytecode '%s'", str(bytecode))
+        result_set = self._client.submit(bytecode, request_options=self._extract_request_options(bytecode))
+        results = result_set.all().result()
+        return RemoteTraversal(iter(results))
+
+    def submitAsync(self, message, bindings=None, request_options=None):
+        warnings.warn(
+            "gremlin_python.driver.driver_remote_connection.DriverRemoteConnection.submitAsync will be replaced by "
+            "gremlin_python.driver.driver_remote_connection.DriverRemoteConnection.submit_async.",
+            DeprecationWarning)
+        self.submit_async(message, bindings, request_options)
+
+    def submit_async(self, bytecode):
+        log.debug("submit_async with bytecode '%s'", str(bytecode))
+        future = Future()
+        future_result_set = self._client.submit_async(bytecode, request_options=self._extract_request_options(bytecode))
+
+        def cb(f):
+            try:
+                result_set = f.result()
+                results = result_set.all().result()
+                future.set_result(RemoteTraversal(iter(results)))
+            except Exception as e:
+                future.set_exception(e)
+
+        future_result_set.add_done_callback(cb)
+        return future
+
+    def is_closed(self):
+        return self._client.is_closed()
+
+    def is_session_bound(self):
+        return self.__session is not None
+
+    def create_session(self):
+        log.info("Creating session based connection")
+        if self.is_session_bound():
+            raise Exception('Connection is already bound to a session - child sessions are not allowed')
+        conn = DriverRemoteConnection(self.__url,
+                                      traversal_source=self.__traversal_source,
+                                      protocol_factory=self.__protocol_factory,
+                                      transport_factory=self.__transport_factory,
+                                      pool_size=self.__pool_size,
+                                      max_workers=self.__max_workers,
+                                      username=self.__username,
+                                      password=self.__password,
+                                      kerberized_service=self.__kerberized_service,
+                                      message_serializer=self.__message_serializer,
+                                      graphson_reader=self.__graphson_reader,
+                                      graphson_writer=self.__graphson_writer,
+                                      headers=self.__headers,
+                                      session=uuid.uuid4(),
+                                      enable_user_agent_on_connect=self.__enable_user_agent_on_connect,
+                                      **self.__transport_kwargs)
+        self.__spawned_sessions.append(conn)
+        return conn
+
+    def remove_session(self, session_based_connection):
+        session_based_connection.close()
+        self.__spawned_sessions.remove(session_based_connection)
+
+    def commit(self):
+        log.info("Submitting commit graph operation.")
+        return self._client.submit(Bytecode.GraphOp.commit())
+
+    def rollback(self):
+        log.info("Submitting rollback graph operation.")
+        return self._client.submit(Bytecode.GraphOp.rollback())
+
+    @staticmethod
+    def _extract_request_options(bytecode):
+        options_strategy = next((x for x in bytecode.source_instructions
+                                 if x[0] == "withStrategies" and type(x[1]) is OptionsStrategy), None)
+        request_options = None
+        if options_strategy:
+            allowed_keys = ['evaluationTimeout', 'scriptEvaluationTimeout', 'batchSize', 'requestId', 'userAgent']
+            request_options = {allowed: options_strategy[1].configuration[allowed] for allowed in allowed_keys
+                               if allowed in options_strategy[1].configuration}
+        return request_options

gremlin_python/driver/protocol.py

@@ -0,0 +1,259 @@
+#
+# 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
+import logging
+import abc
+import base64
+import struct
+
+# import kerberos    Optional dependency imported in relevant codeblock
+
+from gremlin_python.driver import request
+from gremlin_python.driver.resultset import ResultSet
+from gremlin_python.process.translator import Translator
+from gremlin_python.process.traversal import Bytecode
+
+log = logging.getLogger("gremlinpython")
+
+__author__ = 'David M. Brown (davebshow@gmail.com)'
+
+
+class GremlinServerError(Exception):
+    def __init__(self, status):
+        super(GremlinServerError, self).__init__('{0}: {1}'.format(status['code'], status['message']))
+        self._status_attributes = status['attributes']
+        self.status_code = status['code']
+        self.status_message = status['message']
+
+    @property
+    def status_attributes(self):
+        return self._status_attributes
+
+
+class ConfigurationError(Exception):
+    pass
+
+
+class AbstractBaseProtocol(metaclass=abc.ABCMeta):
+
+    @abc.abstractmethod
+    def connection_made(self, transport):
+        self._transport = transport
+
+    @abc.abstractmethod
+    def data_received(self, message, results_dict):
+        pass
+
+    @abc.abstractmethod
+    def write(self, request_id, request_message):
+        pass
+
+
+class GremlinServerWSProtocol(AbstractBaseProtocol):
+    QOP_AUTH_BIT = 1
+    _kerberos_context = None
+    _max_content_length = 10 * 1024 * 1024
+
+    def __init__(self,
+                 message_serializer,
+                 username='', password='',
+                 kerberized_service='',
+                 max_content_length=10 * 1024 * 1024):
+        self._message_serializer = message_serializer
+        self._username = username
+        self._password = password
+        self._kerberized_service = kerberized_service
+        self._max_content_length = max_content_length
+
+    def connection_made(self, transport):
+        super(GremlinServerWSProtocol, self).connection_made(transport)
+
+    def write(self, request_id, request_message):
+        message = self._message_serializer.serialize_message(request_id, request_message)
+        self._transport.write(message)
+
+    def data_received(self, message, results_dict):
+        # if Gremlin Server cuts off then we get a None for the message
+        if message is None:
+            log.error("Received empty message from server.")
+            raise GremlinServerError({'code': 500,
+                                      'message': 'Server disconnected - please try to reconnect', 'attributes': {}})
+
+        message = self._message_serializer.deserialize_message(message)
+        request_id = message['requestId']
+        result_set = results_dict[request_id] if request_id in results_dict else ResultSet(None, None)
+        status_code = message['status']['code']
+        aggregate_to = message['result']['meta'].get('aggregateTo', 'list')
+        data = message['result']['data']
+        result_set.aggregate_to = aggregate_to
+        if status_code == 407:
+            if self._username and self._password:
+                auth_bytes = b''.join([b'\x00', self._username.encode('utf-8'),
+                                       b'\x00', self._password.encode('utf-8')])
+                auth = base64.b64encode(auth_bytes)
+                request_message = request.RequestMessage(
+                    'traversal', 'authentication', {'sasl': auth.decode()})
+            elif self._kerberized_service:
+                request_message = self._kerberos_received(message)
+            else:
+                error_message = 'Gremlin server requires authentication credentials in DriverRemoteConnection. ' \
+                                'For basic authentication provide username and password. ' \
+                                'For kerberos authentication provide the kerberized_service parameter.'
+                log.error(error_message)
+                raise ConfigurationError(error_message)
+            self.write(request_id, request_message)
+            data = self._transport.read()
+            # Allow for auth handshake with multiple steps
+            return self.data_received(data, results_dict)
+        elif status_code == 204:
+            result_set.stream.put_nowait([])
+            del results_dict[request_id]
+            return status_code
+        elif status_code in [200, 206]:
+            result_set.stream.put_nowait(data)
+            if status_code == 200:
+                result_set.status_attributes = message['status']['attributes']
+                del results_dict[request_id]
+            return status_code
+        else:
+            # This message is going to be huge and kind of hard to read, but in the event of an error,
+            # it can provide invaluable info, so space it out appropriately.
+            log.error("\r\nReceived error message '%s'\r\n\r\nWith results dictionary '%s'",
+                      str(message), str(results_dict))
+            del results_dict[request_id]
+            raise GremlinServerError(message['status'])
+
+    def _kerberos_received(self, message):
+        # Inspired by: https://github.com/thobbs/pure-sasl/blob/0.6.2/puresasl/mechanisms.py
+        #              https://github.com/thobbs/pure-sasl/blob/0.6.2/LICENSE
+        try:
+            import kerberos
+        except ImportError:
+            raise ImportError('Please install gremlinpython[kerberos].')
+
+        # First pass: get service granting ticket and return it to gremlin-server
+        if not self._kerberos_context:
+            try:
+                _, kerberos_context = kerberos.authGSSClientInit(
+                    self._kerberized_service, gssflags=kerberos.GSS_C_MUTUAL_FLAG)
+                kerberos.authGSSClientStep(kerberos_context, '')
+                auth = kerberos.authGSSClientResponse(kerberos_context)
+                self._kerberos_context = kerberos_context
+            except kerberos.KrbError as e:
+                raise ConfigurationError(
+                    'Kerberos authentication requires a valid service name in DriverRemoteConnection, '
+                    'as well as a valid tgt (export KRB5CCNAME) or keytab (export KRB5_KTNAME): ' + str(e))
+            return request.RequestMessage('', 'authentication', {'sasl': auth})
+
+        # Second pass: completion of authentication
+        sasl_response = message['status']['attributes']['sasl']
+        if not self._username:
+            result_code = kerberos.authGSSClientStep(self._kerberos_context, sasl_response)
+            if result_code == kerberos.AUTH_GSS_COMPLETE:
+                self._username = kerberos.authGSSClientUserName(self._kerberos_context)
+            return request.RequestMessage('', 'authentication', {'sasl': ''})
+
+        # Third pass: sasl quality of protection (qop) handshake
+
+        # Gremlin-server Krb5Authenticator only supports qop=QOP_AUTH; use ssl for confidentiality.
+        # Handshake content format:
+        # byte 0: the selected qop. 1==auth, 2==auth-int, 4==auth-conf
+        # byte 1-3: the max length for any buffer sent back and forth on this connection. (big endian)
+        # the rest of the buffer: the authorization user name in UTF-8 - not null terminated.
+        kerberos.authGSSClientUnwrap(self._kerberos_context, sasl_response)
+        data = kerberos.authGSSClientResponse(self._kerberos_context)
+        plaintext_data = base64.b64decode(data)
+        assert len(plaintext_data) == 4, "Unexpected response from gremlin server sasl handshake"
+        word, = struct.unpack('!I', plaintext_data)
+        qop_bits = word >> 24
+        assert self.QOP_AUTH_BIT & qop_bits, "Unexpected sasl qop level received from gremlin server"
+
+        name_length = len(self._username)
+        fmt = '!I' + str(name_length) + 's'
+        word = self.QOP_AUTH_BIT << 24 | self._max_content_length
+        out = struct.pack(fmt, word, self._username.encode("utf-8"), )
+        encoded = base64.b64encode(out).decode('ascii')
+        kerberos.authGSSClientWrap(self._kerberos_context, encoded)
+        auth = kerberos.authGSSClientResponse(self._kerberos_context)
+        return request.RequestMessage('', 'authentication', {'sasl': auth})
+
+
+class GremlinServerHTTPProtocol(AbstractBaseProtocol):
+
+    def __init__(self,
+                 message_serializer,
+                 username='', password=''):
+        self._message_serializer = message_serializer
+        self._username = username
+        self._password = password
+
+    def connection_made(self, transport):
+        super(GremlinServerHTTPProtocol, self).connection_made(transport)
+
+    def write(self, request_id, request_message):
+
+        basic_auth = {}
+        if self._username and self._password:
+            basic_auth['username'] = self._username
+            basic_auth['password'] = self._password
+
+        content_type = str(self._message_serializer.version, encoding='utf-8')
+        message = {
+            'headers': {'CONTENT-TYPE': content_type,
+                        'ACCEPT': content_type},
+            'payload': self._message_serializer.serialize_message(request_id, request_message),
+            'auth': basic_auth
+        }
+
+        self._transport.write(message)
+
+    def data_received(self, response, results_dict):
+        # if Gremlin Server cuts off then we get a None for the message
+        if response is None:
+            log.error("Received empty message from server.")
+            raise GremlinServerError({'code': 500,
+                                      'message': 'Server disconnected - please try to reconnect', 'attributes': {}})
+
+        if response['ok']:
+            message = self._message_serializer.deserialize_message(response['content'])
+            request_id = message['requestId']
+            result_set = results_dict[request_id] if request_id in results_dict else ResultSet(None, None)
+            status_code = message['status']['code']
+            aggregate_to = message['result']['meta'].get('aggregateTo', 'list')
+            data = message['result']['data']
+            result_set.aggregate_to = aggregate_to
+
+            if status_code == 204:
+                result_set.stream.put_nowait([])
+                del results_dict[request_id]
+                return status_code
+            elif status_code in [200, 206]:
+                result_set.stream.put_nowait(data)
+                if status_code == 200:
+                    result_set.status_attributes = message['status']['attributes']
+                    del results_dict[request_id]
+                return status_code
+        else:
+            # This message is going to be huge and kind of hard to read, but in the event of an error,
+            # it can provide invaluable info, so space it out appropriately.
+            log.error("\r\nReceived error message '%s'\r\n\r\nWith results dictionary '%s'",
+                      str(response['content']), str(results_dict))
+            body = json.loads(response['content'])
+            del results_dict[body['requestId']]
+            raise GremlinServerError({'code': response['status'], 'message': body['message'], 'attributes': {}})

gremlin_python/driver/remote_connection.py

@@ -0,0 +1,84 @@
+#
+# 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 abc
+
+from gremlin_python.process import traversal
+
+__author__ = 'Marko A. Rodriguez (http://markorodriguez.com), Lyndon Bauto (lyndonb@bitquilltech.com)'
+
+
+class RemoteConnection(object, metaclass=abc.ABCMeta):
+    def __init__(self, url, traversal_source):
+        self._url = url
+        self._traversal_source = traversal_source
+
+    @property
+    def url(self):
+        return self._url
+
+    @property
+    def traversal_source(self):
+        return self._traversal_source
+
+    @abc.abstractmethod
+    def submit(self, bytecode):
+        pass
+
+    def is_closed(self):
+        raise Exception('is_closed() must be implemented')
+
+    def is_session_bound(self):
+        return False
+
+    def create_session(self):
+        raise Exception('createSession() must be implemented');
+
+    def commit(self):
+        raise Exception('commit() must be implemented')
+
+    def rollback(self):
+        raise Exception('rollback() must be implemented')
+
+    def __repr__(self):
+        return "remoteconnection[" + self._url + "," + self._traversal_source + "]"
+
+
+class RemoteTraversal(traversal.Traversal):
+    def __init__(self, traversers):
+        super(RemoteTraversal, self).__init__(None, None, None)
+        self.traversers = traversers
+
+
+class RemoteStrategy(traversal.TraversalStrategy):
+    def __init__(self, remote_connection):
+        # Gave this a fqcn that has a local "py:" prefix since this strategy isn't sent as bytecode to the server.
+        # this is a sort of local-only strategy that actually executes client side. not sure if this prefix is the
+        # right way to name this or not, but it should have a name to identify it.
+        traversal.TraversalStrategy.__init__(self, fqcn="py:RemoteStrategy")
+        self.remote_connection = remote_connection
+
+    def apply(self, traversal):
+        if traversal.traversers is None:
+            remote_traversal = self.remote_connection.submit(traversal.bytecode)
+            traversal.remote_results = remote_traversal
+            traversal.traversers = remote_traversal.traversers
+
+    def apply_async(self, traversal):
+        if traversal.traversers is None:
+            traversal.remote_results = self.remote_connection.submit_async(traversal.bytecode)

gremlin_python/driver/request.py

@@ -0,0 +1,25 @@
+#
+# 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 collections
+
+__author__ = 'David M. Brown (davebshow@gmail.com)'
+
+
+RequestMessage = collections.namedtuple(
+    'RequestMessage', ['processor', 'op', 'args'])

gremlin_python/driver/resultset.py

@@ -0,0 +1,100 @@
+#
+# 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.
+#
+from concurrent.futures import Future
+
+__author__ = 'David M. Brown (davebshow@gmail.com)'
+
+
+class ResultSet:
+
+    def __init__(self, stream, request_id):
+        self._stream = stream
+        self._request_id = request_id
+        self._done = None
+        self._aggregate_to = None
+        self._status_attributes = {}
+
+    @property
+    def aggregate_to(self):
+        return self._aggregate_to
+
+    @aggregate_to.setter
+    def aggregate_to(self, val):
+        self._aggregate_to = val
+
+    @property
+    def status_attributes(self):
+        return self._status_attributes
+
+    @status_attributes.setter
+    def status_attributes(self, val):
+        self._status_attributes = val
+
+    @property
+    def request_id(self):
+        return self._request_id
+
+    @property
+    def stream(self):
+        return self._stream
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        result = self.one()
+        if not result:
+            raise StopIteration
+        return result
+
+    def next(self):
+        return self.__next__()
+
+    @property
+    def done(self):
+        return self._done
+
+    @done.setter
+    def done(self, future):
+        self._done = future
+
+    def one(self):
+        while not self.done.done():
+            if not self.stream.empty():
+                return self.stream.get_nowait()
+        if not self.stream.empty():
+            return self.stream.get_nowait()
+        return self.done.result()
+
+    def all(self):
+        future = Future()
+
+        def cb(f):
+            try:
+                f.result()
+            except Exception as e:
+                future.set_exception(e)
+            else:
+                results = []
+                while not self.stream.empty():
+                    results += self.stream.get_nowait()
+                future.set_result(results)
+
+        self.done.add_done_callback(cb)
+        return future