driver_remote_connection.py 3.44 KB
Newer Older
davebshow's avatar
davebshow committed
1
2
3
4
import asyncio
from urllib.parse import urlparse

from aiogremlin.driver.cluster import Cluster
5
from gremlin_python.driver import serializer
davebshow's avatar
davebshow committed
6
from gremlin_python.driver.remote_connection import RemoteTraversal
davebshow's avatar
davebshow committed
7
8
9
10
11
12


__author__ = 'David M. Brown (davebshow@gmail.com)'


class DriverRemoteConnection:
davebshow's avatar
davebshow committed
13
14
15
16
17
18
19
20
21
    """
    Remote connection to a Gremlin Server. Do not instantiate directly,
    instead use :py:meth:`DriverRemoteConnection.open` or
    :py:meth:`DriverRemoteConnection.using`

    :param aiogremlin.driver.client.Client client:
    :param asyncio.BaseEventLoop loop:
    :param aiogremlin.driver.cluster.Cluster cluster:
    """
davebshow's avatar
davebshow committed
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

    def __init__(self, client, loop, *, cluster=None):
        self._client = client
        self._loop = loop
        self._cluster = cluster

    @property
    def client(self):
        return self._client

    @property
    def config(self):
        return self._cluster.config

    @classmethod
davebshow's avatar
davebshow committed
37
38
39
40
41
42
43
44
45
    async def using(cls, cluster, aliases=None):
        """
        Create a :py:class:`DriverRemoteConnection` using a specific
        :py:class:`Cluster<aiogremlin.driver.cluster.Cluster>`

        :param aiogremlin.driver.cluster.Cluster cluster:
        :param dict aliases: Optional mapping for aliases. Default is `None`.
            Also accepts `str` argument which will be assigned to `g`
        """
davebshow's avatar
davebshow committed
46
        client = await cluster.connect(aliases=aliases)
davebshow's avatar
davebshow committed
47
        loop = cluster._loop
davebshow's avatar
davebshow committed
48
49
50
51
52
        return cls(client, loop)

    @classmethod
    async def open(cls, url=None, aliases=None, loop=None, *,
                   graphson_reader=None, graphson_writer=None, **config):
davebshow's avatar
davebshow committed
53
54
55
56
57
58
59
60
61
62
        """
        :param str url: Optional url for host Gremlin Server

        :param dict aliases: Optional mapping for aliases. Default is `None`.
            Also accepts `str` argument which will be assigned to `g`
        :param asyncio.BaseEventLoop loop:
        :param graphson_reader: Custom graphson_reader
        :param graphson_writer: Custom graphson_writer
        :param config: Optional cluster configuration passed as kwargs or `dict`
        """
davebshow's avatar
davebshow committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        if url:
            parsed_url = urlparse(url)
            config.update({
                'scheme': parsed_url.scheme,
                'hosts': [parsed_url.hostname],
                'port': parsed_url.port})
        if isinstance(aliases, str):
            aliases = {'g': aliases}
        if not loop:
            loop = asyncio.get_event_loop()
        message_serializer = serializer.GraphSONMessageSerializer(
            reader=graphson_reader,
            writer=graphson_writer)
        config.update({'message_serializer': message_serializer})
        cluster = await Cluster.open(loop, aliases=aliases, **config)
        client = await cluster.connect()
        return cls(client, loop, cluster=cluster)

    async def close(self):
davebshow's avatar
davebshow committed
82
83
84
85
        """
        Close underlying cluster if applicable. If created with
        :py:meth:`DriverRemoteConnection.using`, cluster is NOT closed.
        """
davebshow's avatar
davebshow committed
86
87
88
89
        if self._cluster:
            await self._cluster.close()

    async def submit(self, bytecode):
davebshow's avatar
davebshow committed
90
        """Submit bytecode to the Gremlin Server"""
davebshow's avatar
davebshow committed
91
        result_set = await self._client.submit(bytecode)
92
93
94
        results_all = await result_set.all()
        results = results_all.result()
        return RemoteTraversal(iter(results))
davebshow's avatar
davebshow committed
95
96
97
98
99
100
101
102

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()
        self._client = None
        self._cluster = None