o
    RDi                     @  s   d dl mZ d dlZd dlmZ d dlmZ d dlmZm	Z	 ee
Zzd dlmZ W n ey9   d dlmZ Y nw dZG d	d
 d
eZdddZG dd dZG dd dZdS )    )annotationsN)	getLogger)TracebackType)CallableIterator)_Timer)Timeri@8  c                      s,   e Zd ZdZd fdd	Zdd
dZ  ZS )HeartBeatTimerz_A thread which executes a function every client_session_keep_alive_heartbeat_frequency seconds.-client_session_keep_alive_heartbeat_frequencyintfr   returnNonec                   s   |}t  || d| _d S )NT)super__init__daemon)selfr
   r   interval	__class__ V/var/www/Datamplify/venv/lib/python3.10/site-packages/snowflake/connector/time_util.pyr      s   
zHeartBeatTimer.__init__c              
   C  sr   | j  s7| j | j | j  s0z|   W n ty/ } ztd| W Y d }~nd }~ww | j  rd S d S )Nzfailed to heartbeat: %s)finishedis_setwaitr   function	Exceptionloggerdebug)r   er   r   r   run   s   

zHeartBeatTimer.run)r
   r   r   r   r   r   r   r   )__name__
__module____qualname____doc__r   r    __classcell__r   r   r   r   r	      s    r	   r   r   c                   C  s   t t d S )z)Returns the current time in milliseconds.  )r   timer   r   r   r   get_time_millis(   s   r)   c                   @  s8   e Zd ZdZdddZdddZdddZdddZdS )TimerContextManagera=  Context manager class to easily measure execution of a code block.

    Once the context manager finishes, the class should be cast into an int to retrieve
    result.

    Example:

        with TimerContextManager() as measured_time:
            pass
        download_metric = measured_time.get_timing_millis()
    r   r   c                 C  s   d | _ d | _d S N)_start_endr   r   r   r   r   :   s   
zTimerContextManager.__init__c                 C  s   t  | _| S r+   )r)   r,   r.   r   r   r   	__enter__>   s   zTimerContextManager.__enter__exc_typetype[BaseException] | Noneexc_valBaseException | Noneexc_tbTracebackType | Nonec                 C     t  | _d S r+   )r)   r-   )r   r0   r2   r4   r   r   r   __exit__B   s   zTimerContextManager.__exit__r   c                 C  s(   | j du s
| jdu rtd| j| j  S )z$Get measured timing in milliseconds.Nz<Trying to get timing before TimerContextManager has finished)r,   r-   r   r.   r   r   r   get_timing_millisJ   s
   z%TimerContextManager.get_timing_millisNr!   )r   r*   )r0   r1   r2   r3   r4   r5   r   r   r   r   )r"   r#   r$   r%   r   r/   r7   r8   r   r   r   r   r*   -   s    


r*   c                   @  s   e Zd ZdZ			ddd
dZed ddZed!ddZed!ddZed!ddZ	ed"ddZ
d!ddZd#ddZd#ddZdS )$TimeoutBackoffCtxz:Base context for handling timeouts and backoffs on retriesNmax_retry_attempts
int | Nonetimeoutbackoff_generatorIterator | Noner   r   c                 C  s,   || _ || _|| _d| _|  | _d | _d S Nr   )_backoff_generator_max_retry_attempts_timeout_current_retry_count_advance_backoff_current_sleep_time_start_time_millis)r   r;   r=   r>   r   r   r   r   V   s   

zTimeoutBackoffCtx.__init__c                 C  s   | j S r+   )rC   r.   r   r   r   r=   f   s   zTimeoutBackoffCtx.timeoutr   c                 C  
   t | jS r+   )r   rD   r.   r   r   r   current_retry_countj      
z%TimeoutBackoffCtx.current_retry_countc                 C  rH   r+   )r   rF   r.   r   r   r   current_sleep_timen   rJ   z$TimeoutBackoffCtx.current_sleep_timec                 C  sB   | j d u r	td| jd u rtd| jd }t | j  }|| S )NzKStart time not recorded in remaining_time_millis, call set_start_time firstz(Timeout is None in remaining_time_millisr'   )rG   	TypeErrorrC   r)   )r   timeout_milliselapsed_time_millisr   r   r   remaining_time_millisr   s   


z'TimeoutBackoffCtx.remaining_time_millisboolc                 C  s\   | j dur| jdu rtd | j dur| jdk nd}| jdur&| j| jknd}| o-| S )z$Decides whether to retry connection.Nz=Timeout set in TimeoutBackoffCtx, but start time not recordedr   F)rC   rG   r   warningrO   rB   rD   )r   	timed_outretry_attempts_exceededr   r   r   should_retry   s   
zTimeoutBackoffCtx.should_retryc                 C  s   | j d ur
t| j S dS r@   )rA   nextr.   r   r   r   rE      s   z"TimeoutBackoffCtx._advance_backoffc                 C  r6   r+   )r)   rG   r.   r   r   r   set_start_time   s   z TimeoutBackoffCtx.set_start_timec                 C  sB   |  j d7  _ |  | _td| j   td| j d dS )z4Updates retry count and sleep time for another retry   zUpdate retry count to zUpdate sleep time to z secondsN)rD   rE   rF   r   r   r.   r   r   r   	increment   s   
zTimeoutBackoffCtx.increment)NNN)r;   r<   r=   r<   r>   r?   r   r   )r   r<   r9   )r   rP   r!   )r"   r#   r$   r%   r   propertyr=   rI   rK   rO   rT   rE   rV   rX   r   r   r   r   r:   S   s&    

r:   r9   )
__future__r   r(   loggingr   typesr   typingr   r   r"   r   	threadingr   r   ImportError"DEFAULT_MASTER_VALIDITY_IN_SECONDSr	   r)   r*   r:   r   r   r   r   <module>   s    
&