o
    aiK                     @  sD  d dl mZ d dlZd dlZd dlmZmZmZmZm	Z	 d dl
Z
d dlmZmZ d dlmZmZmZ d dlmZ d dlmZmZ d dlmZmZmZ zd dlZd d	lmZ W n eya   ed
w G dd deZedZ d2ddZ!d3ddZ"ededZ#ej$ddd4d"d#Z%ddd5d$d%Z&d6d'd(Z'd3d)d*Z(d7d0d1Z)dS )8    )annotationsN)AnyTypeVarCallable	AwaitableIterator)OPSPANDATA)_check_minimum_versionIntegrationDidNotEnable)Span)add_query_sourcerecord_sql_queries)ensure_integration_enabledparse_versioncapture_internal_exceptions)
BaseCursorzasyncpg not installed.c                   @  s<   e Zd ZdZde ZdZdddddZedddZdS )AsyncPGIntegrationasyncpgzauto.db.F)record_paramsr   boolc                C  s
   |t _d S )N)r   _record_params)selfr    r   ]/home/ubuntu/SoloSpeech/.venv/lib/python3.10/site-packages/sentry_sdk/integrations/asyncpg.py__init__   s   
zAsyncPGIntegration.__init__returnNonec                  C  s|   t tj} tt|  ttjjtj_ttjj	tj_	ttjj
ddtj_
ttjjtj_ttjjtj_ttjjtj_d S )NTexecutemany)r   r   __version__r
   r   _wrap_execute
Connectionexecute_wrap_connection_method_execute_executemany_wrap_cursor_creationcursorprepare_wrap_connect_addrconnect_utils_connect_addr)asyncpg_versionr   r   r   
setup_once!   s    


zAsyncPGIntegration.setup_onceN)r   r   )r   r   )	__name__
__module____qualname__
identifieroriginr   r   staticmethodr/   r   r   r   r   r      s    
r   Tquerystrr   c                 C  s   t dd|  S )Nz\s+ )resubstrip)r7   r   r   r   _normalize_query;   s   r=   f'Callable[..., Awaitable[T]]'c                      d fdd}|S )	Nargs'Any'kwargsr   'T'c                    s   t  td u r | i |I d H S t| dkr$ | i |I d H S t| d }td |d d dtjd} | i |I d H }W d    n1 sJw   Y  t  t	| W d    |S 1 sbw   Y  |S )N      F)r)   r7   params_list
paramstyler    span_origin)

sentry_sdk
get_clientget_integrationr   lenr=   r   r4   r   r   )rA   rC   r7   spanresr>   r   r   _inner@   s.   


z_wrap_execute.<locals>._innerrA   rB   rC   rB   r   rD   r   r>   rQ   r   rP   r   r"   ?   s   r"   	SubCursor)boundFr   r)   'SubCursor | None'rG   'tuple[Any, ...] | None'r    r   'Iterator[Span]'c             	   c  s~    t  t}|d ur|jsd }|rdnd }t|}t| ||||| d utjd}|V  W d    d S 1 s8w   Y  d S )Npyformat)r)   r7   rG   rH   r    record_cursor_reprrI   )rJ   rK   rL   r   r   r=   r   r4   )r)   r7   rG   r    integrationparam_stylerN   r   r   r   _recorda   s$   "r]   c                  s   d fdd}|S )	NrA   rB   rC   r   rD   c                    s   t  td u r| i |I d H S | d }t| dkr"| d nd }td || d}t|| d  | i |I d H }W d    |S 1 sIw   Y  |S )NrF   rE   r   r   )rJ   rK   rL   r   rM   r]   _set_db_datarA   rC   r7   rG   rN   rO   r    r>   r   r   rQ      s   
z'_wrap_connection_method.<locals>._innerrR   r   )r>   r    rQ   r   r`   r   r%   |   s   r%   'Callable[..., T]'c                   s   t t d fdd}|S )	NrA   rB   rC   r   rD   c                    s~   | d }t | dkr| d nd }td ||dd}t|| d   | i |}|d| W d    |S 1 s8w   Y  |S )NrF   rE   Fr   r   z	db.cursor)rM   r]   r^   set_datar_   rP   r   r   rQ      s    


z%_wrap_cursor_creation.<locals>._innerrR   )r   r   rS   r   rP   r   r(      s   r(   c                   r@   )	NrA   rB   rC   r   rD   c               	     s6  t  td u r | i |I d H S |d j}|d j}t jtjdtj	di}|
tjd |d}|rUz|
tj|d  |
tj|d  W n	 tyT   Y nw |
tj| |
tj| t  t jdd|jd	 W d    n1 szw   Y   | i |I d H }W d    |S 1 sw   Y  |S )
Nparamsconnect)opnamer4   
postgresqladdrr   rF   r7   )messagecategorydata)rJ   rK   rL   r   userdatabase
start_spanr   DBr4   rb   r	   	DB_SYSTEMgetSERVER_ADDRESSSERVER_PORT
IndexErrorDB_NAMEDB_USERr   add_breadcrumb_data)rA   rC   rl   rm   rN   rh   rO   rP   r   r   rQ      s@   



z"_wrap_connect_addr.<locals>._innerrR   r   rS   r   rP   r   r+      s   r+   rN   'Span'connrB   r   c                 C  s   |  tjd |j}|r*z|  tj|d  |  tj|d  W n	 ty)   Y nw |jj}|r7|  tj	| |jj
}|rF|  tj| d S d S )Nrg   r   rF   )rb   r	   rp   _addrrr   rs   rt   _paramsrm   ru   rl   rv   )rN   rz   rh   rm   rl   r   r   r   r^      s    r^   )r7   r8   r   r8   )r>   r?   r   r?   )
r)   rV   r7   r8   rG   rW   r    r   r   rX   )r>   r?   r    r   r   r?   )r>   ra   r   ra   )rN   ry   rz   rB   r   r   )*
__future__r   
contextlibr:   typingr   r   r   r   r   rJ   sentry_sdk.constsr   r	   sentry_sdk.integrationsr
   r   r   sentry_sdk.tracingr   sentry_sdk.tracing_utilsr   r   sentry_sdk.utilsr   r   r   r   asyncpg.cursorr   ImportErrorr   r6   r=   r"   rT   contextmanagerr]   r%   r(   r+   r^   r   r   r   r   <module>   s:    



%