1
1
import asyncio
2
2
import functools
3
3
import random
4
+ from collections .abc import AsyncGenerator , Awaitable , Callable , Iterable
4
5
from contextlib import asynccontextmanager
5
6
from dataclasses import dataclass
6
7
from datetime import timedelta
7
8
from textwrap import dedent
8
- from typing import (
9
- Any ,
10
- AsyncGenerator ,
11
- Awaitable ,
12
- Callable ,
13
- Concatenate ,
14
- Iterable ,
15
- ParamSpec ,
16
- TypeVar ,
17
- cast ,
18
- )
9
+ from typing import Any , Concatenate , ParamSpec , TypeVar , cast
19
10
20
11
import psycopg .sql as sql
21
12
from latch_config .config import PostgresConnectionConfig
37
28
DeadlockDetected ,
38
29
DiskFull ,
39
30
DuplicateFile ,
40
- )
41
- from psycopg .errors import Error as PGError
42
- from psycopg .errors import (
43
31
IdleSessionTimeout ,
44
32
InsufficientResources ,
45
33
IoError ,
55
43
TooManyConnections ,
56
44
UndefinedFile ,
57
45
)
46
+ from psycopg .errors import Error as PGError
58
47
from psycopg .rows import AsyncRowFactory , Row , dict_row , kwargs_row
59
48
from psycopg .types .composite import CompositeInfo , register_composite
60
49
from psycopg .types .enum import EnumInfo , register_enum
@@ -158,15 +147,9 @@ class LatchAsyncConnection(AsyncConnection[Row]):
158
147
trace_attributes : Attributes
159
148
160
149
def cursor (
161
- self ,
162
- row_factory : AsyncRowFactory [Any ],
163
- * ,
164
- binary : bool = True ,
150
+ self , row_factory : AsyncRowFactory [Any ], * , binary : bool = True
165
151
) -> AsyncCursor [Any ]:
166
- res = super ().cursor (
167
- row_factory = row_factory ,
168
- binary = binary ,
169
- )
152
+ res = super ().cursor (row_factory = row_factory , binary = binary )
170
153
assert isinstance (res , TracedAsyncCursor )
171
154
res .trace_attributes = self .trace_attributes
172
155
return res
@@ -291,18 +274,12 @@ async def _configure_wrapper(self, conn: AsyncConnection[object]):
291
274
)
292
275
293
276
async def run_setup (cmd : sql .SQL ):
294
- with tracer .start_as_current_span (
295
- "setup command" ,
296
- ):
277
+ with tracer .start_as_current_span ("setup command" ):
297
278
await conn .query_opt (dict , cmd )
298
279
299
280
async def run_composite_setup ():
300
- with tracer .start_as_current_span (
301
- "composite type setup" ,
302
- ):
303
- with tracer .start_as_current_span (
304
- "register composite types" ,
305
- ):
281
+ with tracer .start_as_current_span ("composite type setup" ):
282
+ with tracer .start_as_current_span ("register composite types" ):
306
283
for db_type , f in self .composite_type_map .items ():
307
284
type_info = await CompositeInfo .fetch (conn , db_type )
308
285
if type_info is None :
@@ -312,12 +289,8 @@ async def run_composite_setup():
312
289
register_composite (type_info , conn , f )
313
290
314
291
async def run_enum_setup ():
315
- with tracer .start_as_current_span (
316
- "enum setup" ,
317
- ):
318
- with tracer .start_as_current_span (
319
- "fetch type info" ,
320
- ):
292
+ with tracer .start_as_current_span ("enum setup" ):
293
+ with tracer .start_as_current_span ("fetch type info" ):
321
294
# query from
322
295
# https://github.com/psycopg/psycopg/blob/fd1659118e96a48f22b4e67ff17c2cdab8bd0e84/psycopg/psycopg/_typeinfo.py#L148
323
296
raw_responses = await conn .queryn (
@@ -406,9 +379,7 @@ async def run_enum_setup():
406
379
407
380
async def getconn (self , timeout : float | None = None ) -> AsyncConnection [object ]:
408
381
with tracer .start_as_current_span (
409
- "postgres.connect" ,
410
- kind = SpanKind .CLIENT ,
411
- attributes = self ._trace_attributes ,
382
+ "postgres.connect" , kind = SpanKind .CLIENT , attributes = self ._trace_attributes
412
383
):
413
384
return await super ().getconn (timeout )
414
385
@@ -443,9 +414,7 @@ def mixin_dict(a: dict[str, object], b: dict[str, object]):
443
414
def pg_error_to_dict (x : PGError , * , short : bool = False ):
444
415
diagnostic_obj = {
445
416
"severity" : x .diag .severity ,
446
- "message" : {
447
- "detail" : x .diag .message_detail ,
448
- },
417
+ "message" : {"detail" : x .diag .message_detail },
449
418
}
450
419
451
420
if not short :
0 commit comments