
    \j                        d dl 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
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 d dlmZ  ed	      Z	 d d
lmZ d dlmZmZ d dlmZ d dl m!Z! d dlm#Z# e#rd dlm$Z$m%Z%m&Z&m'Z'm(Z( d dl)m*Z*m+Z+  G d de
      Z,ddZ- G d de      Z.	 	 	 	 	 	 ddZ/ G d d      Z0y# e"$ r	  ed      w xY w)    N)OP
SPANSTATUS)continue_traceget_baggageget_traceparent)IntegrationDidNotEnable)request_body_within_bounds)BAGGAGE_HEADER_NAMESENTRY_TRACE_HEADER_NAMETransactionSource)AnnotatedValuecapture_internal_exceptionsevent_from_exception)TypeVarR)Broker)
Middlewaredefault_middleware)Retry)MessagezDramatiq is not installed)TYPE_CHECKING)AnyCallableDictOptionalUnion)EventHintc                   0    e Zd ZdZdZde Zedd       Zy)DramatiqIntegrationa?  
    Dramatiq integration for Sentry

    Please make sure that you call `sentry_sdk.init` *before* initializing
    your broker, as it monkey patches `Broker.__init__`.

    This integration was originally developed and maintained
    by https://github.com/jacobsvante and later donated to the Sentry
    project.
    dramatiqzauto.queue.Nc                      t                y N)_patch_dramatiq_broker     J/root/env/lib/python3.12/site-packages/sentry_sdk/integrations/dramatiq.py
setup_oncezDramatiqIntegration.setup_once4   s     r'   returnN)__name__
__module____qualname____doc__
identifieroriginstaticmethodr)   r&   r'   r(   r!   r!   %   s,    	 J:,'F! !r'   r!   c                  X    t         j                  	 	 	 	 	 	 	 	 dfd} | t         _        y )Nc                    t        j                         j                  t              }	 |j	                  d      }|t        d t        D              }nt        |      }|:|D cg c]  }t        |t              r| }}|j                  dt                      ||d<    | g|i | y # t
        $ r t        |      dk(  r|d   }g }nd }Y w xY wc c}w )N
middleware   r   c              3   *   K   | ]  } |         y wr$   r&   ).0ms     r(   	<genexpr>zP_patch_dramatiq_broker.<locals>.sentry_patched_broker__init__.<locals>.<genexpr>N   s     >aac>s   )
sentry_sdk
get_clientget_integrationr!   popKeyErrorlenlistr   
isinstanceSentryMiddlewareinsert)selfargskwintegrationr5   r9   original_broker__init__s         r(   sentry_patched_broker__init__z=_patch_dramatiq_broker.<locals>.sentry_patched_broker__init__<   s     !++-==>QR
	"-J >+=>>Jj)J"%/Wz!EU7V!WJWa!1!34%<2t2r2)  	" 4yA~!!W
!
	"  Xs   B- %C;C-!CC)rE   r   rF   r   rG   r   r+   N)r   __init__)rJ   rI   s    @r(   r%   r%   9   s;    $oo33$3,13	3: 4FOr'   c                   Z    e Zd ZdZdZdddddedd	fd
ZddZd	d	d	 	 	 	 	 	 	 	 	 ddZeZ	y	)rC   z
    A Dramatiq middleware that automatically captures and sends
    exceptions to Sentry.

    This is automatically added to every instantiated broker via the
    DramatiqIntegration.
    _sentry_headersbrokerr   message
Message[R]delayr+   Nc                     t        j                         j                  t              }|y t        t               t        t               i|j                  | j                  <   y r$   )
r;   r<   r=   r!   r   r   r   r   optionsSENTRY_HEADERS_NAME)rE   rN   rO   rQ   rH   s        r(   before_enqueuezSentryMiddleware.before_enqueueg   sN     !++-==>QR  $o&75
001r'   c                 "   t        j                         j                  t              }|y t        j                         |_        |j
                  j                         }|j                          |j                  d|j                         |j                  t        ||             |j                  j                  | j                        xs i }d|j                  v ri }t        ||j                   t"        j$                  t&        j(                  t        j*                        }|j-                  t.        j0                         t        j2                  ||j                   t"        j$                  t&        j(                         |j                          y )Ndramatiq_message_idretries)nameopsourcer1   )rY   rZ   r[   )r;   r<   r=   r!   isolation_scope_scope_manager	__enter__clear_breadcrumbs	set_extra
