o
    i                     @   s  d dl Z d dlmZ d dlZd dlmZmZmZ d dlmZm	Z	 d dl
mZmZ d dlmZ d dlmZmZmZ d dlmZmZmZmZmZ d d	lmZ ervd d
lmZmZmZmZmZ d dl m!Z!m"Z"m#Z# d dlm$Z$ ededef dZ%zd dl&m'Z'm(Z(m)Z)m*Z*m+Z+ d dl,m-Z-m.Z.m/Z/ W n e0y   edw e-e.e/fZ1G dd deZ2d&ddZ3d'ddZ4d(dd Z5d)d"d#Z6d&d$d%Z7dS )*    N)datetime)continue_traceget_baggageget_traceparent)OP
SPANSTATUS)DidNotEnableIntegration)should_send_default_pii)BAGGAGE_HEADER_NAMESENTRY_TRACE_HEADER_NAMETransactionSource)capture_internal_exceptionsensure_integration_enabledevent_from_exceptionSENSITIVE_DATA_SUBSTITUTEreraise)TYPE_CHECKING)AnyCallableOptionalUnionTypeVar)EventProcessorEventHint)ExcInfoF.)bound)HueyResultResultGroupTaskPeriodicTask)CancelExecution	RetryTaskTaskLockedExceptionzHuey is not installedc                   @   s(   e Zd ZdZde ZedddZdS )HueyIntegrationhueyzauto.queue.returnNc                   C   s   t   t  d S N)patch_enqueuepatch_execute r-   r-   Y/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/sentry_sdk/integrations/huey.py
setup_once.   s   
zHueyIntegration.setup_oncer)   N)__name__
__module____qualname__
identifieroriginstaticmethodr/   r-   r-   r-   r.   r'   *   s
    
r'   r)   c                     s4   t j tt 						d	 fdd} | t _d S )
Nselfr   taskr"   r)   $Optional[Union[Result, ResultGroup]]c                    sd   t jtj|jtjd t|tst	t
 tt i|jd<  | |W  d    S 1 s+w   Y  d S )N)opnamer5   sentry_headers)
sentry_sdk
start_spanr   QUEUE_SUBMIT_HUEYr;   r'   r5   
isinstancer#   r   r   r   r   kwargs)r7   r8   old_enqueuer-   r.   _sentry_enqueue7   s   

$z&patch_enqueue.<locals>._sentry_enqueue)r7   r   r8   r"   r)   r9   )r   enqueuer   r'   )rD   r-   rB   r.   r+   4   s   
r+   r8   r   r   c                    s   d	 fdd}|S )
Neventr   hintr   r)   Optional[Event]c                    s   t  > | di } j|d<  j jk|d< | di } jt r% jntt r, j	nt jp1d j d|d< W d    | S 1 sDw   Y  | S )Ntagshuey_task_idhuey_task_retryextrar   )r8   argsrA   retryzhuey-job)
r   
setdefaultiddefault_retriesretriesr;   r
   rM   r   rA   )rF   rG   rI   rL   r8   r-   r.   event_processorN   s$   

z._make_event_processor.<locals>.event_processor)rF   r   rG   r   r)   rH   r-   )r8   rT   r-   rS   r.   _make_event_processorM   s   rU   exc_infor   c                 C   sf   t  }| d tv r|jtj d S |jtj t| t 	 j
tjddd\}}|j||d d S )Nr   F)typehandled)client_options	mechanism)rG   )r=   get_current_scopeHUEY_CONTROL_FLOW_EXCEPTIONStransaction
set_statusr   ABORTEDINTERNAL_ERRORr   
get_clientoptionsr'   r4   capture_event)rV   scoperF   rG   r-   r-   r.   _capture_exceptionh   s   

re   funcc                    s   t t d fdd}|S )NrM   r   rA   r)   c                     sB   z
 | i |}W |S  t y    t }t| t|  Y |S w r*   )	ExceptionsysrV   re   r   )rM   rA   resultrV   rf   r-   r.   _sentry_executey   s   
z+_wrap_task_execute.<locals>._sentry_execute)rM   r   rA   r   r)   r   )r   r'   )rf   rk   r-   rj   r.   _wrap_task_executex   s   
rl   c                     s:   t j tt 	 d							d fd	d
} | t _d S )Nr7   r   r8   r"   	timestampOptional[datetime]r)   r   c              	      s
  t  w}t  d|_|  |t| W d    n1 s!w   Y  |jdd }t	|p1i |j
tjtjtjd}|tj t|ddsQt|j|_d|_t |  | ||W  d    W  d    S 1 snw   Y  W d    d S 1 s~w   Y  d S )Nr(   r<   )r;   r:   sourcer5   _sentry_is_patchedFT)r=   isolation_scoper   _nameclear_breadcrumbsadd_event_processorrU   rA   popr   r;   r   QUEUE_TASK_HUEYr   TASKr'   r5   r^   r   OKgetattrrl   executerp   start_transaction)r7   r8   rm   rd   r<   r]   old_executer-   r.   rk      s0   

"z&patch_execute.<locals>._sentry_executer*   )r7   r   r8   r"   rm   rn   r)   r   )r   _executer   r'   )rk   r-   r|   r.   r,      s   
r,   r0   )r8   r   r)   r   )rV   r   r)   N)rf   r   r)   r   )8rh   r   r=   sentry_sdk.apir   r   r   sentry_sdk.constsr   r   sentry_sdk.integrationsr   r	   sentry_sdk.scoper
   sentry_sdk.tracingr   r   r   sentry_sdk.utilsr   r   r   r   r   typingr   r   r   r   r   r   sentry_sdk._typesr   r   r   r   r   huey.apir   r    r!   r"   r#   huey.exceptionsr$   r%   r&   ImportErrorr\   r'   r+   rU   re   rl   r,   r-   r-   r-   r.   <module>   s8    





