o
    پi                     @   sV   d Z ddlZddlZddlmZ ddlZddlmZ ddlm	Z
 G dd de
jZdS )z<Reference implementation for health checking in gRPC Python.    N)MutableMapping)
health_pb2)health_pb2_grpcc                   @   s   e Zd ZU dZeedf ed< eeejf ed< e	ed< ddd	Z
d
ejddfddZd
ejddfddZdedejjddfddZdedejjddfddZdddZdS )HealthServicerz6An AsyncIO implementation of health checking servicer.z-_health_pb2.HealthCheckResponse.ServingStatus_server_status_server_watchers_gracefully_shutting_downreturnNc                 C   s&   dt jji| _ttj| _d| _	d S )N F)
_health_pb2HealthCheckResponseSERVINGr   collectionsdefaultdictasyncio	Conditionr   r   )self r   I/home/ubuntu/.local/lib/python3.10/site-packages/grpc_health/v1/_async.py__init__"   s   
zHealthServicer.__init__requestc                    s<   | j |j}|d u r|tjjI d H  d S tj|dS )Nstatus)	r   getserviceabortgrpc
StatusCode	NOT_FOUNDr   r   )r   r   contextr   r   r   r   Check'   s
   zHealthServicer.Checkc              	      s   | j |j }d }zG|4 I d H & 	 | j|jtjj}||kr,|tj|dI d H  |}| I d H  q1 I d H s<w   Y  W |j| j v rO| j |j= d S d S |j| j v r\| j |j= w )NTr   )	r   r   r   r   r   r   SERVICE_UNKNOWNwritewait)r   r   r   	conditionlast_statusr   r   r   r   Watch1   s,   

zHealthServicer.Watchr   r   c              	      sv   || j v r4| j |}|4 I d H  || j|< |  W d   I d H  d S 1 I d H s-w   Y  d S || j|< d S )N)r   r   r   
notify_all)r   r   r   r$   r   r   r   _setP   s   


.zHealthServicer._setc                    s"   | j rdS | ||I dH  dS )zSets the status of a service.

        Args:
          service: string, the name of the service.
          status: HealthCheckResponse.status enum value indicating the status of
            the service
        N)r   r(   )r   r   r   r   r   r   set]   s   zHealthServicer.setc                    s8   | j rdS d| _ | jD ]}| |tjjI dH  qdS )a
  Permanently sets the status of all services to NOT_SERVING.

        This should be invoked when the server is entering a graceful shutdown
        period. After this method is invoked, future attempts to set the status
        of a service will be ignored.
        NT)r   r   r(   r   r   NOT_SERVING)r   r   r   r   r   enter_graceful_shutdownm   s   
z&HealthServicer.enter_graceful_shutdown)r	   N)__name__
__module____qualname____doc__r   str__annotations__r   r   boolr   r   HealthCheckRequestr    r&   r   ServingStatusr(   r)   r+   r   r   r   r   r      sB   
 





r   )r/   r   r   typingr   r   grpc_health.v1r   r   r   _health_pb2_grpcr   r   r   r   r   <module>   s   