message_idadd_event_processor_make_message_event_processorrS   getrT   r   
actor_namer   QUEUE_TASK_DRAMATIQr   TASKr1   
set_statusr   OKstart_transaction)rE   rN   rO   rH   scopesentry_headerstransactions          r(   before_process_messagez'SentryMiddleware.before_process_messages   s0    ++-==>QR!+!;!;!=&&002!-w/A/AB!!"?"UV ,,T-E-EFL"'N$##%%$))&--
 	z}}-$$##%%$))		
 	r'   )result	exceptionc                   t        j                         j                  t              }|y |j	                  |j
                        }|j                  j                  d      xs |j                  j                  d      }|j                  }t        j                         j                  }	|	sy |d uxr$ |xr t        ||       xr t        |t               }
|
s'|	j                  d d d        |j                  d d d        y t        |t        j                         j                  t        j                  dd      \  }}t        j                   ||       |	j                  t#        |      |d        |j                  t#        |      |d        y )NthrowsF)typehandled)client_options	mechanism)hint)r;   r<   r=   r!   	get_actorre   rS   rd   r]   get_current_scoperm   rB   r   __exit__r   r0   capture_eventrs   )rE   rN   rO   ro   rp   rH   actorrr   scope_managerrm   is_event_capture_requiredeventrw   s                r(   after_process_messagez&SentryMiddleware.after_process_message   s`    !++-==>QR  !3!34$$X.M%--2C2CH2M.. 224@@ T! 1=
9f =>1y%00 	"
 )  tT2""4t4*%002::+66 
t 	  T2T)_i>tI	4@r'   )rN   r   rO   rP   r+   N)
rN   r   rO   rP   ro   zOptional[Any]rp   zOptional[Exception]r+   N)
r,   r-   r.   r/   rT   intrU   rn   r   after_skip_messager&   r'   r(   rC   rC   \   s     ,



)5

>A

	

 J #'+/*A*A *A
  *A )*A 
*AX /r'   rC   c                      d fd}|S )Nc                 |    t               5  t              j                  |        d d d        | S # 1 sw Y   | S xY wr$   )r   DramatiqMessageExtractorextract_into_event)r   rw   rO   s     r(   innerz,_make_message_event_processor.<locals>.inner   s=    (* 	H$W-@@G	H 	H s   1;)r   r   rw   r   r+   zOptional[Event]r&   )rO   rH   r   s   `  r(   rc   rc      s     Lr'   c                   (    e Zd ZddZdefdZddZy)r   r+   Nc                 @    t        |j                               | _        y r$   )dictasdictmessage_data)rE   rO   s     r(   rK   z!DramatiqMessageExtractor.__init__   s     !12r'   c                 R    t        t        j                  | j                              S r$   )r@   jsondumpsr   )rE   s    r(   content_lengthz'DramatiqMessageExtractor.content_length   s    4::d//011r'   c                 $   t        j                         }|j                         sy |j                  di       }|j                  di       }d|d<   d }t	        || j                               st        j                         }n| j                  }||d<   y )Ncontextsr"   rs   data)	r;   r<   	is_active
setdefaultr
   r   r   removed_because_over_size_limitr   )rE   r   clientr   request_infor   s         r(   r   z+DramatiqMessageExtractor.extract_into_event   s    &&(!##J3**:r:)VBF)&$2E2E2GH!AACD$$D#Vr'   )rO   rP   r+   N)r   r   r+   N)r,   r-   r.   rK   r   r   r   r&   r'   r(   r   r      s    32 2$r'   r   r*   )rO   rP   rH   r!   r+   z(Callable[[Event, Hint], Optional[Event]])1r   r;   sentry_sdk.constsr   r   sentry_sdk.apir   r   r   sentry_sdk.integrationsr   r	   $sentry_sdk.integrations._wsgi_commonr
   sentry_sdk.tracingr   r   r   sentry_sdk.utilsr   r   r   typingr   r   dramatiq.brokerr   dramatiq.middlewarer   r   dramatiq.errorsr   dramatiq.messager   ImportErrorr   r   r   r   r   r   sentry_sdk._typesr   r   r!   r%   rC   rc   r   r&   r'   r(   <module>r      s      , G G = K 
 
 CL4&B%( !;;-!+ !( 4Fc/z c/L		(=	/	$ $g  4
2
334s   B3 3C