11import asyncio
22import functools
33import random
4+ from collections .abc import AsyncGenerator , Awaitable , Callable , Iterable
45from contextlib import asynccontextmanager
56from dataclasses import dataclass
67from datetime import timedelta
78from 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
1910
2011import psycopg .sql as sql
2112from latch_config .config import PostgresConnectionConfig
3728 DeadlockDetected ,
3829 DiskFull ,
3930 DuplicateFile ,
40- )
41- from psycopg .errors import Error as PGError
42- from psycopg .errors import (
4331 IdleSessionTimeout ,
4432 InsufficientResources ,
4533 IoError ,
5543 TooManyConnections ,
5644 UndefinedFile ,
5745)
46+ from psycopg .errors import Error as PGError
5847from psycopg .rows import AsyncRowFactory , Row , dict_row , kwargs_row
5948from psycopg .types .composite import CompositeInfo , register_composite
6049from psycopg .types .enum import EnumInfo , register_enum
@@ -158,15 +147,9 @@ class LatchAsyncConnection(AsyncConnection[Row]):
158147 trace_attributes : Attributes
159148
160149 def cursor (
161- self ,
162- row_factory : AsyncRowFactory [Any ],
163- * ,
164- binary : bool = True ,
150+ self , row_factory : AsyncRowFactory [Any ], * , binary : bool = True
165151 ) -> 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 )
170153 assert isinstance (res , TracedAsyncCursor )
171154 res .trace_attributes = self .trace_attributes
172155 return res
@@ -291,18 +274,12 @@ async def _configure_wrapper(self, conn: AsyncConnection[object]):
291274 )
292275
293276 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" ):
297278 await conn .query_opt (dict , cmd )
298279
299280 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" ):
306283 for db_type , f in self .composite_type_map .items ():
307284 type_info = await CompositeInfo .fetch (conn , db_type )
308285 if type_info is None :
@@ -312,12 +289,8 @@ async def run_composite_setup():
312289 register_composite (type_info , conn , f )
313290
314291 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" ):
321294 # query from
322295 # https://github.com/psycopg/psycopg/blob/fd1659118e96a48f22b4e67ff17c2cdab8bd0e84/psycopg/psycopg/_typeinfo.py#L148
323296 raw_responses = await conn .queryn (
@@ -406,9 +379,7 @@ async def run_enum_setup():
406379
407380 async def getconn (self , timeout : float | None = None ) -> AsyncConnection [object ]:
408381 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
412383 ):
413384 return await super ().getconn (timeout )
414385
@@ -443,9 +414,7 @@ def mixin_dict(a: dict[str, object], b: dict[str, object]):
443414def pg_error_to_dict (x : PGError , * , short : bool = False ):
444415 diagnostic_obj = {
445416 "severity" : x .diag .severity ,
446- "message" : {
447- "detail" : x .diag .message_detail ,
448- },
417+ "message" : {"detail" : x .diag .message_detail },
449418 }
450419
451420 if not short :
0 commit comments