Skip to content

Commit d7a37dd

Browse files
robsdeduderobsdedude
authored
Fix application error desyncing auto-commit result (#1333)
Raising an application error inside a `with driver.session() as s:` block that has an on-going auto-commit transaction, i.e., `res = s.run(...)` has been called but `res` has not yet been consumed will lead to the session returning its backing connection to the pool (where it can be picked up by other sessions) without properly finalizing the auto-commit transaction. `res` then holds a reference to (and potentially uses) a connection that is no longer bound to `res`'s parent session. This can lead to all sorts of misbehavior (from data races casing protocol errors all the way to reading records of another query). Reproducer: ```python from __future__ import annotations from neo4j import GraphDatabase, Result from neo4j.debug import watch watch("neo4j") URL = "bolt://localhost:7687" AUTH = ("neo4j", "pass") DB = "neo4j" class MyError(RuntimeError): pass def inspect_records(result: neo4j.Result, n: int, stream_name: str): print(f"\ninspecting stream {stream_name}: fetching {n} records") records = [record["res"] for record in result.fetch(n)] print(f"got {len(records)} records") for record in records: print(repr(record)) assert isinstance(record, str) and record.startswith(f"{stream_name}-") def main(): range_top = 10 query = ( 'UNWIND range(1, $upper) AS i ' 'RETURN "res" + $n + "-" + apoc.text.lpad(toString(i), $pad, "0") ' 'AS res' ) params = {"upper": range_top, "pad": len(str(range_top))} with GraphDatabase.driver(URL, auth=AUTH) as driver: fetch_size = 2 try: with ( driver.session(database=DB, fetch_size=fetch_size) as session ): res1 = session.run(query, parameters={**params, "n": 1}) inspect_records(res1, 3, "res1") raise MyError("test error") except MyError as e: print(f"Caught MyError: {e}") with ( driver.session(database=DB, fetch_size=fetch_size) as session ): res2 = session.run(query, parameters={**params, "n": 2}) inspect_records(res2, 3, "res2") # !!! res1 now reading records from res2 inspect_records(res1, 3, "res1") if __name__ == "__main__": main() ```
1 parent c10e272 commit d7a37dd

4 files changed

Lines changed: 146 additions & 46 deletions

File tree

‎src/neo4j/_async/work/session.py‎

Lines changed: 25 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,6 @@ class AsyncSession(AsyncWorkspace):
9292
# The current auto-commit transaction result, if any.
9393
_auto_result = None
9494

95-
# The state this session is in.
96-
_state_failed = False
97-
9895
_config: SessionConfig
9996
_bookmark_manager: Bookmarks | None
10097
_pipelined_begin: ContextBool
@@ -113,12 +110,12 @@ async def __aenter__(self) -> AsyncSession:
113110
return self
114111

115112
async def __aexit__(self, exception_type, exception_value, traceback):
116-
if exception_type:
117-
if issubclass(exception_type, asyncio.CancelledError):
118-
self._handle_cancellation(message="__aexit__")
119-
self._closed = True
120-
return
121-
self._state_failed = True
113+
if exception_type and issubclass(
114+
exception_type, asyncio.CancelledError
115+
):
116+
self._handle_cancellation(message="__aexit__")
117+
self._closed = True
118+
return
122119
await self.close()
123120

124121
async def _connect(self, access_mode, **acquire_kwargs):
@@ -195,14 +192,21 @@ async def close(self) -> None:
195192
if self._closed:
196193
return
197194
if self._connection:
198-
if self._auto_result and self._state_failed is False:
195+
if self._auto_result:
199196
try:
197+
# The user didn't consume the auto-commit result.
198+
# Therefore, we try to consume it on a best-effort basis
199+
# before closing the parent session.
200200
await self._auto_result.consume()
201-
await self._update_bookmark(self._auto_result._bookmark)
202-
except Exception:
203-
# TODO: Investigate potential non graceful close states
204-
self._auto_result = None
205-
self._state_failed = True
201+
except (Neo4jError, DriverError) as e:
202+
# Consume failed, we'll swallow the error.
203+
# Best practice is to consume the result before closing the
204+
# session.
205+
log.warning(
206+
"Failed to consume outstanding auto-commit "
207+
"transaction result before closing session: %s",
208+
e,
209+
)
206210

207211
if self._transaction:
208212
if self._transaction._closed() is False:
@@ -215,18 +219,16 @@ async def close(self) -> None:
215219
await self._connection.send_all()
216220
await self._connection.fetch_all()
217221
# TODO: Investigate potential non graceful close states
218-
except Neo4jError:
219-
pass
220-
except TransactionError:
221-
pass
222-
except ServiceUnavailable:
223-
pass
224-
except SessionExpired:
222+
except (
223+
Neo4jError,
224+
TransactionError,
225+
ServiceUnavailable,
226+
SessionExpired,
227+
):
225228
pass
226229
finally:
227230
await self._disconnect()
228231

229-
self._state_failed = False
230232
self._closed = True
231233

232234
if AsyncUtil.is_async_code:

‎src/neo4j/_sync/work/session.py‎

Lines changed: 25 additions & 23 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tests/unit/async_/work/test_session.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@
4848
from ...._async_compat import mark_async_test
4949

5050

51+
class ApplicationError(RuntimeError):
52+
pass
53+
54+
5155
@mark_async_test
5256
async def test_session_context_calls_close(mocker):
5357
s = AsyncSession(None, SessionConfig())
@@ -1025,3 +1029,47 @@ async def test_session_run_idempotent_retry(
10251029
assert call_args[0] == "RETURN 1"
10261030
connection_1.telemetry.assert_called_once()
10271031
connection_2.telemetry.assert_not_called()
1032+
1033+
1034+
@mark_async_test
1035+
async def test_auto_commit_is_consumed_on_application_error(
1036+
async_fake_pool,
1037+
async_scripted_connection_generator,
1038+
) -> None:
1039+
bookmark = "res:bm1"
1040+
conn = async_scripted_connection_generator()
1041+
conn.set_script(
1042+
(
1043+
("run", {"on_success": ({"fields": ["n"]},), "on_summary": None}),
1044+
(
1045+
"pull",
1046+
{
1047+
"on_records": ([[1]],),
1048+
"on_success": ({"has_more": True},),
1049+
"on_summary": None,
1050+
},
1051+
),
1052+
(
1053+
"discard",
1054+
{
1055+
"on_success": ({"bookmark": bookmark},),
1056+
"on_summary": None,
1057+
},
1058+
),
1059+
)
1060+
)
1061+
async_fake_pool.buffered_connection_mocks.append(conn)
1062+
async_fake_pool.pool_config.telemetry_disabled = True
1063+
1064+
with pytest.raises(ApplicationError):
1065+
async with AsyncSession(async_fake_pool, SessionConfig()) as session:
1066+
res = await session.run("RETURN 1")
1067+
raise ApplicationError("boom")
1068+
1069+
assert not res._attached
1070+
assert session._auto_result is None
1071+
assert res._bookmark == bookmark
1072+
assert set(session._bookmarks) == {bookmark}
1073+
1074+
received_bookmarks = await session.last_bookmarks()
1075+
assert received_bookmarks.raw_values == {bookmark}

‎tests/unit/sync/work/test_session.py‎

Lines changed: 48 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)