o
    ٰi(                  	   @   s  d dl mZ d dl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mZ d dlmZmZmZ z4d d	lmZmZmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlm Z m!Z!m"Z" d dl#m$Z$ W n e%yu   edw d dl&m'Z' e'rd dl&m(Z(m)Z)m*Z* d dl+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z4m5Z5m6Z6 d dlm7Z7 d dl8m9Z9 dZ:G dd deZ;G dd de	Z<d0ddZ=d0ddZ>d1d"d#Z?d0d$d%Z@d2d)d*ZAee;d+eBd&d'd,d-ddfd.d/ZCdS )3    )deepcopyN)OP)DidNotEnableIntegration)SentryAsgiMiddleware)should_send_default_pii)SOURCE_FOR_STYLETransactionSource)ensure_integration_enabledevent_from_exceptiontransaction_from_function)RequestStarliteState)BaseRouteHandler)DefineMiddleware)get_plugin_for_value)	HTTPRoute)ConnectionDataExtractoris_async_callableRef)	BaseModelzStarlite is not installed)TYPE_CHECKING)AnyOptionalUnion)
ASGIAppHintHTTPReceiveMessage	HTTPScopeMessage
MiddlewareReceiveScopeSendWebSocketReceiveMessage)MiddlewareProtocol)Eventzgeneric Starlite requestc                   @   s(   e Zd ZdZde ZedddZdS )StarliteIntegrationstarlitez
auto.http.returnNc                   C   s   t   t  t  d S )N)patch_app_initpatch_middlewarespatch_http_route_handle r.   r.   T/home/ubuntu/.local/lib/python3.10/site-packages/sentry_sdk/integrations/starlite.py
setup_once5   s   
zStarliteIntegration.setup_oncer*   N)__name__
__module____qualname__
identifieroriginstaticmethodr0   r.   r.   r.   r/   r(   1   s
    
r(   c                       s0   e Zd Zejfdddeddf fddZ  ZS )SentryStarliteASGIMiddlewareappr   span_originr*   Nc                    s   t  j|ddd|dd d S )NFendpointasgi   )r9   unsafe_context_datatransaction_stylemechanism_typer:   asgi_version)super__init__)selfr9   r:   	__class__r.   r/   rC   =   s   
z%SentryStarliteASGIMiddleware.__init__)r2   r3   r4   r(   r6   strrC   __classcell__r.   r.   rE   r/   r8   <   s    r8   r*   c                     s(   t j tt d
 fdd	} | t _dS )a  
    Replaces the Starlite class's `__init__` function in order to inject `after_exception` handlers and set the
    `SentryStarliteASGIMiddleware` as the outmost middleware in the stack.
    See:
    - https://starlite-api.github.io/starlite/usage/0-the-starlite-app/5-application-hooks/#after-exception
    - https://starlite-api.github.io/starlite/usage/7-middleware/0-middleware-intro/
    rD   r   argsr   kwargsr*   Nc                    sd   | dg }|jtgt|tr|n|gd |dpg }tg||d<  | g|R i | d S )Nafter_exception)rK   
middleware)popupdateexception_handler
isinstancelistgetr8   )rD   rI   rJ   rK   rL   old__init__r.   r/   injection_wrapperT   s   z)patch_app_init.<locals>.injection_wrapper)rD   r   rI   r   rJ   r   r*   N)r   rC   r
   r(   )rU   r.   rS   r/   r+   J   s   
r+   c                     s(   t j tt d fdd} | t _d S )NrD   r   r*   list[Middleware]c                    s   dd  | D S )Nc                 S   s   g | ]}t |qS r.   )enable_span_for_middleware).0rL   r.   r.   r/   
<listcomp>n   s    zIpatch_middlewares.<locals>.resolve_middleware_wrapper.<locals>.<listcomp>r.   )rD   old_resolve_middleware_stackr.   r/   resolve_middleware_wrapperl   s   z5patch_middlewares.<locals>.resolve_middleware_wrapper)rD   r   r*   rV   )r   resolve_middlewarer
   r(   )r\   r.   rZ   r/   r,   i   s   
r,   rL   r!   c                    s   t | dr	| tu r| S t| tr| jj n| j 										
	 d fdd} jdv}|r>t| tr;|| j_| S || _| S )N__call__rD   r&   scopeStarliteScopereceiver"   sendr$   r*   c                    s   t  td u r| |I d H S | jj t jtj tj	dR}|
