1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
|
# Copyright (C) 2016-present the asyncpg authors and contributors
# <see AUTHORS file>
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
import asyncio
import os
import platform
import unittest
from asyncpg import _testbase as tb
from asyncpg import exceptions
class TestListeners(tb.ClusterTestCase):
async def test_listen_01(self):
async with self.create_pool(database='postgres') as pool:
async with pool.acquire() as con:
q1 = asyncio.Queue()
q2 = asyncio.Queue()
q3 = asyncio.Queue()
def listener1(*args):
q1.put_nowait(args)
def listener2(*args):
q2.put_nowait(args)
async def async_listener3(*args):
q3.put_nowait(args)
await con.add_listener('test', listener1)
await con.add_listener('test', listener2)
await con.add_listener('test', async_listener3)
await con.execute("NOTIFY test, 'aaaa'")
self.assertEqual(
await q1.get(),
(con, con.get_server_pid(), 'test', 'aaaa'))
self.assertEqual(
await q2.get(),
(con, con.get_server_pid(), 'test', 'aaaa'))
self.assertEqual(
await q3.get(),
(con, con.get_server_pid(), 'test', 'aaaa'))
await con.remove_listener('test', listener2)
await con.remove_listener('test', async_listener3)
await con.execute("NOTIFY test, 'aaaa'")
self.assertEqual(
await q1.get(),
(con, con.get_server_pid(), 'test', 'aaaa'))
with self.assertRaises(asyncio.TimeoutError):
await asyncio.wait_for(q2.get(), timeout=0.05)
await con.reset()
await con.remove_listener('test', listener1)
await con.execute("NOTIFY test, 'aaaa'")
with self.assertRaises(asyncio.TimeoutError):
await asyncio.wait_for(q1.get(), timeout=0.05)
with self.assertRaises(asyncio.TimeoutError):
await asyncio.wait_for(q2.get(), timeout=0.05)
async def test_listen_02(self):
async with self.create_pool(database='postgres') as pool:
async with pool.acquire() as con1, pool.acquire() as con2:
q1 = asyncio.Queue()
def listener1(*args):
q1.put_nowait(args)
await con1.add_listener('ipc', listener1)
await con2.execute("NOTIFY ipc, 'hello'")
self.assertEqual(
await q1.get(),
(con1, con2.get_server_pid(), 'ipc', 'hello'))
await con1.remove_listener('ipc', listener1)
async def test_listen_notletters(self):
async with self.create_pool(database='postgres') as pool:
async with pool.acquire() as con1, pool.acquire() as con2:
q1 = asyncio.Queue()
def listener1(*args):
q1.put_nowait(args)
await con1.add_listener('12+"34', listener1)
await con2.execute("""NOTIFY "12+""34", 'hello'""")
self.assertEqual(
await q1.get(),
(con1, con2.get_server_pid(), '12+"34', 'hello'))
await con1.remove_listener('12+"34', listener1)
async def test_dangling_listener_warns(self):
async with self.create_pool(database='postgres') as pool:
with self.assertWarnsRegex(
exceptions.InterfaceWarning,
'.*Connection.*is being released to the pool but '
'has 1 active notification listener'):
async with pool.acquire() as con:
def listener1(*args):
pass
await con.add_listener('ipc', listener1)
class TestLogListeners(tb.ConnectedTestCase):
@tb.with_connection_options(server_settings={
'client_min_messages': 'notice'
})
async def test_log_listener_01(self):
q1 = asyncio.Queue()
q2 = asyncio.Queue()
def notice_callb(con, message):
# Message fields depend on PG version, hide some values.
dct = message.as_dict()
del dct['server_source_line']
q1.put_nowait((con, type(message), dct))
async def async_notice_callb(con, message):
# Message fields depend on PG version, hide some values.
dct = message.as_dict()
del dct['server_source_line']
q2.put_nowait((con, type(message), dct))
async def raise_notice():
await self.con.execute(
"""DO $$
BEGIN RAISE NOTICE 'catch me!'; END;
$$ LANGUAGE plpgsql"""
)
async def raise_warning():
await self.con.execute(
"""DO $$
BEGIN RAISE WARNING 'catch me!'; END;
$$ LANGUAGE plpgsql"""
)
con = self.con
con.add_log_listener(notice_callb)
con.add_log_listener(async_notice_callb)
expected_msg = {
'context': 'PL/pgSQL function inline_code_block line 2 at RAISE',
'message': 'catch me!',
'server_source_function': 'exec_stmt_raise',
}
expected_msg_notice = {
**expected_msg,
'severity': 'NOTICE',
'severity_en': 'NOTICE',
'sqlstate': '00000',
}
expected_msg_warn = {
**expected_msg,
'severity': 'WARNING',
'severity_en': 'WARNING',
'sqlstate': '01000',
}
if con.get_server_version() < (9, 6):
del expected_msg_notice['context']
del expected_msg_notice['severity_en']
del expected_msg_warn['context']
del expected_msg_warn['severity_en']
await raise_notice()
await raise_warning()
msg = await q1.get()
msg[2].pop('server_source_filename', None)
self.assertEqual(
msg,
(con, exceptions.PostgresLogMessage, expected_msg_notice))
msg = await q1.get()
msg[2].pop('server_source_filename', None)
self.assertEqual(
msg,
(con, exceptions.PostgresWarning, expected_msg_warn))
msg = await q2.get()
msg[2].pop('server_source_filename', None)
self.assertEqual(
msg,
(con, exceptions.PostgresLogMessage, expected_msg_notice))
msg = await q2.get()
msg[2].pop('server_source_filename', None)
self.assertEqual(
msg,
(con, exceptions.PostgresWarning, expected_msg_warn))
con.remove_log_listener(notice_callb)
con.remove_log_listener(async_notice_callb)
await raise_notice()
self.assertTrue(q1.empty())
con.add_log_listener(notice_callb)
await raise_notice()
await q1.get()
self.assertTrue(q1.empty())
await con.reset()
await raise_notice()
self.assertTrue(q1.empty())
@tb.with_connection_options(server_settings={
'client_min_messages': 'notice'
})
async def test_log_listener_02(self):
q1 = asyncio.Queue()
cur_id = None
def notice_callb(con, message):
q1.put_nowait((con, cur_id, message.message))
con = self.con
await con.execute(
"CREATE FUNCTION _test(i INT) RETURNS int LANGUAGE plpgsql AS $$"
" BEGIN"
" RAISE NOTICE '1_%', i;"
" PERFORM pg_sleep(0.1);"
" RAISE NOTICE '2_%', i;"
" RETURN i;"
" END"
"$$"
)
try:
con.add_log_listener(notice_callb)
for cur_id in range(10):
await con.execute("SELECT _test($1)", cur_id)
for cur_id in range(10):
self.assertEqual(
q1.get_nowait(),
(con, cur_id, '1_%s' % cur_id))
self.assertEqual(
q1.get_nowait(),
(con, cur_id, '2_%s' % cur_id))
con.remove_log_listener(notice_callb)
self.assertTrue(q1.empty())
finally:
await con.execute('DROP FUNCTION _test(i INT)')
@tb.with_connection_options(server_settings={
'client_min_messages': 'notice'
})
async def test_log_listener_03(self):
q1 = asyncio.Queue()
async def raise_message(level, code):
await self.con.execute("""
DO $$ BEGIN
RAISE {} 'catch me!' USING ERRCODE = '{}';
END; $$ LANGUAGE plpgsql;
""".format(level, code))
def notice_callb(con, message):
# Message fields depend on PG version, hide some values.
q1.put_nowait(message)
self.con.add_log_listener(notice_callb)
await raise_message('WARNING', '99999')
msg = await q1.get()
self.assertIsInstance(msg, exceptions.PostgresWarning)
self.assertEqual(msg.sqlstate, '99999')
await raise_message('WARNING', '01004')
msg = await q1.get()
self.assertIsInstance(msg, exceptions.StringDataRightTruncation)
self.assertEqual(msg.sqlstate, '01004')
with self.assertRaises(exceptions.InvalidCharacterValueForCastError):
await raise_message('', '22018')
self.assertTrue(q1.empty())
async def test_dangling_log_listener_warns(self):
async with self.create_pool(database='postgres') as pool:
with self.assertWarnsRegex(
exceptions.InterfaceWarning,
'.*Connection.*is being released to the pool but '
'has 1 active log listener'):
async with pool.acquire() as con:
def listener1(*args):
pass
con.add_log_listener(listener1)
@unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
@unittest.skipIf(
platform.system() == 'Windows',
'not compatible with ProactorEventLoop which is default in Python 3.8+')
class TestConnectionTerminationListener(tb.ProxiedClusterTestCase):
async def test_connection_termination_callback_called_on_remote(self):
called = False
async_called = False
def close_cb(con):
nonlocal called
called = True
async def async_close_cb(con):
nonlocal async_called
async_called = True
con = await self.connect()
con.add_termination_listener(close_cb)
con.add_termination_listener(async_close_cb)
self.proxy.close_all_connections()
try:
await con.fetchval('SELECT 1')
except Exception:
pass
self.assertTrue(called)
self.assertTrue(async_called)
async def test_connection_termination_callback_called_on_local(self):
called = False
def close_cb(con):
nonlocal called
called = True
con = await self.connect()
con.add_termination_listener(close_cb)
await con.close()
await asyncio.sleep(0)
self.assertTrue(called)
|
