o
    Á¿i  ã                   @   s¢   d dl Z d dl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 e deƒ ¡ d	efd
d„Zd	eeef fdd„Zdd„ Zdd„ Zdd„ Zdd„ ZdS )é    N)Úconfig)ÚPin)Úunwrap)Útraceré   )ÚSchemaExtractorÚavroÚreturnc                   C   s   t tddƒS )NÚ__version__Ú )Úgetattrr   © r   r   úW/home/ubuntu/.local/lib/python3.10/site-packages/ddtrace/contrib/internal/avro/patch.pyÚget_version   s   r   c                   C   s   ddiS )Nr   Ú*r   r   r   r   r   Ú_supported_versions   s   r   c                  C   sX   t tddƒrdS dt_tj} | ddtƒ | ddtƒ tƒ  tj	j
¡ tƒ  tj	j¡ dS )zPatch the instrumented methodsÚ_datadog_patchFNTzavro.iozDatumReader.readzDatumWriter.write)r   r   r   ÚwraptÚwrap_function_wrapperÚ_traced_deserializeÚ_traced_serializer   ÚontoÚioÚDatumReaderÚDatumWriter)Ú_wr   r   r   Úpatch   s   r   c                   C   s6   t tddƒrdt_ttjjdƒ ttjjdƒ d S d S )Nr   FÚreadÚwrite)r   r   r   r   r   r   r   r   r   r   r   Úunpatch(   s
   ür   c                 C   s„   t js
| |i |¤ŽS t |¡}|r| ¡ s| |i |¤ŽS t ¡ }z| |i |¤ŽW |r4t |j	|tj
¡ S S |rAt |j	|tj
¡ w w ©N)r   Ú_data_streams_enabledr   Úget_fromÚenabledr   Úcurrent_spanr   Úattach_schema_on_spanÚwriters_schemaÚSERIALIZATION)ÚfuncÚinstanceÚargsÚkwargsÚpinÚactiver   r   r   r   3   s   
ÿÿr   c                 C   s˜   t js
| |i |¤ŽS t |¡}|r| ¡ s| |i |¤ŽS t ¡ }z| |i |¤ŽW |}|r8|r9t |j	|tj
¡ S S S |}|rJ|rKt |j	|tj
¡ w w w r    )r   r!   r   r"   r#   r   r$   r   r%   r&   ÚDESERIALIZATION)r(   r)   r*   r+   r,   r-   Úreaderr   r   r   r   E   s    
ÿÿÿr   )r   r   Úddtracer   Úddtrace._trace.pinr   Úddtrace.internal.utils.wrappersr   Úddtrace.tracer   Úschema_iteratorr   Ú_addÚdictÚstrr   r   r   r   r   r   r   r   r   r   Ú<module>   s"    þ