d  						d fdd	}td
t}|dk}|sG|n}d fdd}	td
t}
|
dk}|sa|	n}| |||I d H W  d    S 1 sww   Y  d S )Nopnamer6   starlite.middleware_namerI   r   rJ   r*   2Union[HTTPReceiveMessage, WebSocketReceiveMessage]c                     s   t  td u r| i |I d H S t jtjtdttj	d}|
d  | i |I d H W  d    S 1 s?w   Y  d S Nr4   rc   rf   )
sentry_sdk
get_clientget_integrationr(   
start_spanr   MIDDLEWARE_STARLITE_RECEIVEgetattrrG   r6   set_tag)rI   rJ   span)middleware_namera   r.   r/   _sentry_receive   s   $zNenable_span_for_middleware.<locals>._create_span_call.<locals>._sentry_receiver2   rr   messager    c                    s   t  td u r| I d H S t jtjtdttj	d}|
d  | I d H W  d    S 1 s9w   Y  d S rh   )ri   rj   rk   r(   rl   r   MIDDLEWARE_STARLITE_SENDrn   rG   r6   ro   )rs   rp   )rq   rb   r.   r/   _sentry_send   s   $zKenable_span_for_middleware.<locals>._create_span_call.<locals>._sentry_sendru   )rI   r   rJ   r   r*   rg   )rs   r    r*   N)ri   rj   rk   r(   rF   r2   rl   r   MIDDLEWARE_STARLITEr6   ro   rn   rG   )rD   r_   ra   rb   middleware_spanrr   receive_namereceive_patchednew_receiveru   	send_namesend_patchednew_sendold_call)rq   ra   rb   r/   _create_span_call   s6   $z5enable_span_for_middleware.<locals>._create_span_call)r   )
rD   r&   r_   r`   ra   r"   rb   r$   r*   N)hasattrr8   rP   r   rL   r^   r2   )rL   r   not_yet_patchedr.   r~   r/   rW   v   s2   


5
rW   c                     s2   t j 											 d fd
d} | t _d S )NrD   r   r_   r   ra   r"   rb   r$   r*   c                    s   t  td u r| ||I d H S t  }d j||d}tddd|  d}|I d H d fdd}tj|_	|
| | ||I d H S )Nr9   )r_   ra   rb   T)
parse_bodyparse_querybodyeventr'   _r   r*   c                    s    d}|  di }t dd|d< t r d |d< d ur&|d< d }|jd ur1|j}nt|jtr<|jj}n|j}|d urGt|}dt	d	 i}|sVt
}dtji}| t|||d
 | S )Nroute_handlerrequest_body    content_lengthcookiesdatasourcer;   )r   transactiontransaction_info)rR   lenr   re   rP   fnr   valuer   r   _DEFAULT_TRANSACTION_NAMEr	   ROUTErN   r   )r   r   r   request_infofunctx_nametx_infoextracted_request_datarequest_datar_   r.   r/   event_processor   s4   



zHpatch_http_route_handle.<locals>.handle_wrapper.<locals>.event_processor)r   r'   r   r   r*   r'   )ri   rj   rk   r(   get_isolation_scoperequest_classr   rM   r5   _nameadd_event_processor)rD   r_   ra   rb   sentry_scoper   r   r   
old_handler   r/   handle_wrapper   s$   

#
z/patch_http_route_handle.<locals>.handle_wrapper)
rD   r   r_   r   ra   r"   rb   r$   r*   N)r   handle)r   r.   r   r/   r-      s   
9r-   r_   r`   Optional[dict[str, Any]]c                 C   sh   |  d}|s	d S t|tr|S t|tr| S t|dr"| S t|}|r2t|js2||S d S )Nuserasdict)	rR   rP   dictr   r   r   r   r   to_dict)r_   
scope_userpluginr.   r.   r/   retrieve_user_from_scope  s   




r   excr   r   c                 C   sd   d }t  r	t|}|rt|trt }|| t| t j	t
jddd\}}tj||d d S )NF)typehandled)client_options	mechanism)hint)r   r   rP   r   ri   r   set_userr   rj   optionsr(   r5   capture_event)r   r_   r   	user_infor   r   r   r.   r.   r/   rO     s   


rO   r1   )rL   r!   r*   r!   )r_   r`   r*   r   )Dcopyr   ri   sentry_sdk.constsr   sentry_sdk.integrationsr   r   sentry_sdk.integrations.asgir   sentry_sdk.scoper   sentry_sdk.tracingr   r	   sentry_sdk.utilsr
   r   r   r)   r   r   r   starlite.handlers.baser   starlite.middlewarer   starlite.plugins.baser   starlite.routes.httpr   starlite.utilsr   r   r   pydanticr   ImportErrortypingr   r   r   r   starlite.typesr   r   r   r   r    r!   r"   r#   r`   r$   r%   r&   sentry_sdk._typesr'   r   r(   r8   r+   r,   rW   r-   r   	ExceptionrO   r.   r.   r.   r/   <module>   sF    0



L
? 