o
    [h5                     @   sN  d dl 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dlm	Z	 ddl
mZ ddlmZ G d	d
 d
eZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd  d e d g d!ZG d"d# d#eZejd$d% ZdS )&    N   )event)util)url)DefaultDialect)_distill_params)_DDLCompilesc                   @   s(   e Zd ZdZdZdZdd Zdd ZdS )
AssertRuleFNTc                 C   s   d S N selfexecute_observedr   r   /home/ubuntu/experiments/live_experiments/Pythonexperiments/Otree/venv/lib/python3.10/site-packages/sqlalchemy/testing/assertsql.pyprocess_statement   s   zAssertRule.process_statementc                 C   s   J d)NFz?All statements are complete, but pending assertion rules remainr   r   r   r   r   no_more_statements   s   zAssertRule.no_more_statements)__name__
__module____qualname__is_consumederrormessageconsume_statementr   r   r   r   r   r   r	      s    r	   c                   @      e Zd ZdS )SQLMatchRuleNr   r   r   r   r   r   r   r   $   s    r   c                   @   s"   e Zd ZdZdddZdd ZdS )	CursorSQLFNc                 C   s   || _ || _d S r
   )	statementparams)r   r   r   r   r   r   __init__+      
zCursorSQL.__init__c                 C   sp   |j d }| j|jks| jd ur%| j|jkr%d| j| j|j|jf | _d S |j d d| _|j s6d| _d S d S )Nr   z5Testing for exact SQL %s parameters %s received %s %sT)
statementsr   r   
parametersr   popr   r   )r   r   stmtr   r   r   r   /   s    


zCursorSQL.process_statementr
   )r   r   r   r   r   r   r   r   r   r   r   (   s    
r   c                   @   sF   e Zd ZdddZdd Zdd Zd	d
 Zdd Zdd Zdd Z	dS )CompiledSQLNdefaultc                 C   s   || _ || _|| _d S r
   )r   r   dialect)r   r   r   r'   r   r   r   r   E   s   
zCompiledSQL.__init__c                 C   s   t dd| j}||kS N[\n\t] )resubr   )r   r   received_statementr$   r   r   r   _compare_sqlJ   s   zCompiledSQL._compare_sqlc                 C   sB   | j dkrt S | j dkrddi}ni }t| j  di |S )Nr&   
postgresqlZimplicit_returningTr   )r'   r   r   URLget_dialect)r   r   r   r   r   r   _compile_dialectN   s   


zCompiledSQL._compile_dialectc                    s   |j }| |}t|jjtr|jjj||jdd n|jjj||jj	|jj
|jdd tddt }|j}|sH  g}||fS  fdd|D }||fS )zxreconstruct the statement and params in terms
        of a target dialect, which for CompiledSQL is just DefaultDialect.schema_translate_map)r'   r3   )r'   column_keysinliner3   r)   r*   c                    s   g | ]}  |qS r   )construct_params).0mcompiledr   r   
<listcomp>u   s    
z3CompiledSQL._received_statement.<locals>.<listcomp>)contextr2   
isinstancer:   r   r   compileZexecution_optionsgetr4   r5   r+   r,   r   	text_typer"   r6   )r   r   r<   Zcompare_dialect_received_statementr"   _received_parametersr   r9   r   rA   Y   s4   


zCompiledSQL._received_statementc                 C   s   |j }| |\}}| |}| ||}|r^|d ur^t|}t|}|rX|rXt|d}	tt|D ]\}
}|	D ]}||vsH|| |	| krJ nq:||
=  nq4d}n|rX|s'|s\|r^d}|rhd| _d | _	d S | 
|||d | _	d S )Nr   FT)r-   Zreceived_parameters)r<   rA   _all_paramsr.   listdictr#   	enumerater   r   _failure_message)r   r   r<   rA   rB   r   Z
equivalent
all_paramsZall_receivedparamidxreceivedZ	param_keyr   r   r   r   {   sB   
	
zCompiledSQL.process_statementc                 C   s<   | j rt| j r|  |}n| j }t|ts|g}|S d S r
   )r   r   callabler=   rD   )r   r<   r   r   r   r   rC      s   
zCompiledSQL._all_paramsc                 C   "   d| j ddt|ddf S )NzzTesting for compiled statement %r partial params %s, received %%(received_statement)r with params %%(received_parameters)r%%%)r   replacereprr   Zexpected_paramsr   r   r   rG         zCompiledSQL._failure_messageNr&   )
r   r   r   r   r.   r2   rA   r   rC   rG   r   r   r   r   r%   D   s    
"1r%   c                   @   s&   e Zd Zd	ddZdd Zdd ZdS )
RegexSQLNr&   c                 C   s,   t |  t|| _|| _|| _|| _d S r
   )r   r   r+   r>   regex
orig_regexr   r'   )r   rV   r   r'   r   r   r   r      s
   

zRegexSQL.__init__c                 C   rM   )Nz{Testing for compiled statement ~%r partial params %s, received %%(received_statement)r with params %%(received_parameters)rrN   rO   )rW   rP   rQ   rR   r   r   r   rG      rS   zRegexSQL._failure_messagec                 C   s   t | j|S r
   )boolrV   match)r   r   r-   r   r   r   r.      s   zRegexSQL._compare_sqlrT   )r   r   r   r   rG   r.   r   r   r   r   rU      s    
