Skip to content

Commit 5870e24

Browse files
committed
Fix tx modes
1 parent 4c0afa1 commit 5870e24

File tree

2 files changed

+27
-28
lines changed

2 files changed

+27
-28
lines changed

tests/test_connections.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,29 +30,24 @@ def _test_isolation_level_read_only(
3030
cursor = connection.cursor()
3131
with suppress(dbapi.DatabaseError):
3232
maybe_await(cursor.execute("DROP TABLE foo"))
33-
3433
cursor = connection.cursor()
3534
maybe_await(cursor.execute(
3635
"CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
3736
))
3837

3938
connection.set_isolation_level(isolation_level)
4039
cursor = connection.cursor()
41-
4240
query = "UPSERT INTO foo(id) VALUES (1)"
4341
if read_only:
4442
with pytest.raises(dbapi.DatabaseError):
4543
maybe_await(cursor.execute(query))
46-
4744
else:
4845
maybe_await(cursor.execute(query))
4946

5047
maybe_await(connection.rollback())
5148

5249
connection.set_isolation_level("AUTOCOMMIT")
53-
5450
cursor = connection.cursor()
55-
5651
maybe_await(cursor.execute("DROP TABLE foo"))
5752

5853
def _test_connection(self, connection: dbapi.Connection) -> None:
@@ -211,7 +206,9 @@ def connect() -> dbapi.AsyncConnection:
211206
try:
212207
yield conn
213208
finally:
214-
await greenlet_spawn(conn.close)
209+
def close() -> None:
210+
maybe_await(conn.close())
211+
await greenlet_spawn(close)
215212

216213
@pytest.mark.asyncio
217214
@pytest.mark.parametrize(

ydb_dbapi/connections.py

Lines changed: 24 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -30,26 +30,24 @@ class IsolationLevel(str, Enum):
3030

3131

3232
class _IsolationSettings(NamedTuple):
33-
ydb_mode: ydb.BaseQueryTxMode
33+
ydb_mode: ydb.BaseQueryTxMode | None
3434
interactive: bool
3535

3636

3737
_ydb_isolation_settings_map = {
38-
IsolationLevel.AUTOCOMMIT: _IsolationSettings(
39-
ydb.QuerySerializableReadWrite(), interactive=False
40-
),
38+
IsolationLevel.AUTOCOMMIT: _IsolationSettings(None, interactive=False),
4139
IsolationLevel.SERIALIZABLE: _IsolationSettings(
4240
ydb.QuerySerializableReadWrite(), interactive=True
4341
),
4442
IsolationLevel.ONLINE_READONLY: _IsolationSettings(
45-
ydb.QueryOnlineReadOnly(), interactive=True
43+
ydb.QueryOnlineReadOnly(), interactive=False
4644
),
4745
IsolationLevel.ONLINE_READONLY_INCONSISTENT: _IsolationSettings(
4846
ydb.QueryOnlineReadOnly().with_allow_inconsistent_reads(),
49-
interactive=True,
47+
interactive=False,
5048
),
5149
IsolationLevel.STALE_READONLY: _IsolationSettings(
52-
ydb.QueryStaleReadOnly(), interactive=True
50+
ydb.QueryStaleReadOnly(), interactive=False
5351
),
5452
IsolationLevel.SNAPSHOT_READONLY: _IsolationSettings(
5553
ydb.QuerySnapshotReadOnly(), interactive=True
@@ -78,10 +76,11 @@ def __init__(
7876

7977
self.connection_kwargs: dict = kwargs
8078

81-
self._tx_mode: ydb.BaseQueryTxMode = ydb.QuerySerializableReadWrite()
79+
self._shared_session_pool: bool = False
80+
8281
self._tx_context: TxContext | AsyncTxContext | None = None
82+
self._tx_mode: ydb.BaseQueryTxMode | None = None
8383
self.interactive_transaction: bool = False
84-
self._shared_session_pool: bool = False
8584

8685
if ydb_session_pool is not None:
8786
self._shared_session_pool = True
@@ -99,21 +98,24 @@ def __init__(
9998
self._session: ydb.QuerySession | ydb.aio.QuerySession | None = None
10099

101100
def set_isolation_level(self, isolation_level: IsolationLevel) -> None:
102-
ydb_isolation_settings = _ydb_isolation_settings_map[isolation_level]
103101
if self._tx_context and self._tx_context.tx_id:
104102
raise InternalError(
105103
"Failed to set transaction mode: transaction is already began"
106104
)
105+
106+
ydb_isolation_settings = _ydb_isolation_settings_map[isolation_level]
107+
108+
self._tx_context = None
107109
self._tx_mode = ydb_isolation_settings.ydb_mode
108110
self.interactive_transaction = ydb_isolation_settings.interactive
109111

110112
def get_isolation_level(self) -> str:
111-
if self._tx_mode.name == ydb.QuerySerializableReadWrite().name:
112-
if self.interactive_transaction:
113-
return IsolationLevel.SERIALIZABLE
113+
if self._tx_mode is None:
114114
return IsolationLevel.AUTOCOMMIT
115+
if self._tx_mode.name == ydb.QuerySerializableReadWrite().name:
116+
return IsolationLevel.SERIALIZABLE
115117
if self._tx_mode.name == ydb.QueryOnlineReadOnly().name:
116-
if self._tx_mode.settings.allow_inconsistent_reads:
118+
if self._tx_mode.allow_inconsistent_reads:
117119
return IsolationLevel.ONLINE_READONLY_INCONSISTENT
118120
return IsolationLevel.ONLINE_READONLY
119121
if self._tx_mode.name == ydb.QueryStaleReadOnly().name:
@@ -123,6 +125,12 @@ def get_isolation_level(self) -> str:
123125
msg = f"{self._tx_mode.name} is not supported"
124126
raise NotSupportedError(msg)
125127

128+
def _maybe_init_tx(
129+
self, session: ydb.QuerySession | ydb.aio.QuerySession
130+
) -> None:
131+
if self._tx_context is None and self._tx_mode is not None:
132+
self._tx_context = session.transaction(self._tx_mode)
133+
126134

127135
class Connection(BaseConnection):
128136
_driver_cls = ydb.Driver
@@ -154,10 +162,7 @@ def cursor(self) -> Cursor:
154162
if self._session is None:
155163
raise RuntimeError("Connection is not ready, use wait_ready.")
156164

157-
if self.interactive_transaction:
158-
self._tx_context = self._session.transaction(self._tx_mode)
159-
else:
160-
self._tx_context = None
165+
self._maybe_init_tx(self._session)
161166

162167
self._current_cursor = self._cursor_cls(
163168
session=self._session,
@@ -281,10 +286,7 @@ def cursor(self) -> AsyncCursor:
281286
if self._session is None:
282287
raise RuntimeError("Connection is not ready, use wait_ready.")
283288

284-
if self.interactive_transaction:
285-
self._tx_context = self._session.transaction(self._tx_mode)
286-
else:
287-
self._tx_context = None
289+
self._maybe_init_tx(self._session)
288290

289291
self._current_cursor = self._cursor_cls(
290292
session=self._session,

0 commit comments

Comments
 (0)