Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/firebolt/db/connection.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

from asyncio import new_event_loop
from functools import wraps
from inspect import cleandoc
from types import TracebackType
Expand Down Expand Up @@ -37,7 +38,7 @@ class Connection(AsyncBaseConnection):
"""
)

__slots__ = AsyncBaseConnection.__slots__ + ("_closing_lock",)
__slots__ = AsyncBaseConnection.__slots__ + ("_closing_lock", "_loop")

cursor_class = Cursor

Expand All @@ -46,6 +47,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
# Holding this lock for write means that connection is closing itself.
# cursor() should hold this lock for read to read/write state
self._closing_lock = RWLockWrite()
self._loop = new_event_loop()

@wraps(AsyncBaseConnection.cursor)
def cursor(self) -> Cursor:
Expand All @@ -57,7 +59,9 @@ def cursor(self) -> Cursor:
@wraps(AsyncBaseConnection._aclose)
def close(self) -> None:
with self._closing_lock.gen_wlock():
return async_to_sync(super()._aclose)()
if not self.closed:
self._loop.run_until_complete(self._aclose())
self._loop.close()

# Context manager support
def __enter__(self) -> Connection:
Expand Down