1
1
from __future__ import annotations
2
2
3
3
from collections .abc import AsyncGenerator
4
+ from collections .abc import Generator
4
5
from contextlib import suppress
5
6
6
7
import pytest
9
10
import ydb_dbapi as dbapi
10
11
11
12
12
- class BaseDBApiTestSuit :
13
- async def _test_isolation_level_read_only (
13
+ class BaseSyncDBApiTestSuit :
14
+ def _test_isolation_level_read_only (
14
15
self ,
15
16
connection : dbapi .Connection ,
16
17
isolation_level : str ,
17
18
read_only : bool ,
19
+ ) -> None :
20
+ connection .set_isolation_level ("AUTOCOMMIT" )
21
+ with connection .cursor () as cursor : # noqa: SIM117
22
+ with suppress (dbapi .DatabaseError ):
23
+ cursor .execute ("DROP TABLE foo" )
24
+
25
+ with connection .cursor () as cursor :
26
+ cursor .execute (
27
+ "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
28
+ )
29
+
30
+ connection .set_isolation_level (isolation_level )
31
+
32
+ with connection .cursor () as cursor :
33
+ query = "UPSERT INTO foo(id) VALUES (1)"
34
+ if read_only :
35
+ with pytest .raises (dbapi .DatabaseError ):
36
+ cursor .execute (query )
37
+ cursor .finish_query ()
38
+
39
+ else :
40
+ cursor .execute (query )
41
+
42
+ connection .rollback ()
43
+
44
+ connection .set_isolation_level ("AUTOCOMMIT" )
45
+
46
+ with connection .cursor () as cursor :
47
+ cursor .execute ("DROP TABLE foo" )
48
+
49
+ def _test_connection (self , connection : dbapi .Connection ) -> None :
50
+ connection .commit ()
51
+ connection .rollback ()
52
+
53
+ cur = connection .cursor ()
54
+ with suppress (dbapi .DatabaseError ):
55
+ cur .execute ("DROP TABLE foo" )
56
+ cur .finish_query ()
57
+
58
+ assert not connection .check_exists ("/local/foo" )
59
+ with pytest .raises (dbapi .ProgrammingError ):
60
+ connection .describe ("/local/foo" )
61
+
62
+ cur .execute ("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))" )
63
+ cur .finish_query ()
64
+
65
+ assert connection .check_exists ("/local/foo" )
66
+
67
+ col = (connection .describe ("/local/foo" )).columns [0 ]
68
+ assert col .name == "id"
69
+ assert col .type == ydb .PrimitiveType .Int64
70
+
71
+ cur .execute ("DROP TABLE foo" )
72
+ cur .close ()
73
+
74
+ def _test_cursor_raw_query (self , connection : dbapi .Connection ) -> None :
75
+ cur = connection .cursor ()
76
+ assert cur
77
+
78
+ with suppress (dbapi .DatabaseError ):
79
+ cur .execute ("DROP TABLE test" )
80
+ cur .finish_query ()
81
+
82
+ cur .execute (
83
+ "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
84
+ )
85
+ cur .finish_query ()
86
+
87
+ cur .execute (
88
+ """
89
+ DECLARE $data AS List<Struct<id:Int64, text: Utf8>>;
90
+
91
+ INSERT INTO test SELECT id, text FROM AS_TABLE($data);
92
+ """ ,
93
+ {
94
+ "$data" : ydb .TypedValue (
95
+ [
96
+ {"id" : 17 , "text" : "seventeen" },
97
+ {"id" : 21 , "text" : "twenty one" },
98
+ ],
99
+ ydb .ListType (
100
+ ydb .StructType ()
101
+ .add_member ("id" , ydb .PrimitiveType .Int64 )
102
+ .add_member ("text" , ydb .PrimitiveType .Utf8 )
103
+ ),
104
+ )
105
+ },
106
+ )
107
+ cur .finish_query ()
108
+
109
+ cur .execute ("DROP TABLE test" )
110
+
111
+ cur .close ()
112
+
113
+ def _test_errors (self , connection : dbapi .Connection ) -> None :
114
+ with pytest .raises (dbapi .InterfaceError ):
115
+ dbapi .connect (
116
+ "localhost:2136" , # type: ignore
117
+ database = "/local666" , # type: ignore
118
+ )
119
+
120
+ cur = connection .cursor ()
121
+
122
+ with suppress (dbapi .DatabaseError ):
123
+ cur .execute ("DROP TABLE test" )
124
+ cur .finish_query ()
125
+
126
+ with pytest .raises (dbapi .DataError ):
127
+ cur .execute ("SELECT 18446744073709551616" )
128
+
129
+ with pytest .raises (dbapi .DataError ):
130
+ cur .execute ("SELECT * FROM 拉屎" )
131
+
132
+ with pytest .raises (dbapi .DataError ):
133
+ cur .execute ("SELECT floor(5 / 2)" )
134
+
135
+ with pytest .raises (dbapi .ProgrammingError ):
136
+ cur .execute ("SELECT * FROM test" )
137
+
138
+ cur .execute ("CREATE TABLE test(id Int64, PRIMARY KEY (id))" )
139
+ cur .finish_query ()
140
+
141
+ cur .execute ("INSERT INTO test(id) VALUES(1)" )
142
+ cur .finish_query ()
143
+
144
+ with pytest .raises (dbapi .IntegrityError ):
145
+ cur .execute ("INSERT INTO test(id) VALUES(1)" )
146
+
147
+ cur .execute ("DROP TABLE test" )
148
+ cur .close ()
149
+
150
+
151
+ class TestConnection (BaseSyncDBApiTestSuit ):
152
+ @pytest .fixture
153
+ def connection (
154
+ self , connection_kwargs : dict
155
+ ) -> Generator [dbapi .Connection ]:
156
+ conn = dbapi .connect (** connection_kwargs ) # ignore: typing
157
+ try :
158
+ yield conn
159
+ finally :
160
+ conn .close ()
161
+
162
+ @pytest .mark .parametrize (
163
+ ("isolation_level" , "read_only" ),
164
+ [
165
+ (dbapi .IsolationLevel .SERIALIZABLE , False ),
166
+ (dbapi .IsolationLevel .AUTOCOMMIT , False ),
167
+ (dbapi .IsolationLevel .ONLINE_READONLY , True ),
168
+ (dbapi .IsolationLevel .ONLINE_READONLY_INCONSISTENT , True ),
169
+ (dbapi .IsolationLevel .STALE_READONLY , True ),
170
+ (dbapi .IsolationLevel .SNAPSHOT_READONLY , True ),
171
+ ],
172
+ )
173
+ def test_isolation_level_read_only (
174
+ self ,
175
+ isolation_level : str ,
176
+ read_only : bool ,
177
+ connection : dbapi .Connection ,
178
+ ) -> None :
179
+ self ._test_isolation_level_read_only (
180
+ connection , isolation_level , read_only
181
+ )
182
+
183
+ def test_connection (self , connection : dbapi .Connection ) -> None :
184
+ self ._test_connection (connection )
185
+
186
+ def test_cursor_raw_query (self , connection : dbapi .Connection ) -> None :
187
+ self ._test_cursor_raw_query (connection )
188
+
189
+ def test_errors (self , connection : dbapi .Connection ) -> None :
190
+ self ._test_errors (connection )
191
+
192
+
193
+ class BaseAsyncDBApiTestSuit :
194
+ async def _test_isolation_level_read_only (
195
+ self ,
196
+ connection : dbapi .AsyncConnection ,
197
+ isolation_level : str ,
198
+ read_only : bool ,
18
199
) -> None :
19
200
connection .set_isolation_level ("AUTOCOMMIT" )
20
201
async with connection .cursor () as cursor :
@@ -45,7 +226,9 @@ async def _test_isolation_level_read_only(
45
226
async with connection .cursor () as cursor :
46
227
await cursor .execute ("DROP TABLE foo" )
47
228
48
- async def _test_connection (self , connection : dbapi .Connection ) -> None :
229
+ async def _test_connection (
230
+ self , connection : dbapi .AsyncConnection
231
+ ) -> None :
49
232
await connection .commit ()
50
233
await connection .rollback ()
51
234
@@ -73,7 +256,7 @@ async def _test_connection(self, connection: dbapi.Connection) -> None:
73
256
await cur .close ()
74
257
75
258
async def _test_cursor_raw_query (
76
- self , connection : dbapi .Connection
259
+ self , connection : dbapi .AsyncConnection
77
260
) -> None :
78
261
cur = connection .cursor ()
79
262
assert cur
@@ -113,9 +296,9 @@ async def _test_cursor_raw_query(
113
296
114
297
await cur .close ()
115
298
116
- async def _test_errors (self , connection : dbapi .Connection ) -> None :
299
+ async def _test_errors (self , connection : dbapi .AsyncConnection ) -> None :
117
300
with pytest .raises (dbapi .InterfaceError ):
118
- await dbapi .connect (
301
+ await dbapi .async_connect (
119
302
"localhost:2136" , # type: ignore
120
303
database = "/local666" , # type: ignore
121
304
)
@@ -151,12 +334,12 @@ async def _test_errors(self, connection: dbapi.Connection) -> None:
151
334
await cur .close ()
152
335
153
336
154
- class TestAsyncConnection (BaseDBApiTestSuit ):
337
+ class TestAsyncConnection (BaseAsyncDBApiTestSuit ):
155
338
@pytest_asyncio .fixture
156
339
async def connection (
157
340
self , connection_kwargs : dict
158
- ) -> AsyncGenerator [dbapi .Connection ]:
159
- conn = await dbapi .connect (** connection_kwargs ) # ignore: typing
341
+ ) -> AsyncGenerator [dbapi .AsyncConnection ]:
342
+ conn = await dbapi .async_connect (** connection_kwargs ) # ignore: typing
160
343
try :
161
344
yield conn
162
345
finally :
@@ -178,22 +361,22 @@ async def test_isolation_level_read_only(
178
361
self ,
179
362
isolation_level : str ,
180
363
read_only : bool ,
181
- connection : dbapi .Connection ,
364
+ connection : dbapi .AsyncConnection ,
182
365
) -> None :
183
366
await self ._test_isolation_level_read_only (
184
367
connection , isolation_level , read_only
185
368
)
186
369
187
370
@pytest .mark .asyncio
188
- async def test_connection (self , connection : dbapi .Connection ) -> None :
371
+ async def test_connection (self , connection : dbapi .AsyncConnection ) -> None :
189
372
await self ._test_connection (connection )
190
373
191
374
@pytest .mark .asyncio
192
375
async def test_cursor_raw_query (
193
- self , connection : dbapi .Connection
376
+ self , connection : dbapi .AsyncConnection
194
377
) -> None :
195
378
await self ._test_cursor_raw_query (connection )
196
379
197
380
@pytest .mark .asyncio
198
- async def test_errors (self , connection : dbapi .Connection ) -> None :
381
+ async def test_errors (self , connection : dbapi .AsyncConnection ) -> None :
199
382
await self ._test_errors (connection )
0 commit comments