
    j                        d dl mZ d dlZd dlZd dlmZ d dlmZ d dlm	Z	 dZ
dZdZd"dZd#dZd$dZd%dZd&dZd'dZd(dZdddde	dd)d!ZdS )*    )annotationsN)Any)
task_queue)HEALTH_STALE_TASK_SECONDS<      valuestr | int | Nonereturnintc                    	 t          | pt                    }n# t          t          f$ r
 t          }Y nw xY wt	          dt          |t                              S )N   )r   DEFAULT_TASK_LIMIT	TypeError
ValueErrormaxminMAX_TASK_LIMIT)r	   limits     ?/home/longshao/multi-rider-rag/shared/task_queue_diagnostics.pynormalize_task_limitr      s_    #E//00z" # # #"#q#e^,,---s    44
str | Nonestrc                H    t          | pd                                          S )N )r   strip)r	   s    r   _clean_filterr      s     u{!!###    r   c                    t          | pd                                          }|sdS t          |          dk    rdS t          |          dk    r|d d          d|dd           S |d d          d|dd           S )	Nr      z***
      z...)r   r   len)r	   texts     r   _mask_ownerr'      s    u{!!##D r
4yyA~~u
4yyBrr(**tBCCy***2A2h&&49&&&r   dict[str, Any]c                    	 t          j        | pd          }n# t          $ r i cY S w xY wt          |t                    r|ni S )Nz{})jsonloads	Exception
isinstancedict)r	   payloads     r   _parse_payloadr0   '   sY    *U]d++   			 $//777R7s    ((start_tsend_ts
int | Nonec                    	 t          | pd          }t          |pd          }n# t          t          f$ r Y d S w xY w|r|sd S t          d||z
            S Nr   r   r   r   r   )r1   r2   startends       r   _duration_betweenr9   /   sz    HM""&+Az"   tt  tq#+s   "% ::now_tsc                    	 t          |pd          }n# t          t          f$ r Y d S w xY w|sd S t          d| |z
            S r5   r6   )r:   r1   r7   s      r   _duration_sincer<   :   sb    HM""z"   tt tq&5.!!!s    ))rowstale_after_secondsc                   t          | d                   }t          | d         pd          }| d         }| d         }| d         }|dk    }|o| pt          |          ||z
  k     }	|rt          ||          nt	          ||          }
|rt	          ||          nt          ||          }|rt          ||          nt	          ||          }| d         | d	         |t          |                    d
          pd          t          | d                   t          | d                   ||||
||t          | d         pd          |	t          | d         pd          d t                   dS )Nr/   statusr   
claimed_tsfinished_ts
created_tsrunningid	task_typejob_id	owner_keyowner_ipretriesr   error)task_idrF   r@   rG   rH   rI   rC   rA   rB   wait_secondsrun_secondstotal_secondsrJ   stalerK   )r0   r   r   r9   r<   getr'   ERROR_PREVIEW_CHARS)r=   r:   r>   r/   r@   rA   rB   rC   
is_runningis_stalerM   rN   rO   s                r   _serialize_taskrU   D   s   S^,,GX$"%%F\"Jm$K\"J9$J`:~_Z6L_C_1_H 	1*j111VZ00  	8
+++z;77  	1*k222VZ00  t9%gkk(++1r22 [!122J00  "$"&s9~*++S\'R(()=*=)=>  r   )rF   r@   r   r:   r>   rF   r@   r   c                   t          pt          j                              t          |          }t          |           }t          |          }z
  }g }	g }
|r*|	                    d           |
                    |           |r*|	                    d           |
                    |           |	rdd                    |	           nd}t          j                    5 }|                    d          	                                }|                    d          
                                }|                    d          
                                }|                    d	|f          	                                }|                    d
| dg |
|          
                                }d d d            n# 1 swxY w Y   d |D             }|||dt          ||d         nd          |                    dd          |                    dd          |                    dd          |                    dd          t          ||d         nd          dd |D             d |D             fd|D             dS )Nztask_type=?zstatus=?zWHERE z AND r   z(SELECT COUNT(*) AS total FROM task_queuez
            SELECT status, COUNT(*) AS count
            FROM task_queue
            GROUP BY status
            ORDER BY status ASC
            z
            SELECT task_type, status, COUNT(*) AS count
            FROM task_queue
            GROUP BY task_type, status
            ORDER BY task_type ASC, status ASC
            z
            SELECT COUNT(*) AS total
            FROM task_queue
            WHERE status='running'
              AND (claimed_ts IS NULL OR claimed_ts < ?)
            z>
            SELECT *
            FROM task_queue
            zO
            ORDER BY created_ts DESC, id DESC
            LIMIT ?
            c                J    i | ] }|d          t          |d         pd          !S )r@   countr   r   .0r=   s     r   
<dictcomp>z+get_task_queue_snapshot.<locals>.<dictcomp>   s/    RRRsS]CG(9$:$:RRRr   )rF   r@   r   totalr   pendingrD   	completedfailed)r]   r^   rD   r_   r`   stale_runningc                N    g | ]"}|d          t          |d         pd          d#S )r@   rX   r   )r@   rX   rY   rZ   s     r   
<listcomp>z+get_task_queue_snapshot.<locals>.<listcomp>   sE     
 
 
 8}s3w<3D1/E/EFF
 
 
r   c                \    g | ])}|d          |d         t          |d         pd          d*S )rF   r@   rX   r   )rF   r@   rX   rY   rZ   s     r   rc   z+get_task_queue_snapshot.<locals>.<listcomp>   sT     
 
 
 	 !-h-S\.Q// 
 
 
r   c                2    g | ]}t          |          S  )rU   )r[   r=   r:   r>   s     r   rc   z+get_task_queue_snapshot.<locals>.<listcomp>   s&    YYY/#v/BCCYYYr   )generated_tsfiltersr>   totals	by_statusby_type_statustasks)r   timer   r   appendjoinr   _connectexecutefetchonefetchallrQ   )rF   r@   r   r:   r>   
safe_limitsafe_task_typesafe_statuscutoffclausesparams	where_sqlconn	total_rowstatus_rowstype_status_rows	stale_row	task_rowsstatus_countss      ``              r   get_task_queue_snapshotr   p   s@    &49;;''F%e,,J"9--N''K))FGF &}%%%n%%% #z"""k"""4;C0g..000I				 $$LL!KLLUUWW	ll
 
 (** 	  <<
 
 (** 	 LL I
 
 (** 	 LL    "f!j!	
 	
 (** 	7$ $ $ $ $ $ $ $ $ $ $ $ $ $ $L SRkRRRM'!
 

  3y/D7++!LL$((A66$((A66&**;::#''!44 y7L7!3!3RSTT
 

 
"
 
 

 
 (
 
 
 ZYYYYyYYY9  s   "CF==GG)r	   r
   r   r   )r	   r   r   r   )r	   r   r   r   )r	   r   r   r(   )r1   r   r2   r   r   r3   )r:   r   r1   r   r   r3   )r=   r   r:   r   r>   r   r   r(   )rF   r   r@   r   r   r
   r:   r3   r>   r   r   r(   )
__future__r   r*   rm   typingr   sharedr   shared.healthr   r   r   rR   r   r   r'   r0   r9   r<   rU   r   rf   r   r   <module>r      s@   " " " " " "               3 3 3 3 3 3   . . . .$ $ $ $' ' ' '8 8 8 8   " " " ") ) ) )\ !"8\ \ \ \ \ \ \ \r   