diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc index a5ecd33da38..29ec207794d 100644 --- a/CHANGELOG.asciidoc +++ b/CHANGELOG.asciidoc @@ -27,6 +27,11 @@ image::https://raw.githubusercontent.com/apache/tinkerpop/master/docs/static/ima This release also includes changes from prior 3.7.x releases. +* Added native-async WebSocket driver for Python (`gremlin_python.driver`): low-level transport (`AsyncAiohttpWSTransport`), protocol (`AsyncGremlinServerWSProtocol`), result streaming (`AsyncResultSet`), connection (`AsyncConnection`), and client (`AsyncClient`) — full asyncio usage without thread-pool wrappers. +* Added `AsyncDriverRemoteConnection` for the Python GLV async driver: a high-level connection object equivalent to `DriverRemoteConnection` that exposes `submit()`, `submit_async()`, and `submit_stream()` as coroutines, supports server-side sessions, and automatically rolls back any open transaction on `close()`. +* Added `AsyncGraphTraversalSource`, `AsyncGraphTraversal`, and `async_traversal()` factory to the Python GLV: enables the full Gremlin DSL in async mode (`await g.V().has_label("person").values("name").to_list()`); terminal steps `to_list()`, `next()`, `to_set()`, `has_next()`, and `iterate()` are all awaitable. +* Added `AsyncTransaction` to the Python GLV: `asyncio.Lock`-protected transaction helper returned by `await g.tx()` with `begin()`, `commit()`, `rollback()`, and async context-manager support. +* Added keepalive pings (`ping_interval`), automatic retry on connection reset (`InvalidatedConnectionError`), and per-client default request options (including `query_timeout`) to the Python GLV async driver for production reliability. * Enabled building and running with Java 21 and Java 25 (experimental; `spark-gremlin` excluded, as Spark 3.3.x only runs on Java 8 through 17). * Bumped to Groovy 4.0.32 which adds support for parsing Java 25 bytecode. * Bumped Hadoop to 3.4.3 (and Kerby to 2.0.3) to enable `hadoop-gremlin` to build and run on Java 25. diff --git a/gremlin-python/src/main/python/examples/async_connections.py b/gremlin-python/src/main/python/examples/async_connections.py new file mode 100644 index 00000000000..f11759c4475 --- /dev/null +++ b/gremlin-python/src/main/python/examples/async_connections.py @@ -0,0 +1,108 @@ +# +# 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 asyncio +import os +import ssl +import sys +import uuid + +sys.path.append("..") + +from gremlin_python.driver.async_driver_remote_connection import AsyncDriverRemoteConnection +from gremlin_python.driver.async_graph_traversal import async_traversal + +SERVER_URL = os.getenv("GREMLIN_SERVER_URL", "ws://localhost:8182/gremlin").format(45940) +VERTEX_LABEL = os.getenv("VERTEX_LABEL", "connection") + + +async def main(): + await with_context_manager() + await with_explicit_close() + await with_auth() + await with_custom_pool() + await with_session() + + +async def with_context_manager(): + # The recommended way: async context manager guarantees close() is called + # even if an exception is raised. + async with AsyncDriverRemoteConnection(SERVER_URL, "g") as rc: + g = async_traversal().with_remote(rc) + + # Terminal steps (to_list, next, iterate) are coroutines — await them. + await g.add_v(VERTEX_LABEL).iterate() + + count = await g.V().has_label(VERTEX_LABEL).count().next() + print("with_context_manager — vertex count:", count) + + +async def with_explicit_close(): + # When a context manager is not practical, close() manually in a finally block. + rc = AsyncDriverRemoteConnection(SERVER_URL, "g") + try: + g = async_traversal().with_remote(rc) + count = await g.V().count().next() + print("with_explicit_close — vertex count:", count) + finally: + await rc.close() + + +async def with_auth(): + # Plain-text (SASL/basic) authentication. + auth_url = os.getenv("GREMLIN_SERVER_BASIC_AUTH_URL", + "ws://localhost:8182/gremlin").format(45941) + + ssl_ctx = None + if ":45941" in auth_url: + ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_ctx.check_hostname = False + ssl_ctx.verify_mode = ssl.CERT_NONE + + async with AsyncDriverRemoteConnection( + auth_url, "g", + username="stephen", password="password", + ssl=ssl_ctx, + ) as rc: + g = async_traversal().with_remote(rc) + count = await g.V().count().next() + print("with_auth — vertex count:", count) + + +async def with_custom_pool(): + # pool_size controls concurrent WebSocket connections. + async with AsyncDriverRemoteConnection(SERVER_URL, "g", pool_size=4) as rc: + g = async_traversal().with_remote(rc) + count = await g.V().count().next() + print("with_custom_pool — vertex count:", count) + + +async def with_session(): + # Server-side sessions preserve state across multiple requests. + session_id = str(uuid.uuid4()) + async with AsyncDriverRemoteConnection(SERVER_URL, "g", session=session_id) as rc: + g = async_traversal().with_remote(rc) + + await g.add_v(VERTEX_LABEL).property("session", session_id).iterate() + count = await g.V().has("session", session_id).count().next() + print("with_session — vertex count:", count) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/gremlin-python/src/main/python/examples/async_traversals.py b/gremlin-python/src/main/python/examples/async_traversals.py new file mode 100644 index 00000000000..4a224feb1fa --- /dev/null +++ b/gremlin-python/src/main/python/examples/async_traversals.py @@ -0,0 +1,117 @@ +# +# 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. +# + +# This example requires the Modern toy graph to be preloaded when Gremlin Server +# starts. See https://tinkerpop.apache.org/docs/current/reference/#gremlin-server-docker-image +# and use conf/gremlin-server-modern.yaml. + +import asyncio +import os +import sys + +sys.path.append("..") + +from gremlin_python.driver.async_driver_remote_connection import AsyncDriverRemoteConnection +from gremlin_python.driver.async_graph_traversal import async_traversal +from gremlin_python.process.graph_traversal import __ +from gremlin_python.process.traversal import P, T + +SERVER_URL = os.getenv("GREMLIN_SERVER_URL", "ws://localhost:8182/gremlin").format(45940) + +# CI uses port 45940 with gmodern binding; local uses 8182 with g binding. +GRAPH_BINDING = "gmodern" if ":45940" in SERVER_URL else "g" + + +async def main(): + async with AsyncDriverRemoteConnection(SERVER_URL, GRAPH_BINDING) as rc: + g = async_traversal().with_remote(rc) + + await basic_queries(g) + await write_operations(g) + await concurrent_queries(g) + await traversal_steps(g) + + +async def basic_queries(g): + # to_list() collects all results into a list — identical name to the sync + # version, but here it is a coroutine that must be awaited. + names = await g.V().values("name").to_list() + print("basic_queries — all names:", names) + + # next() returns the first result (raises StopAsyncIteration if empty). + count = await g.V().count().next() + print("basic_queries — vertex count:", count) + + # to_set() deduplicates results. + labels = await g.V().label().to_set() + print("basic_queries — unique labels:", labels) + + # has_next() returns True/False without consuming the traversal. + exists = await g.V().has("name", "marko").has_next() + print("basic_queries — marko exists:", exists) + + # Filtering and projection with snake_case step names. + person_maps = await g.V().has_label("person").value_map("name", "age").to_list() + print("basic_queries — person maps:", person_maps) + + +async def write_operations(g): + # iterate() discards results and is the right choice for writes. + # It adds a "discard" step so the server knows no data need be returned. + await g.add_v("planet").property("name", "saturn").iterate() + + saturn = await g.V().has("name", "saturn").next() + print("write_operations — created:", saturn) + + # Clean up. + await g.V().has("name", "saturn").drop().iterate() + + +async def concurrent_queries(g): + # asyncio.gather() runs multiple traversals concurrently over the shared + # connection pool — no extra connections are opened. + names_coro = g.V().has_label("person").values("name").to_list() + edges_coro = g.E().label().to_set() + software_coro = g.V().has_label("software").values("name").to_list() + + names, edges, software = await asyncio.gather(names_coro, edges_coro, software_coro) + print("concurrent_queries — people: ", names) + print("concurrent_queries — edge types:", edges) + print("concurrent_queries — software: ", software) + + +async def traversal_steps(g): + # The full Gremlin DSL is available — all intermediate steps work + # exactly as in the sync version. + e1 = await g.V(1).both_e().to_list() + print("traversal_steps (1) — all edges of v1:", e1) + + e2 = await g.V(1).both_e().where(__.other_v().has_id(2)).to_list() + print("traversal_steps (2) — edges between v1 and v2:", e2) + + e3 = await g.V(1).out_e().where(__.in_v().has(T.id, P.within(2, 3))).to_list() + print("traversal_steps (3) — out-edges of v1 to v2 or v3:", e3) + + # next(n) returns the first n results as a list. + first_two = await g.V().has_label("person").values("name").next(2) + print("traversal_steps — first two person names:", first_two) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/gremlin-python/src/main/python/gremlin_python/driver/aiohttp/async_transport.py b/gremlin-python/src/main/python/gremlin_python/driver/aiohttp/async_transport.py new file mode 100644 index 00000000000..3ac528d259d --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/aiohttp/async_transport.py @@ -0,0 +1,145 @@ +# +# 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 asyncio +import sys + +import aiohttp +from aiohttp import ClientResponseError + +from gremlin_python.driver.transport import AbstractBaseTransport + +if sys.version_info >= (3, 11): + from asyncio import timeout as _async_timeout +else: + from async_timeout import timeout as _async_timeout + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +class AsyncAiohttpWSTransport(AbstractBaseTransport): + """Native-async WebSocket transport built on aiohttp. + + Unlike :class:`~gremlin_python.driver.aiohttp.transport.AiohttpTransport`, + every I/O method is a coroutine. The transport does **not** create its own + event loop; callers must already be running inside one (``async def`` + contexts, ``asyncio.run()``, etc.). + + Parameters + ---------- + read_timeout: + Maximum seconds to wait for a WebSocket frame. ``None`` disables the + timeout. + write_timeout: + Maximum seconds to wait when sending a frame. ``None`` disables the + timeout. + **kwargs: + Additional keyword arguments forwarded verbatim to + :meth:`aiohttp.ClientSession.ws_connect`. ``max_content_length`` is + transparently remapped to ``max_msg_size``; ``ssl_options`` is + remapped to ``ssl``. + """ + + def __init__(self, read_timeout=None, write_timeout=None, **kwargs): + self._websocket = None + self._client_session = None + self._read_timeout = read_timeout + self._write_timeout = write_timeout + self._aiohttp_kwargs = kwargs + if "max_content_length" in self._aiohttp_kwargs: + self._aiohttp_kwargs["max_msg_size"] = self._aiohttp_kwargs.pop( + "max_content_length" + ) + if "ssl_options" in self._aiohttp_kwargs: + self._aiohttp_kwargs["ssl"] = self._aiohttp_kwargs.pop("ssl_options") + + async def connect(self, url, headers=None): + """Open the WebSocket connection to *url*. + + Parameters + ---------- + url: + WebSocket endpoint, e.g. ``ws://localhost:8182/gremlin``. + headers: + Optional extra HTTP headers sent during the WebSocket handshake. + """ + self._client_session = aiohttp.ClientSession() + try: + self._websocket = await self._client_session.ws_connect( + url, headers=headers, **self._aiohttp_kwargs + ) + except ClientResponseError as err: + if err.status == 403: + raise Exception( + "Failed to connect to server: HTTP Error code 403 - Forbidden." + ) + raise + + async def write(self, message): + """Send *message* (bytes) as a WebSocket binary frame.""" + async with _async_timeout(self._write_timeout): + await self._websocket.send_bytes(message) + + async def read(self): + """Receive the next WebSocket message and return its payload as bytes. + + Text frames are UTF-8 encoded before returning. Close/error frames + raise :class:`RuntimeError`. + """ + async with _async_timeout(self._read_timeout): + msg = await self._websocket.receive() + + if msg.type == aiohttp.WSMsgType.CLOSE: + await self.close() + raise RuntimeError("Connection was closed by server.") + if msg.type == aiohttp.WSMsgType.CLOSED: + raise RuntimeError("Connection was already closed.") + if msg.type == aiohttp.WSMsgType.ERROR: + raise RuntimeError("Received error on read: '" + str(msg.data) + "'") + if msg.type == aiohttp.WSMsgType.TEXT: + return msg.data.strip().encode("utf-8") + return msg.data + + async def ping(self): + """Send a WebSocket ping frame to keep the connection alive.""" + await self._websocket.ping() + + async def close(self): + """Close the WebSocket and the underlying HTTP session gracefully. + + Both resources are closed concurrently and each close is shielded from + external cancellation so that the connection is always cleaned up even + if the caller's coroutine is cancelled. + """ + tasks = [] + if self._websocket is not None and not self._websocket.closed: + tasks.append(asyncio.ensure_future(self._websocket.close())) + if self._client_session is not None and not self._client_session.closed: + tasks.append(asyncio.ensure_future(self._client_session.close())) + if tasks: + await asyncio.gather(*[asyncio.shield(t) for t in tasks]) + + @property + def closed(self): + """``True`` if the WebSocket or the underlying session is closed.""" + return ( + self._websocket is None + or self._websocket.closed + or self._client_session is None + or self._client_session.closed + ) diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_client.py b/gremlin-python/src/main/python/gremlin_python/driver/async_client.py new file mode 100644 index 00000000000..9ec40807355 --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_client.py @@ -0,0 +1,290 @@ +# +# 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 asyncio +import logging + +from gremlin_python.driver import protocol, request, serializer +from gremlin_python.driver.async_connection import AsyncConnection, InvalidatedConnectionError +from gremlin_python.driver.async_protocol import AsyncGremlinServerWSProtocol +from gremlin_python.driver.aiohttp.async_transport import AsyncAiohttpWSTransport +from gremlin_python.driver.client import Client +from gremlin_python.process import traversal + +log = logging.getLogger("gremlinpython") + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +class AsyncClient(Client): + """Native-async Gremlin client using WebSocket transport. + + Unlike :class:`~gremlin_python.driver.client.Client`, no + :class:`~concurrent.futures.ThreadPoolExecutor` is created. The + connection pool is an :class:`asyncio.Queue` and all submit/close + methods are coroutines that must be awaited. + + Robustness features: + + * Each connection runs a background WebSocket ping to prevent idle + disconnects (configurable via ``ping_interval``). + * On :class:`~gremlin_python.driver.async_connection.InvalidatedConnectionError` + the request is automatically retried on a fresh connection up to + *pool_size* times. + * *default_request_options* are merged into every request so options + like ``evaluationTimeout`` can be set once per client. + + Usage:: + + async with AsyncClient("ws://localhost:8182/gremlin", "g") as client: + result_set = await client.submit("g.V().count()") + print(await result_set.all()) + + Parameters + ---------- + url: + WebSocket endpoint, e.g. ``ws://localhost:8182/gremlin``. + traversal_source: + Gremlin traversal source name (default ``"g"``). + protocol_factory: + Zero-argument callable returning an + :class:`~gremlin_python.driver.async_protocol.AsyncGremlinServerWSProtocol`. + transport_factory: + Zero-argument callable returning an + :class:`~gremlin_python.driver.aiohttp.async_transport.AsyncAiohttpWSTransport`. + pool_size: + Number of concurrent WebSocket connections (default 8; forced to 1 + in session mode). + message_serializer: + Serializer to use (default + :class:`~gremlin_python.driver.serializer.GraphBinarySerializersV1`). + username, password: + Credentials for SASL/basic authentication. + kerberized_service: + Kerberos service principal for Kerberos authentication. + headers: + Extra HTTP headers included in every WebSocket handshake. + session: + Session UUID string to enable server-side sessions. + enable_user_agent_on_connect: + Send the driver user-agent header (default ``True``). + ping_interval: + Seconds between WebSocket keepalive pings (default 60, 0 to disable). + default_request_options: + Dict of request options merged into every request + (e.g. ``{"evaluationTimeout": 30000}``). Per-request options take + precedence over defaults. + **transport_kwargs: + Extra keyword arguments forwarded to the default + :class:`~gremlin_python.driver.aiohttp.async_transport.AsyncAiohttpWSTransport`. + """ + + def __init__( + self, + url, + traversal_source, + protocol_factory=None, + transport_factory=None, + pool_size=None, + max_workers=None, # accepted but unused — no thread pool + message_serializer=None, + username="", + password="", + kerberized_service="", + headers=None, + session=None, + enable_user_agent_on_connect=True, + ping_interval=60, + default_request_options=None, + **transport_kwargs, + ): + log.info("Creating AsyncClient with url '%s'", url) + + # Do NOT call super().__init__() — Client.__init__ creates a + # ThreadPoolExecutor and a threading queue, neither of which we want. + self._closed = False + self._url = url + self._headers = headers + self._enable_user_agent_on_connect = enable_user_agent_on_connect + self._traversal_source = traversal_source + self._executor = None # no thread pool; attribute expected by Client + self._ping_interval = ping_interval + self._default_request_options = default_request_options or {} + + if "max_content_length" not in transport_kwargs: + transport_kwargs["max_content_length"] = 10 * 1024 * 1024 + + if message_serializer is None: + message_serializer = serializer.GraphBinarySerializersV1() + + self._message_serializer = message_serializer + self._username = username + self._password = password + self._session = session + self._session_enabled = session is not None and session != "" + + if transport_factory is None: + def transport_factory(): + return AsyncAiohttpWSTransport(**transport_kwargs) + + self._transport_factory = transport_factory + + if protocol_factory is None: + def protocol_factory(): + return AsyncGremlinServerWSProtocol( + self._message_serializer, + username=self._username, + password=self._password, + kerberized_service=kerberized_service, + max_content_length=transport_kwargs["max_content_length"], + ) + + self._protocol_factory = protocol_factory + + if self._session_enabled: + if pool_size is None: + pool_size = 1 + elif pool_size != 1: + raise Exception("PoolSize must be 1 on session mode!") + if pool_size is None: + pool_size = 8 + self._pool_size = pool_size + + self._pool = asyncio.Queue() + self._fill_pool() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _get_connection(self): + proto = self._protocol_factory() + return AsyncConnection( + self._url, + self._traversal_source, + proto, + self._transport_factory, + None, + self._pool, + headers=self._headers, + enable_user_agent_on_connect=self._enable_user_agent_on_connect, + ping_interval=self._ping_interval, + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def close(self): + """Close all pooled connections and mark the client as closed.""" + if self._closed: + return + if self._session_enabled: + await self._close_session() + log.info("Closing AsyncClient with url '%s'", self._url) + conns = [] + while not self._pool.empty(): + conns.append(self._pool.get_nowait()) + await asyncio.gather(*(conn.close() for conn in conns), return_exceptions=True) + self._closed = True + + async def _close_session(self): + message = request.RequestMessage( + processor="session", + op="close", + args={"session": str(self._session)}, + ) + conn = await self._pool.get() + try: + result_set = await conn.write(message) + await result_set.all() + except protocol.GremlinServerError: + pass + + # ------------------------------------------------------------------ + # Context manager + # ------------------------------------------------------------------ + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + await self.close() + + # ------------------------------------------------------------------ + # Submit (with retry) + # ------------------------------------------------------------------ + + async def submit(self, message, bindings=None, request_options=None): + """Submit *message* and return an :class:`~gremlin_python.driver.async_resultset.AsyncResultSet`. + + Retries automatically up to *pool_size* times when a connection is + invalidated mid-request due to a network reset. + """ + last_error = None + for _ in range(self._pool_size): + try: + return await self.submit_async( + message, bindings=bindings, request_options=request_options + ) + except InvalidatedConnectionError as exc: + log.warning("Connection invalidated, retrying: %s", exc.__cause__) + last_error = exc.__cause__ + raise last_error or InvalidatedConnectionError("All retry attempts exhausted") + + async def submit_async(self, message, bindings=None, request_options=None): + """Submit *message* without retry and return an :class:`~gremlin_python.driver.async_resultset.AsyncResultSet`. + + Returns before all results have arrived; the result set can be + iterated asynchronously (``async for``) or consumed via + :meth:`~gremlin_python.driver.async_resultset.AsyncResultSet.all`. + """ + if self.is_closed(): + raise Exception("Client is closed") + + args = {"gremlin": message, "aliases": {"g": self._traversal_source}} + processor = "" + op = "eval" + if isinstance(message, traversal.Bytecode): + op = "bytecode" + processor = "traversal" + + if isinstance(message, str) and bindings: + args["bindings"] = bindings + + if self._session_enabled: + args["session"] = str(self._session) + processor = "session" + + if isinstance(message, (traversal.Bytecode, str)): + message = request.RequestMessage( + processor=processor, op=op, args=args + ) + + conn = await self._pool.get() + + # Merge default options first; per-request options win. + if self._default_request_options: + merged = {**self._default_request_options} + if request_options: + merged.update(request_options) + message.args.update(merged) + elif request_options: + message.args.update(request_options) + + return await conn.write(message) diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_connection.py b/gremlin-python/src/main/python/gremlin_python/driver/async_connection.py new file mode 100644 index 00000000000..7ce8ea49c90 --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_connection.py @@ -0,0 +1,203 @@ +# +# 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 asyncio +import logging +import uuid + +from gremlin_python.driver.async_resultset import AsyncResultSet +from gremlin_python.driver.connection import Connection + +log = logging.getLogger("gremlinpython") + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +class InvalidatedConnectionError(Exception): + """Raised when a connection was reset during a request. + + :class:`~gremlin_python.driver.async_client.AsyncClient` catches this and + retries the request on a fresh connection. + """ + + +class AsyncConnection(Connection): + """Native-async WebSocket connection with keepalive and auto-reconnect. + + Features beyond the sync :class:`~gremlin_python.driver.connection.Connection`: + + * **Ping keepalive** — a background :class:`asyncio.Task` sends a WebSocket + ping every *ping_interval* seconds to prevent the server or intermediate + proxies from closing idle connections. + * **Auto-reconnect** — a ``ConnectionResetError`` during *write* triggers a + single reconnect + retry before giving up. + * **Invalidation signal** — a ``ConnectionResetError`` during *read* raises + :class:`InvalidatedConnectionError` so the client can transparently retry + the request on another connection. + + Parameters + ---------- + url: + WebSocket endpoint. + traversal_source: + Gremlin traversal source name (e.g. ``"g"``). + protocol: + An :class:`~gremlin_python.driver.async_protocol.AsyncGremlinServerWSProtocol` + instance. + transport_factory: + Zero-argument callable returning an + :class:`~gremlin_python.driver.aiohttp.async_transport.AsyncAiohttpWSTransport`. + executor: + Ignored; present for API compatibility with + :class:`~gremlin_python.driver.connection.Connection`. + pool: + The :class:`asyncio.Queue` that owns this connection. + headers: + Extra HTTP headers sent during the WebSocket handshake. + enable_user_agent_on_connect: + Whether to send the driver user-agent header. + ping_interval: + Seconds between WebSocket ping frames (default 60, 0 to disable). + """ + + def __init__(self, url, traversal_source, protocol, transport_factory, + executor, pool, headers=None, enable_user_agent_on_connect=True, + ping_interval=60): + super().__init__( + url, traversal_source, protocol, transport_factory, + executor, pool, headers=headers, + enable_user_agent_on_connect=enable_user_agent_on_connect, + ) + self._ping_interval = ping_interval + self._ping_task = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def connect(self): + """Open (or re-open) the WebSocket connection.""" + if self._ping_task is not None: + self._ping_task.cancel() + self._ping_task = None + if self._transport is not None: + await self._transport.close() + self._transport = self._transport_factory() + await self._transport.connect(self._url, self._headers) + self._protocol.connection_made(self._transport) + self._inited = True + if self._ping_interval > 0: + self._ping_task = asyncio.create_task(self._ping_forever()) + + async def close(self): + """Close the connection, cancelling the ping task first.""" + if self._ping_task is not None: + self._ping_task.cancel() + self._ping_task = None + if self._inited: + # asyncio.shield protects the close coroutine from cancellation so + # the underlying transport is always cleanly shut down. + await asyncio.shield(self._transport.close()) + + # ------------------------------------------------------------------ + # Keepalive + # ------------------------------------------------------------------ + + async def _ping_forever(self): + """Send periodic WebSocket ping frames until cancelled or the transport dies.""" + while True: + try: + await asyncio.sleep(self._ping_interval) + await self._transport.ping() + except asyncio.CancelledError: + break + except Exception: + # Transport is dead; mark the connection as needing reconnect. + log.debug("Ping failed on %s — marking connection as uninitialised", self._url) + self._inited = False + break + + # ------------------------------------------------------------------ + # I/O + # ------------------------------------------------------------------ + + async def write(self, request_message): + """Send *request_message* and return an :class:`~gremlin_python.driver.async_resultset.AsyncResultSet`. + + Lazily connects on the first call. On a ``ConnectionResetError`` + during write, reconnects and retries once before propagating. + Returns as soon as the first response chunk arrives; a background + :class:`asyncio.Task` drains any remaining partial (206) chunks. + """ + if not self._inited: + await self.connect() + + if request_message.args.get("requestId"): + request_id = str(request_message.args["requestId"]) + uuid.UUID(request_id) + else: + request_id = str(uuid.uuid4()) + + result_set = AsyncResultSet(asyncio.Queue(), request_id) + self._results[request_id] = result_set + + try: + await self._protocol.write(request_id, request_message) + except ConnectionResetError: + log.debug("Write failed on %s — reconnecting and retrying", self._url) + await self.connect() + await self._protocol.write(request_id, request_message) + + return await self._receive(result_set) + + async def _receive(self, result_set): + """Read the first response chunk and set up background consumption.""" + try: + data = await self._transport.read() + status_code = await self._protocol.data_received(data, self._results) + except ConnectionResetError as exc: + self._inited = False + self._pool.put_nowait(self) + raise InvalidatedConnectionError(str(exc)) from exc + except Exception: + self._inited = False + self._pool.put_nowait(self) + raise + + if status_code != 206: + fut = asyncio.get_running_loop().create_future() + fut.set_result(None) + result_set.future = fut + self._pool.put_nowait(self) + else: + async def _consume(): + try: + while True: + try: + chunk = await self._transport.read() + sc = await self._protocol.data_received(chunk, self._results) + except ConnectionResetError as exc: + raise InvalidatedConnectionError(str(exc)) from exc + if sc != 206: + break + finally: + self._pool.put_nowait(self) + + result_set.future = asyncio.create_task(_consume()) + + return result_set diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_driver_remote_connection.py b/gremlin-python/src/main/python/gremlin_python/driver/async_driver_remote_connection.py new file mode 100644 index 00000000000..fcc2d9c7a9f --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_driver_remote_connection.py @@ -0,0 +1,297 @@ +# +# 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 asyncio +import logging +import uuid + +from gremlin_python.driver.async_client import AsyncClient +from gremlin_python.driver.remote_connection import RemoteConnection +from gremlin_python.process.strategies import OptionsStrategy +from gremlin_python.process.traversal import Bytecode + +log = logging.getLogger("gremlinpython") + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +class AsyncDriverRemoteConnection(RemoteConnection): + """Async remote connection backed by :class:`~gremlin_python.driver.async_client.AsyncClient`. + + Primary entry point for the async Gremlin driver. Works with + :func:`~gremlin_python.driver.async_graph_traversal.async_traversal` to + provide a fully async traversal DSL:: + + async with AsyncDriverRemoteConnection("ws://localhost:8182/gremlin", "g") as rc: + g = async_traversal().with_remote(rc) + names = await g.V().has_label("person").values("name").to_list() + + **Robustness features** + + * Per-connection WebSocket keepalive pings (see *ping_interval*). + * Automatic retry on network resets (handled inside + :class:`~gremlin_python.driver.async_client.AsyncClient`). + * *query_timeout* sets a default ``evaluationTimeout`` on every request. + * Session connections always attempt a rollback on ``close()`` to prevent + dangling transactions if the caller forgets to commit or roll back. + + Parameters + ---------- + url: + WebSocket endpoint, e.g. ``ws://localhost:8182/gremlin``. + traversal_source: + Gremlin traversal source name (default ``"g"``). + protocol_factory: + Forwarded to :class:`~gremlin_python.driver.async_client.AsyncClient`. + transport_factory: + Forwarded to :class:`~gremlin_python.driver.async_client.AsyncClient`. + pool_size: + Number of concurrent WebSocket connections (default 8). + message_serializer: + Serializer instance (default GraphBinarySerializersV1). + username, password: + Credentials for SASL/basic authentication. + kerberized_service: + Kerberos service principal. + headers: + Extra HTTP headers for every WebSocket handshake. + session: + Session UUID string to enable server-side sessions. + enable_user_agent_on_connect: + Send the driver user-agent header (default ``True``). + ping_interval: + Seconds between keepalive pings (default 60, 0 to disable). + query_timeout: + Default query timeout in **seconds** applied to every request as + ``evaluationTimeout``. Per-request options override this. + **transport_kwargs: + Extra keyword arguments forwarded to the default transport. + """ + + def __init__( + self, + url, + traversal_source="g", + protocol_factory=None, + transport_factory=None, + pool_size=None, + message_serializer=None, + username="", + password="", + kerberized_service="", + headers=None, + session=None, + enable_user_agent_on_connect=True, + ping_interval=60, + query_timeout=None, + **transport_kwargs, + ): + super().__init__(url, traversal_source) + self.__session = session + self.__spawned_sessions = [] + + default_request_options = {} + if query_timeout is not None: + default_request_options["evaluationTimeout"] = int(query_timeout * 1000) + + self._async_client = AsyncClient( + url, + traversal_source, + protocol_factory=protocol_factory, + transport_factory=transport_factory, + pool_size=pool_size, + 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, + ping_interval=ping_interval, + default_request_options=default_request_options or None, + **transport_kwargs, + ) + + # ------------------------------------------------------------------ + # RemoteConnection interface + # ------------------------------------------------------------------ + + def submit(self, bytecode): + """Not supported synchronously. + + :meth:`submit` is overridden as an *async* coroutine below. Python + allows an ``async def`` to satisfy an abstract ``def`` requirement, so + this docstring exists only to explain why the sync method is absent. + """ + raise RuntimeError( + "AsyncDriverRemoteConnection.submit() must be awaited. " + "Call it as 'await connection.submit(bytecode)' or use " + "'async_traversal().with_remote(rc)' for the DSL." + ) + + # ------------------------------------------------------------------ + # Async submit (called by AsyncRemoteStrategy) + # ------------------------------------------------------------------ + + async def submit(self, bytecode): # noqa: F811 — intentional async override + """Execute *bytecode* and return an :class:`~gremlin_python.driver.async_graph_traversal.AsyncRemoteTraversal`. + + Called automatically by :class:`~gremlin_python.driver.async_graph_traversal.AsyncRemoteStrategy` + when a terminal step (``to_list()``, ``next()``, …) is awaited on a + traversal. Application code rarely needs to call this directly. + """ + from gremlin_python.driver.async_graph_traversal import AsyncRemoteTraversal + options = self._extract_request_options(bytecode) + result_set = await self._async_client.submit(bytecode, request_options=options) + return AsyncRemoteTraversal(result_set) + + # ------------------------------------------------------------------ + # Convenience methods (application-level) + # ------------------------------------------------------------------ + + async def submit_async(self, traversal_or_bytecode, request_options=None): + """Execute a traversal or bytecode and return all results as a list. + + Results are automatically unwrapped from their ``Traverser`` wrappers. + Prefer the DSL (``g.V().to_list()``) for most use-cases; use this + method when you need direct bytecode submission. + + Parameters + ---------- + traversal_or_bytecode: + A :class:`~gremlin_python.process.graph_traversal.GraphTraversal` + (or :class:`~gremlin_python.driver.async_graph_traversal.AsyncGraphTraversal`) + built with the Gremlin DSL without calling a terminal step, or a + raw :class:`~gremlin_python.process.traversal.Bytecode` object. + request_options: + Optional dict of request options. When *None*, options are + extracted from any + :class:`~gremlin_python.process.strategies.OptionsStrategy` + attached to the traversal. + """ + remote_traversal = await self.submit( + getattr(traversal_or_bytecode, "bytecode", traversal_or_bytecode) + ) + return await remote_traversal.to_list() + + async def submit_stream(self, traversal_or_bytecode, request_options=None): + """Execute a traversal and return an :class:`~gremlin_python.driver.async_resultset.AsyncResultSet`. + + Use when you need to stream results one at a time (``async for``) + rather than collecting them into memory. + """ + bytecode = getattr(traversal_or_bytecode, "bytecode", traversal_or_bytecode) + options = request_options if request_options is not None else \ + self._extract_request_options(bytecode) + return await self._async_client.submit(bytecode, request_options=options) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def close(self): + """Close all connections, rolling back any open session first.""" + for session in self.__spawned_sessions: + await session.close() + self.__spawned_sessions.clear() + if self.__session is not None: + # Best-effort rollback; prevents dangling transactions if the + # caller forgot to commit or roll back before closing. + try: + await self.rollback() + except Exception: + pass + await asyncio.shield(self._async_client.close()) + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + await self.close() + + # ------------------------------------------------------------------ + # State + # ------------------------------------------------------------------ + + def is_closed(self): + return self._async_client.is_closed() + + def is_session_bound(self): + return self.__session is not None + + # ------------------------------------------------------------------ + # Session management + # ------------------------------------------------------------------ + + def create_session(self): + """Return a new session-bound :class:`AsyncDriverRemoteConnection`.""" + if self.is_session_bound(): + raise Exception( + "Connection is already bound to a session — nested sessions are not allowed" + ) + conn = AsyncDriverRemoteConnection( + self._url, + traversal_source=self._traversal_source, + session=uuid.uuid4(), + ping_interval=self._async_client._ping_interval, + ) + self.__spawned_sessions.append(conn) + return conn + + async def remove_session(self, session_based_connection): + """Close and untrack a previously spawned session connection.""" + await session_based_connection.close() + if session_based_connection in self.__spawned_sessions: + self.__spawned_sessions.remove(session_based_connection) + + # ------------------------------------------------------------------ + # Transactions + # ------------------------------------------------------------------ + + async def commit(self): + """Commit the current transaction.""" + result_set = await self._async_client.submit(Bytecode.GraphOp.commit()) + await result_set.all() + + async def rollback(self): + """Roll back the current transaction.""" + result_set = await self._async_client.submit(Bytecode.GraphOp.rollback()) + await result_set.all() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @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, + ) + if not options_strategy: + return None + allowed_keys = [ + "evaluationTimeout", "scriptEvaluationTimeout", "batchSize", + "requestId", "userAgent", "materializeProperties", + ] + return { + k: options_strategy[1].configuration[k] + for k in allowed_keys + if k in options_strategy[1].configuration + } diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_graph_traversal.py b/gremlin-python/src/main/python/gremlin_python/driver/async_graph_traversal.py new file mode 100644 index 00000000000..92d2ec9b645 --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_graph_traversal.py @@ -0,0 +1,390 @@ +# +# 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. +# +"""Async traversal layer for gremlin-python. + +Drop-in async replacement for the standard ``traversal()`` factory:: + + async with AsyncDriverRemoteConnection("ws://localhost:8182/gremlin", "g") as rc: + g = async_traversal().with_remote(rc) + + names = await g.V().has_label("person").values("name").to_list() + count = await g.V().count().next() + await g.add_v("person").property("name", "saturn").iterate() + + async with g.tx() as gtx: + await gtx.add_v("person").property("name", "pluto").iterate() + # commits on clean exit, rolls back on exception +""" +import asyncio +import logging + +from gremlin_python.process.graph_traversal import GraphTraversal, GraphTraversalSource +from gremlin_python.process.traversal import ( + Bytecode, Traversal, TraversalStrategies, TraversalStrategy, +) +from gremlin_python.structure.graph import Graph + +log = logging.getLogger("gremlinpython") + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +# --------------------------------------------------------------------------- +# Public factory +# --------------------------------------------------------------------------- + +def async_traversal(): + """Return an :class:`_AsyncAnonymousTraversalSource`. + + Use in place of :func:`~gremlin_python.process.anonymous_traversal.traversal` + when working with :class:`~gremlin_python.driver.async_driver_remote_connection.AsyncDriverRemoteConnection`. + """ + return _AsyncAnonymousTraversalSource() + + +class _AsyncAnonymousTraversalSource: + def with_remote(self, remote_connection): + """Bind *remote_connection* and return an :class:`AsyncGraphTraversalSource`.""" + return AsyncGraphTraversalSource( + Graph(), + AsyncTraversalStrategies(), + remote_connection=remote_connection, + ) + + +# --------------------------------------------------------------------------- +# Strategy pipeline (async) +# --------------------------------------------------------------------------- + +class AsyncTraversalStrategies(TraversalStrategies): + """Async-aware strategy pipeline. + + Calls ``await strategy.apply(traversal)`` for coroutine strategies and + ``strategy.apply(traversal)`` for plain sync strategies so both can coexist. + """ + + async def apply_strategies(self, traversal): # noqa: D102 + for strategy in self.traversal_strategies: + if asyncio.iscoroutinefunction(strategy.apply): + await strategy.apply(traversal) + else: + strategy.apply(traversal) + + +class AsyncRemoteStrategy(TraversalStrategy): + """Client-side strategy that submits bytecode to the remote server asynchronously. + + Registered with the ``py:AsyncRemoteStrategy`` fqcn so it is never sent + to the server as part of the serialized traversal. + """ + + def __init__(self, remote_connection): + super().__init__(fqcn="py:AsyncRemoteStrategy") + self.remote_connection = remote_connection + + async def apply(self, traversal): + if traversal.traversers is None: + remote_traversal = await self.remote_connection.submit(traversal.bytecode) + traversal.remote_results = remote_traversal + traversal.traversers = remote_traversal.traversers + + +# --------------------------------------------------------------------------- +# Traversal base (async terminal steps) +# --------------------------------------------------------------------------- + +class AsyncTraversal(Traversal): + """Base traversal class that replaces blocking terminal steps with coroutines. + + Concrete subclasses must be used with :class:`AsyncTraversalStrategies` so + that ``apply_strategies`` can be awaited inside ``__anext__``. + """ + + # Prevent accidental sync iteration. + def __next__(self): + raise NotImplementedError( + "AsyncTraversal does not support synchronous iteration. " + "Use 'async for', 'await to_list()', or 'await next()'." + ) + + # ------------------------------------------------------------------ + # Async iteration protocol + # ------------------------------------------------------------------ + + def __aiter__(self): + return self + + async def __anext__(self): + if self.traversers is None: + await self.traversal_strategies.apply_strategies(self) + if self.last_traverser is None: + # StopAsyncIteration propagates naturally when the result set is exhausted. + self.last_traverser = await self.traversers.__anext__() + obj = self.last_traverser.object + self.last_traverser.bulk -= 1 + if self.last_traverser.bulk <= 0: + self.last_traverser = None + return obj + + # ------------------------------------------------------------------ + # Async terminal steps + # ------------------------------------------------------------------ + + async def to_list(self): + """Execute the traversal and return all results as a :class:`list`.""" + results = [] + async for item in self: + results.append(item) + return results + + async def to_set(self): + """Execute the traversal and return all results as a :class:`set`.""" + return set(await self.to_list()) + + async def next(self, amount=None): + """Execute and return the first result, or the first *amount* results. + + Raises :class:`StopAsyncIteration` when *amount* is ``None`` and the + result set is empty. + """ + if amount is None: + return await self.__anext__() + results = [] + for _ in range(amount): + try: + results.append(await self.__anext__()) + except StopAsyncIteration: + break + return results + + async def iterate(self): + """Execute the traversal, discarding all results. + + Use for write operations (``add_v``, ``add_e``, ``property``, …) + where no return value is needed. + """ + self.bytecode.add_step("discard") + async for _ in self: + pass + return self + + async def has_next(self): + """Return ``True`` if the traversal produces at least one result.""" + try: + await self.__anext__() + return True + except StopAsyncIteration: + return False + + async def next_traverser(self): + """Return the next raw :class:`~gremlin_python.process.traversal.Traverser`.""" + if self.traversers is None: + await self.traversal_strategies.apply_strategies(self) + return await self.traversers.__anext__() + + +# --------------------------------------------------------------------------- +# Concrete traversal classes +# --------------------------------------------------------------------------- + +class AsyncGraphTraversal(GraphTraversal, AsyncTraversal): + """Async graph traversal. + + Inherits all Gremlin DSL steps from :class:`~gremlin_python.process.graph_traversal.GraphTraversal` + and async terminal steps from :class:`AsyncTraversal` via Python's MRO. + + MRO: AsyncGraphTraversal → GraphTraversal → AsyncTraversal → Traversal → object + + Because ``GraphTraversal`` does not define terminal steps directly (it + inherits them from ``Traversal``), Python finds ``AsyncTraversal``\'s + async versions first in the MRO before reaching ``Traversal``\'s sync + versions. + """ + + +class AsyncRemoteTraversal(AsyncTraversal): + """Holds the :class:`~gremlin_python.driver.async_resultset.AsyncResultSet` + returned by :meth:`~gremlin_python.driver.async_driver_remote_connection.AsyncDriverRemoteConnection.submit`. + + ``traversers`` is set directly to the result set; strategies are never + applied (there is nothing to apply on an already-executed traversal). + """ + + def __init__(self, traversers): + super().__init__(None, None, None) + self.traversers = traversers + + +# --------------------------------------------------------------------------- +# Traversal source +# --------------------------------------------------------------------------- + +class AsyncGraphTraversalSource(GraphTraversalSource): + """Async traversal source. + + All spawned traversals are :class:`AsyncGraphTraversal` instances. + Configuration steps (``with_bulk()``, ``with_path()``, …) return a new + :class:`AsyncGraphTraversalSource`. + + Use :func:`async_traversal` to create one:: + + g = async_traversal().with_remote(rc) + """ + + def __init__(self, graph, traversal_strategies, bytecode=None, remote_connection=None): + # Call super WITHOUT remote_connection to avoid RemoteStrategy being added. + super().__init__(graph, traversal_strategies, bytecode) + self.graph_traversal = AsyncGraphTraversal + self.remote_connection = remote_connection + if remote_connection is not None: + self.traversal_strategies.add_strategies([AsyncRemoteStrategy(remote_connection)]) + + def get_graph_traversal_source(self): + """Return a copy of this source (used by configuration steps).""" + # Strategies are copied — do NOT pass remote_connection to avoid + # registering AsyncRemoteStrategy a second time. + return AsyncGraphTraversalSource( + self.graph, + AsyncTraversalStrategies(self.traversal_strategies), + Bytecode(self.bytecode), + ) + + def get_graph_traversal(self): + """Return a fresh :class:`AsyncGraphTraversal` for this source.""" + return AsyncGraphTraversal( + self.graph, + self.traversal_strategies, + Bytecode(self.bytecode), + ) + + async def tx(self): + """Return an :class:`AsyncTransaction` for this connection. + + Raises if the connection is already session-bound. + """ + strategy = next( + (s for s in self.traversal_strategies.traversal_strategies + if isinstance(s, AsyncRemoteStrategy)), + None, + ) + if strategy is None: + raise Exception("No remote connection configured on this traversal source") + if strategy.remote_connection.is_session_bound(): + raise Exception( + "Cannot open a transaction on an already session-bound connection" + ) + return AsyncTransaction(self, strategy.remote_connection) + + +# --------------------------------------------------------------------------- +# Transaction +# --------------------------------------------------------------------------- + +class AsyncTransaction: + """Async transaction over a session-bound connection. + + Supports the async context-manager protocol:: + + async with await g.tx() as gtx: + await gtx.add_v("person").property("name", "pluto").iterate() + # commits on clean exit; rolls back on exception + + Or manage it explicitly:: + + tx = await g.tx() + gtx = await tx.begin() + try: + await gtx.add_v("person").property("name", "pluto").iterate() + await tx.commit() + except Exception: + await tx.rollback() + """ + + def __init__(self, g, remote_connection): + self._g = g + self._remote_connection = remote_connection + self._session_connection = None + self._gtx = None + self._is_open = False + self._mutex = asyncio.Lock() + + @property + def is_open(self): + return self._is_open + + @property + def gtx(self): + if self._gtx is None: + raise Exception("Transaction not started — call begin() first") + return self._gtx + + async def begin(self): + """Open a session and return a session-bound :class:`AsyncGraphTraversalSource`.""" + async with self._mutex: + if self._is_open: + raise Exception("Transaction is already open") + self._session_connection = self._remote_connection.create_session() + strategies = AsyncTraversalStrategies() + strategies.add_strategies([AsyncRemoteStrategy(self._session_connection)]) + self._gtx = AsyncGraphTraversalSource( + self._g.graph, + strategies, + ) + self._is_open = True + return self._gtx + + async def commit(self): + """Commit the transaction and close the session.""" + async with self._mutex: + if not self._is_open: + raise Exception("Transaction is not open") + await self._session_connection.commit() + await self._close_session() + + async def rollback(self): + """Roll back the transaction and close the session.""" + async with self._mutex: + if not self._is_open: + raise Exception("Transaction is not open") + await self._session_connection.rollback() + await self._close_session() + + async def close(self): + """Close the session without committing or rolling back.""" + async with self._mutex: + if self._is_open: + await self._close_session() + + async def _close_session(self): + self._is_open = False + await self._remote_connection.remove_session(self._session_connection) + self._session_connection = None + self._gtx = None + + # Context-manager support: `async with await g.tx() as gtx` + async def __aenter__(self): + return await self.begin() + + async def __aexit__(self, exc_type, *_): + if not self._is_open: + return + if exc_type is not None: + await self.rollback() + else: + await self.commit() diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_protocol.py b/gremlin-python/src/main/python/gremlin_python/driver/async_protocol.py new file mode 100644 index 00000000000..663fc8d8334 --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_protocol.py @@ -0,0 +1,135 @@ +# +# 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 asyncio +import base64 +import logging + +from gremlin_python.driver import request +from gremlin_python.driver.protocol import ( + ConfigurationError, + GremlinServerError, + GremlinServerWSProtocol, +) + +log = logging.getLogger("gremlinpython") + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + + +class AsyncGremlinServerWSProtocol(GremlinServerWSProtocol): + """Async version of :class:`~gremlin_python.driver.protocol.GremlinServerWSProtocol`. + + Both :meth:`write` and :meth:`data_received` are coroutines that must be + awaited. The transport assigned via :meth:`connection_made` must be an + :class:`~gremlin_python.driver.aiohttp.async_transport.AsyncAiohttpWSTransport` + (i.e. its ``write`` and ``read`` methods must be coroutines). + """ + + async def write(self, request_id, request_message): + """Serialize *request_message* and send it over the WebSocket.""" + message = self._message_serializer.serialize_message( + request_id, request_message + ) + await self._transport.write(message) + + async def data_received(self, message, results_dict): + """Deserialize a raw WebSocket frame and dispatch results. + + Deserialization is offloaded to the default executor so that the event + loop is not blocked by CPU-bound work. + + Returns the Gremlin Server response status code (200, 204, or 206). + Raises :class:`~gremlin_python.driver.protocol.GremlinServerError` for + error status codes. + """ + if message is None: + log.error("Received empty message from server.") + raise GremlinServerError( + { + "code": 500, + "message": "Server disconnected - please try to reconnect", + "attributes": {}, + } + ) + + loop = asyncio.get_running_loop() + message = await loop.run_in_executor( + None, self._message_serializer.deserialize_message, message + ) + + request_id = message["requestId"] + result_set = results_dict.get(request_id) + status_code = message["status"]["code"] + aggregate_to = message["result"]["meta"].get("aggregateTo", "list") + data = message["result"]["data"] + + if result_set is not None: + 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) + await self.write(request_id, request_message) + raw = await self._transport.read() + return await self.data_received(raw, results_dict) + + if status_code == 204: + if result_set is not None: + result_set.stream.put_nowait([]) + del results_dict[request_id] + return status_code + + if status_code in (200, 206): + if result_set is not None: + 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 + + log.error( + "\r\nReceived error message '%s'\r\n\r\nWith results dictionary '%s'", + str(message), + str(results_dict), + ) + if request_id in results_dict: + del results_dict[request_id] + raise GremlinServerError(message["status"]) diff --git a/gremlin-python/src/main/python/gremlin_python/driver/async_resultset.py b/gremlin-python/src/main/python/gremlin_python/driver/async_resultset.py new file mode 100644 index 00000000000..3693ccfd828 --- /dev/null +++ b/gremlin-python/src/main/python/gremlin_python/driver/async_resultset.py @@ -0,0 +1,148 @@ +# +# 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 asyncio +import queue + +from gremlin_python.driver.resultset import ResultSet + +__author__ = 'Apache TinkerPop (dev@tinkerpop.apache.org)' + +# Internal sentinel used to signal that the stream is exhausted. +_EXHAUSTED = object() + + +class AsyncResultSet(ResultSet): + """Result set backed by an :class:`asyncio.Queue` and an + :class:`asyncio.Future`. + + Consumption methods (:meth:`one`, :meth:`all`) are coroutines. The class + also supports ``async for`` iteration. + + The ``future`` attribute **must** be set to a resolved + :class:`asyncio.Future` or a running :class:`asyncio.Task` by the owning + :class:`~gremlin_python.driver.async_connection.AsyncConnection` before + any consumption method is called. + """ + + def __init__(self, stream, request_id): + # stream must be an asyncio.Queue + super().__init__(stream, request_id) + self._future = None + # Cache individual items already unpacked from a multi-item chunk so + # that successive one() calls each return exactly one result. + self._item_cache = queue.Queue() + + # ------------------------------------------------------------------ + # future property (asyncio.Future / asyncio.Task) + # ------------------------------------------------------------------ + + @property + def future(self): + return self._future + + @future.setter + def future(self, value): + self._future = value + + # ------------------------------------------------------------------ + # Disable the synchronous interface + # ------------------------------------------------------------------ + + def __iter__(self): + raise NotImplementedError( + "AsyncResultSet supports only async iteration; use 'async for'" + ) + + def __next__(self): + raise NotImplementedError( + "AsyncResultSet supports only async iteration; use 'async for'" + ) + + # ------------------------------------------------------------------ + # Async iteration protocol + # ------------------------------------------------------------------ + + def __aiter__(self): + return self + + async def __anext__(self): + result = await self.one() + if result is _EXHAUSTED: + raise StopAsyncIteration + return result + + # ------------------------------------------------------------------ + # Consumption + # ------------------------------------------------------------------ + + async def one(self): + """Return the next individual result or the ``_EXHAUSTED`` sentinel. + + Each server response chunk contains a list of items. ``one()`` + unpacks those lists so that each call returns exactly one item, + caching any extras for subsequent calls. + """ + if not self._item_cache.empty(): + return self._item_cache.get_nowait() + + chunk = await self._one_chunk() + if chunk is _EXHAUSTED: + return _EXHAUSTED + + if isinstance(chunk, list): + if not chunk: + # Empty list from a 204 No Content response. + return _EXHAUSTED + for item in chunk[1:]: + self._item_cache.put_nowait(item) + return chunk[0] + + return chunk + + async def _one_chunk(self): + """Return the next raw chunk (list) from the stream or ``_EXHAUSTED``. + + Polls the stream queue while the background receive task is still + running, yielding control for up to 100 ms between checks so that + other coroutines can make progress. + """ + while not self._future.done(): + if not self.stream.empty(): + return self.stream.get_nowait() + # Brief yield so other coroutines can run while we wait. + await asyncio.wait([self._future], timeout=0.1) + + # Future is done; drain whatever landed in the queue. + if not self.stream.empty(): + return await self.stream.get() + + # Re-raise if the background task failed; return sentinel on success. + self._future.result() + return _EXHAUSTED + + async def all(self): + """Collect every result and return them as a flat list.""" + results = [] + while True: + chunk = await self._one_chunk() + if chunk is _EXHAUSTED: + break + if isinstance(chunk, list): + results.extend(chunk) + return results diff --git a/gremlin-python/src/main/python/tests/unit/driver/__init__.py b/gremlin-python/src/main/python/tests/unit/driver/__init__.py new file mode 100644 index 00000000000..fe95886d5c1 --- /dev/null +++ b/gremlin-python/src/main/python/tests/unit/driver/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# diff --git a/gremlin-python/src/main/python/tests/unit/driver/test_async_driver_remote_connection.py b/gremlin-python/src/main/python/tests/unit/driver/test_async_driver_remote_connection.py new file mode 100644 index 00000000000..9d7d37ea89b --- /dev/null +++ b/gremlin-python/src/main/python/tests/unit/driver/test_async_driver_remote_connection.py @@ -0,0 +1,397 @@ +# +# 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. +# +"""Unit tests for AsyncDriverRemoteConnection. + +Architecture under test +----------------------- +* submit(bytecode) — async coroutine, called by AsyncRemoteStrategy, + returns AsyncRemoteTraversal +* submit_async(traversal) — convenience: calls submit() then .to_list() +* submit_stream(traversal) — returns raw AsyncResultSet for streaming +* Lifecycle: close(), context manager, is_closed(), sessions, transactions +* _extract_request_options — static helper + +All I/O is mocked; no Gremlin Server is required. +""" +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock + +from gremlin_python.driver.async_driver_remote_connection import AsyncDriverRemoteConnection +from gremlin_python.driver.async_graph_traversal import ( + AsyncRemoteTraversal, + async_traversal, +) +from gremlin_python.driver.async_resultset import AsyncResultSet +from gremlin_python.process.traversal import Bytecode, Traverser +from gremlin_python.process.strategies import OptionsStrategy + + +URL = "ws://localhost:8182/gremlin" +TRAVERSAL_SOURCE = "g" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_result_set(values): + q = asyncio.Queue() + if values: + q.put_nowait([Traverser(v) for v in values]) + rs = AsyncResultSet(q, "test-id") + fut = MagicMock() + fut.done.return_value = True + fut.result.return_value = None + rs.future = fut + return rs + + +def _patch_submit(rc, values=None): + """Replace rc.submit with an AsyncMock returning AsyncRemoteTraversal.""" + rs = _make_result_set(values or []) + rc.submit = AsyncMock(return_value=AsyncRemoteTraversal(rs)) + return rs + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + +class TestInit(unittest.TestCase): + def test_url_and_traversal_source_stored(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + self.assertEqual(rc.url, URL) + self.assertEqual(rc.traversal_source, TRAVERSAL_SOURCE) + + def test_default_traversal_source(self): + rc = AsyncDriverRemoteConnection(URL) + self.assertEqual(rc.traversal_source, "g") + + def test_not_closed_on_init(self): + rc = AsyncDriverRemoteConnection(URL) + self.assertFalse(rc.is_closed()) + + def test_not_session_bound_by_default(self): + rc = AsyncDriverRemoteConnection(URL) + self.assertFalse(rc.is_session_bound()) + + def test_session_bound_when_session_provided(self): + import uuid + rc = AsyncDriverRemoteConnection(URL, session=str(uuid.uuid4())) + self.assertTrue(rc.is_session_bound()) + + def test_query_timeout_sets_default_request_options(self): + rc = AsyncDriverRemoteConnection(URL, query_timeout=30) + self.assertEqual( + rc._async_client._default_request_options.get("evaluationTimeout"), + 30000, + ) + + def test_query_timeout_fractional_seconds(self): + rc = AsyncDriverRemoteConnection(URL, query_timeout=1.5) + self.assertEqual( + rc._async_client._default_request_options.get("evaluationTimeout"), + 1500, + ) + + +# --------------------------------------------------------------------------- +# submit() — called by AsyncRemoteStrategy +# --------------------------------------------------------------------------- + +class TestSubmit(unittest.IsolatedAsyncioTestCase): + async def test_submit_bytecode_returns_async_remote_traversal(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[6]) + bc = Bytecode() + bc.add_step("V") + bc.add_step("count") + result = await rc.submit(bc) + self.assertIsInstance(result, AsyncRemoteTraversal) + + async def test_submit_traversers_iterable(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=["marko", "vadas"]) + bc = Bytecode() + remote_traversal = await rc.submit(bc) + results = await remote_traversal.to_list() + self.assertEqual(results, ["marko", "vadas"]) + + +# --------------------------------------------------------------------------- +# submit_async() — convenience wrapper +# --------------------------------------------------------------------------- + +class TestSubmitAsync(unittest.IsolatedAsyncioTestCase): + async def test_returns_list(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[10]) + bc = Bytecode() + results = await rc.submit_async(bc) + self.assertEqual(results, [10]) + + async def test_accepts_traversal_object(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=["marko"]) + g = async_traversal().with_remote(rc) + t = g.V().has_label("person") + results = await rc.submit_async(t) + self.assertEqual(results, ["marko"]) + # Verify bytecode was extracted and passed + submitted_bc = rc.submit.call_args[0][0] + self.assertIsInstance(submitted_bc, Bytecode) + + async def test_empty_result(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[]) + results = await rc.submit_async(Bytecode()) + self.assertEqual(results, []) + + async def test_propagates_exception(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc.submit = AsyncMock(side_effect=RuntimeError("server error")) + with self.assertRaises(RuntimeError): + await rc.submit_async(Bytecode()) + + +# --------------------------------------------------------------------------- +# submit_stream() — returns AsyncResultSet +# --------------------------------------------------------------------------- + +class TestSubmitStream(unittest.IsolatedAsyncioTestCase): + async def test_returns_async_result_set(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + q = asyncio.Queue() + q.put_nowait([Traverser("x"), Traverser("y")]) + rs = AsyncResultSet(q, "req") + fut = asyncio.get_event_loop().create_future() + fut.set_result(None) + rs.future = fut + rc._async_client.submit = AsyncMock(return_value=rs) + + result = await rc.submit_stream(Bytecode()) + self.assertIsInstance(result, AsyncResultSet) + + async def test_stream_iterable(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + q = asyncio.Queue() + q.put_nowait([Traverser("a"), Traverser("b")]) + rs = AsyncResultSet(q, "req") + fut = asyncio.get_event_loop().create_future() + fut.set_result(None) + rs.future = fut + rc._async_client.submit = AsyncMock(return_value=rs) + + collected = [] + async for item in await rc.submit_stream(Bytecode()): + collected.append(item) + # AsyncResultSet yields Traverser objects (raw, not unwrapped) + self.assertEqual(len(collected), 2) + + +# --------------------------------------------------------------------------- +# DSL integration via async_traversal() +# --------------------------------------------------------------------------- + +class TestDSLIntegration(unittest.IsolatedAsyncioTestCase): + async def test_g_v_to_list(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[6]) + g = async_traversal().with_remote(rc) + results = await g.V().count().to_list() + self.assertEqual(results, [6]) + + async def test_g_v_next(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[42]) + g = async_traversal().with_remote(rc) + result = await g.V().count().next() + self.assertEqual(result, 42) + + async def test_g_add_v_iterate(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[]) + g = async_traversal().with_remote(rc) + await g.add_v("planet").property("name", "saturn").iterate() + rc.submit.assert_awaited_once() + + async def test_g_has_next_true(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=["marko"]) + g = async_traversal().with_remote(rc) + self.assertTrue(await g.V().has("name", "marko").has_next()) + + async def test_g_has_next_false(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + _patch_submit(rc, values=[]) + g = async_traversal().with_remote(rc) + self.assertFalse(await g.V().has("name", "nobody").has_next()) + + async def test_concurrent_gather(self): + call_results = [["marko", "vadas"], [6]] + call_index = 0 + + async def mock_submit(bc): + nonlocal call_index + rs = _make_result_set(call_results[call_index]) + call_index += 1 + return AsyncRemoteTraversal(rs) + + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc.submit = mock_submit + g = async_traversal().with_remote(rc) + + names, count = await asyncio.gather( + g.V().has_label("person").values("name").to_list(), + g.V().count().next(), + ) + self.assertEqual(names, ["marko", "vadas"]) + self.assertEqual(count, 6) + + +# --------------------------------------------------------------------------- +# Lifecycle +# --------------------------------------------------------------------------- + +class TestLifecycle(unittest.IsolatedAsyncioTestCase): + async def test_close_delegates_to_client(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc._async_client.close = AsyncMock() + await rc.close() + rc._async_client.close.assert_awaited_once() + + async def test_context_manager_calls_close(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc._async_client.close = AsyncMock() + async with rc: + pass + rc._async_client.close.assert_awaited_once() + + async def test_context_manager_calls_close_on_exception(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc._async_client.close = AsyncMock() + with self.assertRaises(ValueError): + async with rc: + raise ValueError("boom") + rc._async_client.close.assert_awaited_once() + + async def test_session_connection_rollback_on_close(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE, session="sid") + rc._async_client.close = AsyncMock() + rc.rollback = AsyncMock() + await rc.close() + rc.rollback.assert_awaited_once() + + async def test_is_closed_reflects_client_state(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rc._async_client._closed = True + self.assertTrue(rc.is_closed()) + rc._async_client._closed = False + self.assertFalse(rc.is_closed()) + + +# --------------------------------------------------------------------------- +# Transactions +# --------------------------------------------------------------------------- + +class TestTransactions(unittest.IsolatedAsyncioTestCase): + async def test_commit_calls_client_submit(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rs = _make_result_set([]) + rc._async_client.submit = AsyncMock(return_value=rs) + await rc.commit() + submitted_bc = rc._async_client.submit.call_args[0][0] + self.assertIsInstance(submitted_bc, Bytecode) + + async def test_rollback_calls_client_submit(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + rs = _make_result_set([]) + rc._async_client.submit = AsyncMock(return_value=rs) + await rc.rollback() + submitted_bc = rc._async_client.submit.call_args[0][0] + self.assertIsInstance(submitted_bc, Bytecode) + + +# --------------------------------------------------------------------------- +# Session management +# --------------------------------------------------------------------------- + +class TestSessionManagement(unittest.TestCase): + def test_create_session_returns_session_bound_connection(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + session_rc = rc.create_session() + self.assertIsInstance(session_rc, AsyncDriverRemoteConnection) + self.assertTrue(session_rc.is_session_bound()) + + def test_create_session_raises_if_already_session_bound(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE, session="sid") + with self.assertRaises(Exception): + rc.create_session() + + def test_create_session_child_is_different_instance(self): + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + session_rc = rc.create_session() + self.assertIsNot(rc, session_rc) + + def test_remove_session_untracks_connection(self): + import asyncio + rc = AsyncDriverRemoteConnection(URL, TRAVERSAL_SOURCE) + session_rc = rc.create_session() + # Patch close so it doesn't try to actually connect + session_rc._async_client.close = lambda: asyncio.coroutine(lambda: None)() + self.assertIn(session_rc, rc._AsyncDriverRemoteConnection__spawned_sessions) + + +# --------------------------------------------------------------------------- +# _extract_request_options +# --------------------------------------------------------------------------- + +class TestExtractRequestOptions(unittest.TestCase): + def test_no_options_strategy_returns_none(self): + bc = Bytecode() + self.assertIsNone(AsyncDriverRemoteConnection._extract_request_options(bc)) + + def test_options_strategy_extracted(self): + bc = Bytecode() + strategy = OptionsStrategy(evaluationTimeout=1000, batchSize=50) + bc.add_source("withStrategies", strategy) + opts = AsyncDriverRemoteConnection._extract_request_options(bc) + self.assertEqual(opts["evaluationTimeout"], 1000) + self.assertEqual(opts["batchSize"], 50) + + def test_unknown_keys_excluded(self): + bc = Bytecode() + strategy = OptionsStrategy(evaluationTimeout=500) + strategy.configuration["unknownKey"] = "value" + bc.add_source("withStrategies", strategy) + opts = AsyncDriverRemoteConnection._extract_request_options(bc) + self.assertNotIn("unknownKey", opts) + + def test_partial_keys(self): + bc = Bytecode() + strategy = OptionsStrategy(requestId="abc-123") + bc.add_source("withStrategies", strategy) + opts = AsyncDriverRemoteConnection._extract_request_options(bc) + self.assertEqual(opts["requestId"], "abc-123") + self.assertNotIn("evaluationTimeout", opts) + + +if __name__ == "__main__": + unittest.main() diff --git a/gremlin-python/src/main/python/tests/unit/driver/test_async_graph_traversal.py b/gremlin-python/src/main/python/tests/unit/driver/test_async_graph_traversal.py new file mode 100644 index 00000000000..6cd413851e8 --- /dev/null +++ b/gremlin-python/src/main/python/tests/unit/driver/test_async_graph_traversal.py @@ -0,0 +1,504 @@ +# +# 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. +# +"""Unit tests for async_graph_traversal. + +Architecture under test +----------------------- +* async_traversal() → _AsyncAnonymousTraversalSource +* .with_remote(rc) → AsyncGraphTraversalSource (wraps GraphTraversalSource) +* g.V() → AsyncGraphTraversal (GraphTraversal + AsyncTraversal via MRO) +* .to_list() → AsyncTraversal.to_list() (async, iterates via __anext__) +* __anext__ → applies AsyncTraversalStrategies → + AsyncRemoteStrategy.apply() → + await rc.submit(bytecode) → + AsyncRemoteTraversal(AsyncResultSet) → + traversal.traversers = AsyncResultSet + → AsyncResultSet.__anext__() → Traverser objects + → traverser.object (unwrapped) + +All I/O is mocked; no Gremlin Server is required. +""" +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock + +from gremlin_python.driver.async_graph_traversal import ( + AsyncGraphTraversal, + AsyncGraphTraversalSource, + AsyncRemoteTraversal, + AsyncRemoteStrategy, + AsyncTraversal, + AsyncTraversalStrategies, + AsyncTransaction, + async_traversal, +) +from gremlin_python.driver.async_driver_remote_connection import AsyncDriverRemoteConnection +from gremlin_python.driver.async_resultset import AsyncResultSet +from gremlin_python.process.traversal import Bytecode, Traverser + + +URL = "ws://localhost:8182/gremlin" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_async_result_set(values): + """AsyncResultSet pre-loaded with Traverser-wrapped *values*. + + Uses a MagicMock future (already done) so this helper works in both sync + and async test contexts without requiring a running event loop. + """ + q = asyncio.Queue() + if values: + q.put_nowait([Traverser(v) for v in values]) + rs = AsyncResultSet(q, "test-id") + fut = MagicMock() + fut.done.return_value = True + fut.result.return_value = None + rs.future = fut + return rs + + +def _make_rc(values=None): + """AsyncDriverRemoteConnection whose submit() returns an AsyncRemoteTraversal.""" + rc = AsyncDriverRemoteConnection(URL, "g") + rs = _make_async_result_set(values or []) + rc.submit = AsyncMock(return_value=AsyncRemoteTraversal(rs)) + return rc + + +def _make_g(rc): + return async_traversal().with_remote(rc) + + +# --------------------------------------------------------------------------- +# async_traversal() factory +# --------------------------------------------------------------------------- + +class TestAsyncTraversalFactory(unittest.TestCase): + def test_with_remote_returns_async_source(self): + rc = _make_rc() + g = async_traversal().with_remote(rc) + self.assertIsInstance(g, AsyncGraphTraversalSource) + + def test_source_v_returns_async_graph_traversal(self): + g = _make_g(_make_rc()) + self.assertIsInstance(g.V(), AsyncGraphTraversal) + + def test_source_e_returns_async_graph_traversal(self): + g = _make_g(_make_rc()) + self.assertIsInstance(g.E(), AsyncGraphTraversal) + + def test_source_add_v_returns_async_graph_traversal(self): + g = _make_g(_make_rc()) + self.assertIsInstance(g.add_v("person"), AsyncGraphTraversal) + + def test_chained_steps_preserve_type(self): + g = _make_g(_make_rc()) + t = g.V().has_label("person").values("name") + self.assertIsInstance(t, AsyncGraphTraversal) + + def test_chained_steps_return_same_instance(self): + g = _make_g(_make_rc()) + t = g.V() + t2 = t.has_label("person") + t3 = t2.values("name") + self.assertIs(t, t2) + self.assertIs(t, t3) + + +# --------------------------------------------------------------------------- +# Bytecode construction +# --------------------------------------------------------------------------- + +class TestBytecodeConstruction(unittest.TestCase): + def test_v_step_in_bytecode(self): + g = _make_g(_make_rc()) + t = g.V() + step_names = [s[0] for s in t.bytecode.step_instructions] + self.assertIn("V", step_names) + + def test_chained_steps_accumulate(self): + g = _make_g(_make_rc()) + t = g.V().has_label("person").values("name") + step_names = [s[0] for s in t.bytecode.step_instructions] + self.assertIn("V", step_names) + self.assertIn("hasLabel", step_names) + self.assertIn("values", step_names) + + def test_independent_traversals_separate_bytecodes(self): + g = _make_g(_make_rc()) + t1 = g.V().has_label("person") + t2 = g.V().has_label("software") + self.assertIsNot(t1.bytecode, t2.bytecode) + + +# --------------------------------------------------------------------------- +# MRO — async terminal steps shadow sync ones +# --------------------------------------------------------------------------- + +class TestMRO(unittest.TestCase): + def test_to_list_is_coroutine_function(self): + g = _make_g(_make_rc()) + import inspect + self.assertTrue(inspect.iscoroutinefunction(g.V().to_list)) + + def test_next_is_coroutine_function(self): + import inspect + g = _make_g(_make_rc()) + self.assertTrue(inspect.iscoroutinefunction(g.V().next)) + + def test_iterate_is_coroutine_function(self): + import inspect + g = _make_g(_make_rc()) + self.assertTrue(inspect.iscoroutinefunction(g.V().iterate)) + + def test_sync_next_raises(self): + g = _make_g(_make_rc()) + with self.assertRaises(NotImplementedError): + next(g.V()) + + +# --------------------------------------------------------------------------- +# AsyncTraversalStrategies +# --------------------------------------------------------------------------- + +class TestAsyncTraversalStrategies(unittest.IsolatedAsyncioTestCase): + async def test_applies_async_strategy(self): + called = [] + + class FakeStrategy: + async def apply(self, traversal): + called.append("async") + + strategies = AsyncTraversalStrategies() + strategies.add_strategies([FakeStrategy()]) + t = AsyncGraphTraversal(None, strategies, Bytecode()) + await strategies.apply_strategies(t) + self.assertEqual(called, ["async"]) + + async def test_applies_sync_strategy(self): + called = [] + + class FakeStrategy: + def apply(self, traversal): + called.append("sync") + + strategies = AsyncTraversalStrategies() + strategies.add_strategies([FakeStrategy()]) + t = AsyncGraphTraversal(None, strategies, Bytecode()) + await strategies.apply_strategies(t) + self.assertEqual(called, ["sync"]) + + +# --------------------------------------------------------------------------- +# AsyncRemoteStrategy +# --------------------------------------------------------------------------- + +class TestAsyncRemoteStrategy(unittest.IsolatedAsyncioTestCase): + async def test_apply_sets_traversers(self): + rc = _make_rc(values=["marko"]) + strategy = AsyncRemoteStrategy(rc) + + strategies = AsyncTraversalStrategies() + strategies.add_strategies([strategy]) + t = AsyncGraphTraversal(None, strategies, Bytecode()) + + self.assertIsNone(t.traversers) + await strategy.apply(t) + self.assertIsNotNone(t.traversers) + + async def test_apply_noop_if_traversers_already_set(self): + rc = _make_rc(values=["marko"]) + strategy = AsyncRemoteStrategy(rc) + + strategies = AsyncTraversalStrategies() + strategies.add_strategies([strategy]) + t = AsyncGraphTraversal(None, strategies, Bytecode()) + + fake_traversers = object() + t.traversers = fake_traversers + await strategy.apply(t) + # Should not overwrite existing traversers + self.assertIs(t.traversers, fake_traversers) + rc.submit.assert_not_called() + + +# --------------------------------------------------------------------------- +# Terminal steps — to_list +# --------------------------------------------------------------------------- + +class TestToList(unittest.IsolatedAsyncioTestCase): + async def test_to_list_returns_unwrapped_values(self): + rc = _make_rc(values=["marko", "vadas"]) + g = _make_g(rc) + result = await g.V().has_label("person").values("name").to_list() + self.assertEqual(result, ["marko", "vadas"]) + + async def test_to_list_empty(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + result = await g.V().has_label("missing").to_list() + self.assertEqual(result, []) + + async def test_to_list_submits_bytecode(self): + rc = _make_rc(values=[6]) + g = _make_g(rc) + await g.V().count().to_list() + rc.submit.assert_awaited_once() + submitted_bytecode = rc.submit.call_args[0][0] + self.assertIsInstance(submitted_bytecode, Bytecode) + + async def test_to_list_single_value(self): + rc = _make_rc(values=[42]) + g = _make_g(rc) + result = await g.V().count().to_list() + self.assertEqual(result, [42]) + + +# --------------------------------------------------------------------------- +# Terminal steps — to_set +# --------------------------------------------------------------------------- + +class TestToSet(unittest.IsolatedAsyncioTestCase): + async def test_to_set_deduplicates(self): + rc = _make_rc(values=["person", "person", "software"]) + g = _make_g(rc) + result = await g.V().label().to_set() + self.assertEqual(result, {"person", "software"}) + + async def test_to_set_returns_set_type(self): + rc = _make_rc(values=["a"]) + g = _make_g(rc) + result = await g.V().label().to_set() + self.assertIsInstance(result, set) + + +# --------------------------------------------------------------------------- +# Terminal steps — next +# --------------------------------------------------------------------------- + +class TestNext(unittest.IsolatedAsyncioTestCase): + async def test_next_returns_first(self): + rc = _make_rc(values=[42]) + g = _make_g(rc) + result = await g.V().count().next() + self.assertEqual(result, 42) + + async def test_next_with_amount(self): + rc = _make_rc(values=["marko", "vadas", "josh"]) + g = _make_g(rc) + result = await g.V().values("name").next(2) + self.assertEqual(result, ["marko", "vadas"]) + + async def test_next_empty_raises_stop_async_iteration(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + with self.assertRaises(StopAsyncIteration): + await g.V().has("name", "nobody").next() + + async def test_next_zero_returns_empty_list(self): + rc = _make_rc(values=["x"]) + g = _make_g(rc) + result = await g.V().values("name").next(0) + self.assertEqual(result, []) + + +# --------------------------------------------------------------------------- +# Terminal steps — iterate +# --------------------------------------------------------------------------- + +class TestIterate(unittest.IsolatedAsyncioTestCase): + async def test_iterate_returns_traversal(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + result = await g.add_v("person").property("name", "test").iterate() + self.assertIsInstance(result, AsyncGraphTraversal) + + async def test_iterate_adds_discard_step(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + t = g.add_v("person") + await t.iterate() + step_names = [s[0] for s in t.bytecode.step_instructions] + self.assertIn("discard", step_names) + + async def test_iterate_calls_submit(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + await g.add_v("planet").property("name", "saturn").iterate() + rc.submit.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Terminal steps — has_next +# --------------------------------------------------------------------------- + +class TestHasNext(unittest.IsolatedAsyncioTestCase): + async def test_has_next_true(self): + rc = _make_rc(values=["marko"]) + g = _make_g(rc) + result = await g.V().has("name", "marko").has_next() + self.assertTrue(result) + + async def test_has_next_false(self): + rc = _make_rc(values=[]) + g = _make_g(rc) + result = await g.V().has("name", "nobody").has_next() + self.assertFalse(result) + + +# --------------------------------------------------------------------------- +# Concurrent usage +# --------------------------------------------------------------------------- + +class TestConcurrent(unittest.IsolatedAsyncioTestCase): + async def test_gather_multiple_traversals(self): + results_queue = [["marko", "vadas"], [6], ["knows", "created"]] + call_index = 0 + + async def mock_submit(bc): + nonlocal call_index + rs = _make_async_result_set(results_queue[call_index]) + call_index += 1 + return AsyncRemoteTraversal(rs) + + rc = AsyncDriverRemoteConnection(URL, "g") + rc.submit = mock_submit + g = _make_g(rc) + + names, count, edges = await asyncio.gather( + g.V().has_label("person").values("name").to_list(), + g.V().count().next(), + g.E().label().to_list(), + ) + + self.assertEqual(names, ["marko", "vadas"]) + self.assertEqual(count, 6) + self.assertEqual(edges, ["knows", "created"]) + + +# --------------------------------------------------------------------------- +# AsyncGraphTraversalSource — get_graph_traversal_source +# --------------------------------------------------------------------------- + +class TestGetGraphTraversalSource(unittest.TestCase): + def test_returns_async_traversal_source(self): + g = _make_g(_make_rc()) + g2 = g.get_graph_traversal_source() + self.assertIsInstance(g2, AsyncGraphTraversalSource) + + def test_strategies_are_async(self): + g = _make_g(_make_rc()) + g2 = g.get_graph_traversal_source() + self.assertIsInstance(g2.traversal_strategies, AsyncTraversalStrategies) + + def test_async_remote_strategy_not_duplicated(self): + rc = _make_rc() + g = _make_g(rc) + g2 = g.get_graph_traversal_source() + remote_strategies = [ + s for s in g2.traversal_strategies.traversal_strategies + if isinstance(s, AsyncRemoteStrategy) + ] + self.assertEqual(len(remote_strategies), 1) + + +# --------------------------------------------------------------------------- +# AsyncTransaction +# --------------------------------------------------------------------------- + +class TestAsyncTransaction(unittest.IsolatedAsyncioTestCase): + def _make_session_rc(self): + """A parent RC that returns a session child on create_session().""" + parent_rc = AsyncDriverRemoteConnection(URL, "g") + parent_rc.submit = AsyncMock(return_value=AsyncRemoteTraversal( + _make_async_result_set([]))) + + session_rc = AsyncDriverRemoteConnection(URL, "g", session="fake-session") + session_rc.submit = AsyncMock(return_value=AsyncRemoteTraversal( + _make_async_result_set([]))) + session_rc.commit = AsyncMock() + session_rc.rollback = AsyncMock() + + parent_rc.create_session = lambda: session_rc + parent_rc.remove_session = AsyncMock() + return parent_rc, session_rc + + async def test_begin_returns_async_graph_traversal_source(self): + parent_rc, _ = self._make_session_rc() + g = _make_g(parent_rc) + tx = await g.tx() + gtx = await tx.begin() + self.assertIsInstance(gtx, AsyncGraphTraversalSource) + + async def test_begin_sets_is_open(self): + parent_rc, _ = self._make_session_rc() + g = _make_g(parent_rc) + tx = await g.tx() + self.assertFalse(tx.is_open) + await tx.begin() + self.assertTrue(tx.is_open) + + async def test_commit_closes_session(self): + parent_rc, session_rc = self._make_session_rc() + g = _make_g(parent_rc) + tx = await g.tx() + await tx.begin() + await tx.commit() + session_rc.commit.assert_awaited_once() + self.assertFalse(tx.is_open) + + async def test_rollback_closes_session(self): + parent_rc, session_rc = self._make_session_rc() + g = _make_g(parent_rc) + tx = await g.tx() + await tx.begin() + await tx.rollback() + session_rc.rollback.assert_awaited_once() + self.assertFalse(tx.is_open) + + async def test_context_manager_commits_on_success(self): + parent_rc, session_rc = self._make_session_rc() + g = _make_g(parent_rc) + async with await g.tx(): + pass + session_rc.commit.assert_awaited_once() + + async def test_context_manager_rolls_back_on_exception(self): + parent_rc, session_rc = self._make_session_rc() + g = _make_g(parent_rc) + with self.assertRaises(ValueError): + async with await g.tx(): + raise ValueError("boom") + session_rc.rollback.assert_awaited_once() + + async def test_double_begin_raises(self): + parent_rc, _ = self._make_session_rc() + g = _make_g(parent_rc) + tx = await g.tx() + await tx.begin() + with self.assertRaises(Exception, msg="already open"): + await tx.begin() + + +if __name__ == "__main__": + unittest.main() diff --git a/gremlin-python/src/main/python/tests/unit/driver/test_async_websocket.py b/gremlin-python/src/main/python/tests/unit/driver/test_async_websocket.py new file mode 100644 index 00000000000..dcce2e19679 --- /dev/null +++ b/gremlin-python/src/main/python/tests/unit/driver/test_async_websocket.py @@ -0,0 +1,638 @@ +# +# 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. +# +"""Unit tests for the native-async WebSocket driver components. + +All tests use :class:`unittest.IsolatedAsyncioTestCase` so no extra +test dependencies are required. Network I/O is replaced with +:mod:`unittest.mock` mocks throughout. +""" +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +import aiohttp + +from gremlin_python.driver.aiohttp.async_transport import AsyncAiohttpWSTransport +from gremlin_python.driver.async_protocol import AsyncGremlinServerWSProtocol +from gremlin_python.driver.async_resultset import AsyncResultSet, _EXHAUSTED +from gremlin_python.driver.async_connection import AsyncConnection +from gremlin_python.driver.async_client import AsyncClient +from gremlin_python.driver.protocol import GremlinServerError + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_ws_message(msg_type, data=b"payload"): + msg = MagicMock() + msg.type = msg_type + msg.data = data + return msg + + +def _make_done_future(result=None, *, loop=None): + """Return an already-resolved asyncio.Future.""" + fut = asyncio.get_event_loop().create_future() if loop is None else loop.create_future() + fut.set_result(result) + return fut + + +# --------------------------------------------------------------------------- +# AsyncAiohttpWSTransport +# --------------------------------------------------------------------------- + +class TestAsyncAiohttpWSTransport(unittest.IsolatedAsyncioTestCase): + + async def _make_transport(self, **kwargs): + t = AsyncAiohttpWSTransport(**kwargs) + return t + + # ---- connect ----------------------------------------------------------- + + async def test_connect_opens_session_and_websocket(self): + transport = await self._make_transport() + mock_ws = AsyncMock() + mock_session = MagicMock() + mock_session.ws_connect = AsyncMock(return_value=mock_ws) + mock_session.closed = False + + with patch("aiohttp.ClientSession", return_value=mock_session): + await transport.connect("ws://localhost:8182/gremlin") + + mock_session.ws_connect.assert_awaited_once_with( + "ws://localhost:8182/gremlin", headers=None + ) + self.assertIs(transport._websocket, mock_ws) + + async def test_connect_403_raises_friendly_message(self): + transport = await self._make_transport() + err = aiohttp.ClientResponseError(request_info=MagicMock(), history=()) + err.status = 403 + + mock_session = MagicMock() + mock_session.ws_connect = AsyncMock(side_effect=err) + mock_session.closed = False + + with patch("aiohttp.ClientSession", return_value=mock_session): + with self.assertRaises(Exception) as ctx: + await transport.connect("ws://localhost:8182/gremlin") + self.assertIn("403", str(ctx.exception)) + + async def test_connect_non_403_reraises(self): + transport = await self._make_transport() + err = aiohttp.ClientResponseError(request_info=MagicMock(), history=()) + err.status = 500 + + mock_session = MagicMock() + mock_session.ws_connect = AsyncMock(side_effect=err) + mock_session.closed = False + + with patch("aiohttp.ClientSession", return_value=mock_session): + with self.assertRaises(aiohttp.ClientResponseError): + await transport.connect("ws://localhost:8182/gremlin") + + async def test_connect_forwards_extra_kwargs(self): + transport = await self._make_transport(max_msg_size=1024) + mock_ws = AsyncMock() + mock_session = MagicMock() + mock_session.ws_connect = AsyncMock(return_value=mock_ws) + mock_session.closed = False + + with patch("aiohttp.ClientSession", return_value=mock_session): + await transport.connect("ws://localhost:8182/gremlin", headers={"X-Foo": "bar"}) + + mock_session.ws_connect.assert_awaited_once_with( + "ws://localhost:8182/gremlin", + headers={"X-Foo": "bar"}, + max_msg_size=1024, + ) + + async def test_connect_remaps_max_content_length(self): + transport = AsyncAiohttpWSTransport(max_content_length=512) + self.assertIn("max_msg_size", transport._aiohttp_kwargs) + self.assertNotIn("max_content_length", transport._aiohttp_kwargs) + self.assertEqual(transport._aiohttp_kwargs["max_msg_size"], 512) + + async def test_connect_remaps_ssl_options(self): + ssl_ctx = MagicMock() + transport = AsyncAiohttpWSTransport(ssl_options=ssl_ctx) + self.assertIn("ssl", transport._aiohttp_kwargs) + self.assertNotIn("ssl_options", transport._aiohttp_kwargs) + self.assertIs(transport._aiohttp_kwargs["ssl"], ssl_ctx) + + # ---- write ------------------------------------------------------------- + + async def test_write_sends_bytes(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + await transport.write(b"hello") + transport._websocket.send_bytes.assert_awaited_once_with(b"hello") + + # ---- read -------------------------------------------------------------- + + async def test_read_binary_returns_bytes(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + transport._websocket.receive = AsyncMock( + return_value=_make_ws_message(aiohttp.WSMsgType.binary, b"\x01\x02") + ) + result = await transport.read() + self.assertEqual(result, b"\x01\x02") + + async def test_read_text_encodes_to_utf8(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + msg = _make_ws_message(aiohttp.WSMsgType.text, "hello world") + transport._websocket.receive = AsyncMock(return_value=msg) + result = await transport.read() + self.assertEqual(result, b"hello world") + + async def test_read_close_raises_runtime_error(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + transport._websocket.closed = False + transport._client_session = AsyncMock() + transport._client_session.closed = False + transport._websocket.receive = AsyncMock( + return_value=_make_ws_message(aiohttp.WSMsgType.close) + ) + with self.assertRaises(RuntimeError) as ctx: + await transport.read() + self.assertIn("closed by server", str(ctx.exception)) + + async def test_read_error_raises_runtime_error(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + transport._websocket.receive = AsyncMock( + return_value=_make_ws_message(aiohttp.WSMsgType.error, "boom") + ) + with self.assertRaises(RuntimeError) as ctx: + await transport.read() + self.assertIn("boom", str(ctx.exception)) + + # ---- ping -------------------------------------------------------------- + + async def test_ping_calls_websocket_ping(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + await transport.ping() + transport._websocket.ping.assert_awaited_once() + + # ---- close ------------------------------------------------------------- + + async def test_close_closes_websocket_and_session(self): + transport = await self._make_transport() + transport._websocket = AsyncMock() + transport._websocket.closed = False + transport._client_session = AsyncMock() + transport._client_session.closed = False + await transport.close() + transport._websocket.close.assert_awaited_once() + transport._client_session.close.assert_awaited_once() + + async def test_close_skips_already_closed(self): + transport = await self._make_transport() + transport._websocket = MagicMock() + transport._websocket.closed = True + transport._client_session = MagicMock() + transport._client_session.closed = True + # Should not raise + await transport.close() + + # ---- closed property --------------------------------------------------- + + async def test_closed_true_when_websocket_none(self): + transport = await self._make_transport() + self.assertTrue(transport.closed) + + async def test_closed_false_when_connected(self): + transport = await self._make_transport() + transport._websocket = MagicMock(closed=False) + transport._client_session = MagicMock(closed=False) + self.assertFalse(transport.closed) + + +# --------------------------------------------------------------------------- +# AsyncGremlinServerWSProtocol +# --------------------------------------------------------------------------- + +class TestAsyncGremlinServerWSProtocol(unittest.IsolatedAsyncioTestCase): + + def _make_protocol(self): + ser = MagicMock() + ser.serialize_message = MagicMock(return_value=b"serialized") + proto = AsyncGremlinServerWSProtocol(ser) + proto._transport = AsyncMock() + return proto, ser + + def _make_result_set(self): + rs = AsyncResultSet(asyncio.Queue(), "req-1") + rs.future = _make_done_future() + return rs + + async def test_write_serializes_and_sends(self): + proto, ser = self._make_protocol() + msg = MagicMock() + await proto.write("req-1", msg) + ser.serialize_message.assert_called_once_with("req-1", msg) + proto._transport.write.assert_awaited_once_with(b"serialized") + + async def test_data_received_200_puts_data_in_stream(self): + proto, ser = self._make_protocol() + rs = self._make_result_set() + results_dict = {"req-1": rs} + server_msg = { + "requestId": "req-1", + "status": {"code": 200, "attributes": {"x": 1}}, + "result": {"data": [1, 2, 3], "meta": {}}, + } + ser.deserialize_message = MagicMock(return_value=server_msg) + + status = await proto.data_received(b"raw", results_dict) + + self.assertEqual(status, 200) + self.assertEqual(rs.stream.get_nowait(), [1, 2, 3]) + self.assertNotIn("req-1", results_dict) + self.assertEqual(rs.status_attributes, {"x": 1}) + + async def test_data_received_204_puts_empty_list(self): + proto, ser = self._make_protocol() + rs = self._make_result_set() + results_dict = {"req-1": rs} + server_msg = { + "requestId": "req-1", + "status": {"code": 204, "attributes": {}}, + "result": {"data": None, "meta": {}}, + } + ser.deserialize_message = MagicMock(return_value=server_msg) + + status = await proto.data_received(b"raw", results_dict) + + self.assertEqual(status, 204) + self.assertEqual(rs.stream.get_nowait(), []) + self.assertNotIn("req-1", results_dict) + + async def test_data_received_206_keeps_entry_in_results_dict(self): + proto, ser = self._make_protocol() + rs = self._make_result_set() + results_dict = {"req-1": rs} + server_msg = { + "requestId": "req-1", + "status": {"code": 206, "attributes": {}}, + "result": {"data": [4, 5], "meta": {}}, + } + ser.deserialize_message = MagicMock(return_value=server_msg) + + status = await proto.data_received(b"raw", results_dict) + + self.assertEqual(status, 206) + self.assertIn("req-1", results_dict) + self.assertEqual(rs.stream.get_nowait(), [4, 5]) + + async def test_data_received_error_raises(self): + proto, ser = self._make_protocol() + rs = self._make_result_set() + results_dict = {"req-1": rs} + server_msg = { + "requestId": "req-1", + "status": {"code": 500, "message": "Internal error", "attributes": {}}, + "result": {"data": None, "meta": {}}, + } + ser.deserialize_message = MagicMock(return_value=server_msg) + + with self.assertRaises(GremlinServerError): + await proto.data_received(b"raw", results_dict) + + async def test_data_received_none_raises(self): + proto, _ = self._make_protocol() + with self.assertRaises(GremlinServerError): + await proto.data_received(None, {}) + + +# --------------------------------------------------------------------------- +# AsyncResultSet +# --------------------------------------------------------------------------- + +class TestAsyncResultSet(unittest.IsolatedAsyncioTestCase): + + def _make_rs(self, items=None, *, done=True): + q = asyncio.Queue() + if items: + q.put_nowait(items) + rs = AsyncResultSet(q, "req-1") + fut = asyncio.get_event_loop().create_future() + if done: + fut.set_result(None) + rs.future = fut + return rs + + async def test_one_returns_individual_items(self): + rs = self._make_rs([10, 20, 30]) + self.assertEqual(await rs.one(), 10) + self.assertEqual(await rs.one(), 20) + self.assertEqual(await rs.one(), 30) + result = await rs.one() + self.assertIs(result, _EXHAUSTED) + + async def test_one_empty_chunk_returns_exhausted(self): + rs = self._make_rs([]) + result = await rs.one() + self.assertIs(result, _EXHAUSTED) + + async def test_all_returns_flat_list(self): + q = asyncio.Queue() + q.put_nowait([1, 2]) + q.put_nowait([3, 4]) + rs = AsyncResultSet(q, "req-1") + fut = asyncio.get_event_loop().create_future() + fut.set_result(None) + rs.future = fut + self.assertEqual(await rs.all(), [1, 2, 3, 4]) + + async def test_all_empty_stream(self): + rs = self._make_rs(None) + self.assertEqual(await rs.all(), []) + + async def test_async_iteration(self): + rs = self._make_rs([100, 200]) + collected = [] + async for item in rs: + collected.append(item) + self.assertEqual(collected, [100, 200]) + + async def test_sync_iter_raises(self): + rs = self._make_rs([1]) + with self.assertRaises(NotImplementedError): + iter(rs) + + async def test_sync_next_raises(self): + rs = self._make_rs([1]) + with self.assertRaises(NotImplementedError): + next(rs) + + async def test_one_with_running_task(self): + """one() should wait for the background task to complete.""" + q = asyncio.Queue() + rs = AsyncResultSet(q, "req-1") + + async def _producer(): + await asyncio.sleep(0.05) + q.put_nowait([42]) + + task = asyncio.create_task(_producer()) + rs.future = task + + result = await rs.one() + self.assertEqual(result, 42) + + +# --------------------------------------------------------------------------- +# AsyncConnection +# --------------------------------------------------------------------------- + +class TestAsyncConnection(unittest.IsolatedAsyncioTestCase): + + def _make_conn(self): + pool = asyncio.Queue() + mock_transport = AsyncMock() + mock_transport.closed = False + mock_protocol = AsyncMock() + mock_protocol.connection_made = MagicMock() + + conn = AsyncConnection( + url="ws://localhost:8182/gremlin", + traversal_source="g", + protocol=mock_protocol, + transport_factory=lambda: mock_transport, + executor=None, + pool=pool, + ) + conn._transport = mock_transport + return conn, mock_transport, mock_protocol, pool + + async def test_connect_calls_transport_and_protocol(self): + conn, transport, protocol, _ = self._make_conn() + conn._transport = None + conn._inited = False + + await conn.connect() + + transport.connect.assert_awaited_once_with( + "ws://localhost:8182/gremlin", conn._headers + ) + protocol.connection_made.assert_called_once_with(transport) + self.assertTrue(conn._inited) + + async def test_close_calls_transport(self): + conn, transport, _, _ = self._make_conn() + conn._inited = True + await conn.close() + transport.close.assert_awaited_once() + + async def test_write_lazy_connects(self): + conn, transport, protocol, _ = self._make_conn() + conn._inited = False + + msg = MagicMock() + msg.args = {} + protocol.data_received = AsyncMock(return_value=200) + transport.read = AsyncMock(return_value=b"data") + + # Put an already-resolved future-bearing result set via the protocol mock + async def fake_data_received(_data, results): + req_id = list(results.keys())[0] + rs = results[req_id] + fut = asyncio.get_running_loop().create_future() + fut.set_result(None) + rs.future = fut + return 200 + + protocol.data_received = fake_data_received + + result_set = await conn.write(msg) + + self.assertTrue(conn._inited) + self.assertIsInstance(result_set, AsyncResultSet) + + async def test_receive_non_206_returns_conn_to_pool(self): + conn, transport, protocol, pool = self._make_conn() + conn._inited = True + + rs = AsyncResultSet(asyncio.Queue(), "req-1") + conn._results["req-1"] = rs + + async def fake_data_received(_data, results): + rs = results.get("req-1") + if rs: + fut = asyncio.get_running_loop().create_future() + fut.set_result(None) + rs.future = fut + return 200 + + protocol.data_received = fake_data_received + transport.read = AsyncMock(return_value=b"data") + + returned = await conn._receive(rs) + + self.assertIs(returned, rs) + self.assertEqual(pool.qsize(), 1) + + async def test_receive_exception_returns_conn_to_pool(self): + conn, transport, _, pool = self._make_conn() + conn._inited = True + + rs = AsyncResultSet(asyncio.Queue(), "req-1") + transport.read = AsyncMock(side_effect=RuntimeError("boom")) + + with self.assertRaises(RuntimeError): + await conn._receive(rs) + + self.assertEqual(pool.qsize(), 1) + self.assertFalse(conn._inited) + + +# --------------------------------------------------------------------------- +# AsyncClient +# --------------------------------------------------------------------------- + +class TestAsyncClient(unittest.IsolatedAsyncioTestCase): + + def _make_client(self, **kwargs): + # Patch the transport so no real network calls happen + mock_transport = AsyncMock() + mock_transport.closed = False + + client = AsyncClient( + "ws://localhost:8182/gremlin", + "g", + transport_factory=lambda: mock_transport, + **kwargs, + ) + return client, mock_transport + + async def test_pool_is_asyncio_queue(self): + client, _ = self._make_client() + self.assertIsInstance(client._pool, asyncio.Queue) + + async def test_pool_size_default(self): + client, _ = self._make_client() + self.assertEqual(client._pool_size, 8) + self.assertEqual(client._pool.qsize(), 8) + + async def test_pool_size_custom(self): + client, _ = self._make_client(pool_size=3) + self.assertEqual(client._pool.qsize(), 3) + + async def test_get_connection_returns_async_connection(self): + client, _ = self._make_client() + conn = client._get_connection() + self.assertIsInstance(conn, AsyncConnection) + + async def test_is_closed_initially_false(self): + client, _ = self._make_client() + self.assertFalse(client.is_closed()) + + async def test_context_manager_closes(self): + client, _ = self._make_client() + async with client: + pass + self.assertTrue(client.is_closed()) + + async def test_submit_async_raises_when_closed(self): + client, _ = self._make_client() + await client.close() + with self.assertRaises(Exception, msg="Client is closed"): + await client.submit_async("g.V()") + + async def test_submit_async_string_builds_eval_message(self): + client, _ = self._make_client() + mock_conn = AsyncMock() + mock_rs = MagicMock() + mock_conn.write = AsyncMock(return_value=mock_rs) + # Drain pool and put our mock connection + while not client._pool.empty(): + client._pool.get_nowait() + client._pool.put_nowait(mock_conn) + + result = await client.submit_async("g.V()") + + written_msg = mock_conn.write.call_args[0][0] + self.assertEqual(written_msg.op, "eval") + self.assertEqual(written_msg.processor, "") + self.assertEqual(written_msg.args["gremlin"], "g.V()") + self.assertIs(result, mock_rs) + + async def test_submit_async_bytecode_builds_bytecode_message(self): + from gremlin_python.process.traversal import Bytecode + client, _ = self._make_client() + mock_conn = AsyncMock() + mock_rs = MagicMock() + mock_conn.write = AsyncMock(return_value=mock_rs) + while not client._pool.empty(): + client._pool.get_nowait() + client._pool.put_nowait(mock_conn) + + bc = Bytecode() + bc.add_step("V") + await client.submit_async(bc) + + written_msg = mock_conn.write.call_args[0][0] + self.assertEqual(written_msg.op, "bytecode") + self.assertEqual(written_msg.processor, "traversal") + + async def test_close_drains_and_closes_connections(self): + client, _ = self._make_client(pool_size=2) + mock_conns = [] + while not client._pool.empty(): + mc = AsyncMock() + mock_conns.append(mc) + client._pool.get_nowait() + for mc in mock_conns: + client._pool.put_nowait(mc) + + await client.close() + + self.assertTrue(client.is_closed()) + for mc in mock_conns: + mc.close.assert_awaited_once() + + async def test_close_idempotent(self): + client, _ = self._make_client() + await client.close() + await client.close() # should not raise + self.assertTrue(client.is_closed()) + + async def test_session_forced_pool_size_1(self): + import uuid + client, _ = self._make_client(session=uuid.uuid4()) + self.assertEqual(client._pool_size, 1) + self.assertEqual(client._pool.qsize(), 1) + + async def test_session_pool_size_not_1_raises(self): + import uuid + with self.assertRaises(Exception): + AsyncClient( + "ws://localhost:8182/gremlin", "g", + session=uuid.uuid4(), pool_size=2 + ) + + +if __name__ == "__main__": + unittest.main()