rU   c                       s4   e Zd Zdd Zdd Z fddZdd Z  ZS )	
DialectSQLc                 C   s   |j jS r
   )r<   r'   r   r   r   r   r2      s   zDialectSQL._compile_dialectc                 C   s   t dd|}||kS r(   )r+   r,   )r   	real_stmtreceived_stmtr$   r   r   r   _compare_no_space   s   zDialectSQL._compare_no_spacec                    sJ   t t| |\}}|jD ]}| |j|r nqtd| ||jjfS )NzICan't locate compiled statement %r in list of statements actually invoked)	superrZ   rA   r!   r]   r   AssertionErrorr<   Zcompiled_parameters)r   r   r\   Zreceived_paramsr[   	__class__r   r   rA      s   
zDialectSQL._received_statementc                 C   sz   t dd| j}|jjj}|dkrt dd|}||kS d }|dkr%d}n|dkr,d	}n|d
kr2d }t d||}||kS )Nr)   r*   Zpyformatz	:([\w_]+)z%(\1)sZqmark?formatz%snumeric)r+   r,   r   r<   r'   
paramstyle)r   r   r-   r$   re   replr   r   r   r.      s   
zDialectSQL._compare_sql)r   r   r   r2   r]   rA   r.   __classcell__r   r   r`   r   rZ      s
    rZ   c                   @   $   e Zd Zdd Zdd Zdd ZdS )CountStatementsc                 C   s   || _ d| _d S Nr   count_statement_count)r   rl   r   r   r   r   
  r    zCountStatements.__init__c                 C   s   |  j d7  _ d S )N   )rm   r   r   r   r   r     s   z!CountStatements.process_statementc                 C   s$   | j | jkrJ d| j | jf d S )NFz,desired statement count %d does not match %drk   r   r   r   r   r     s   z"CountStatements.no_more_statementsN)r   r   r   r   r   r   r   r   r   r   ri   	  s    ri   c                   @   s   e Zd Zdd Zdd ZdS )AllOfc                 G      t || _d S r
   )setrulesr   rr   r   r   r   r        zAllOf.__init__c                 C   sl   t | jD ]%}d |_|| |jr!| j| | jsd| _ d S |js*d | _ d S qt | jd j| _d S NTr   )rD   rr   r   r   r   discardr   r   ruler   r   r   r     s   
zAllOf.process_statementN)r   r   r   r   r   r   r   r   r   ro     s    ro   c                       s,   e Zd Zdd Zdd Z fddZ  ZS )EachOfc                 G   rp   r
   )rD   rr   rs   r   r   r   r   /  rt   zEachOf.__init__c                 C   s^   | j r%| j d }|| |jr| j d n|jr|j| _|jr"n| j s| j s-d| _d S d S )Nr   T)rr   r   r   r#   r   r   rw   r   r   r   r   2  s   



zEachOf.process_statementc                    s@   | j r| j d js| j d   d S | j rtt|   d S d S rj   )rr   r   r   r^   ry   r   r`   r   r   r   @  s
   zEachOf.no_more_statements)r   r   r   r   r   r   rg   r   r   r`   r   ry   .  s    ry   c                   @      e Zd Zdd ZdS )Orc                 C   s>   | j D ]}|| |jrd| _ d S qt| j d j| _d S ru   )rr   r   r   rD   r   rw   r   r   r   r   H  s   

zOr.process_statementN)r   r   r   r   r   r   r   r   r{   G      r{   c                   @   rz   )SQLExecuteObservedc                 C   s"   || _ || _t||| _g | _d S r
   )r<   clauseelementr   r"   r!   )r   r<   r~   multiparamsr   r   r   r   r   S  s   
zSQLExecuteObserved.__init__N)r   r   r   r   r   r   r   r   r}   R  r|   r}   c                   @   r   )SQLCursorExecuteObservedNr   r   r   r   r   r   Z  s    r   )r   r"   r<   executemanyc                   @   rh   )SQLAsserterc                 C   s
   g | _ d S r
   )accumulatedr   r   r   r   r   d  s   
zSQLAsserter.__init__c                 C   s   | j | _| ` d S r
   )r   _finalr   r   r   r   _closeg  s   zSQLAsserter._closec                 G   sh   t | }t| j}|r#|d}|| |jrn
|jr!J |j|s|r)J d|js2|  d S d S )Nr   Fz Additional SQL statements remain)ry   rD   r   r#   r   r   r   r   )r   rr   rx   Zobservedr   r   r   r   assert_k  s   



zSQLAsserter.assert_N)r   r   r   r   r   r   r   r   r   r   r   c  s    r   c                 #   s    t   g t| dfdd}t| d fdd}z V  W t| d| t| d|    d S t| d| t| d|    w )NZbefore_executec                    s   |||f d d < d S r
   r   )connr~   r   r   )origr   r   connection_execute  s   z)assert_engine.<locals>.connection_executeZafter_cursor_executec                    sj   |sd S  j r j d j|u r j d }nt|d d d } j | |jt|||| d S )Nr   rn   r   )r   r<   r}   appendr!   r   )r   cursorr   r"   r<   r   ZobsZasserterr   r   r   cursor_execute  s   z%assert_engine.<locals>.cursor_execute)r   r   Zlistens_forremover   )enginer   r   r   r   r   assert_engine|  s   


r   )collections
contextlibr+   r*   r   r   r   r   Zengine.defaultr   Zengine.utilr   Zschemar   objectr	   r   r   r%   rU   rZ   ri   ro   ry   r{   r}   
namedtupler   r   contextmanagerr   r   r   r   r   <module>   s<    .
	