U
    h$3                     @   s~  d Z ddlZddlZddlZddlZddlZddlmZ ddlm	Z	 z(ddl
mZ ddl
mZ ddl
mZ W n  ek
r   d Z ZZY nX zddlmZ W n ek
r   Y nX dd	lmZ dd
lmZ ddlmZ edZdd ZG dd deZeddZG dd deZdd ZG dd deZG dd deeZG dd deZ G dd de eZ!z"ddl"m#Z# G d d! d!e e#Z$W n ek
r   dZ$Y nX z"dd"l%m&Z& G d#d$ d$ee&Z'W n ek
r   dZ'Y nX G d%d& d&eZ(G d'd( d(e(eZ)z"dd)l*m+Z+ G d*d+ d+e(e+Z,W n ek
r:   dZ,Y nX z"dd,l*m-Z- G d-d. d.e(e-Z.W n ek
rx   dZ.Y nX dS )/aG  
Lightweight connection pooling for peewee.

In a multi-threaded application, up to `max_connections` will be opened. Each
thread (or, if using gevent, greenlet) will have it's own connection.

In a single-threaded application, only one connection will be created. It will
be continually recycled until either it exceeds the stale timeout or is closed
explicitly (using `.manual_close()`).

By default, all your application needs to do is ensure that connections are
closed when you are finished with them, and they will be returned to the pool.
For web applications, this typically means that at the beginning of a request,
you will open a connection, and when you return a response, you will close the
connection.

Simple Postgres pool example code:

    # Use the special postgresql extensions.
    from playhouse.pool import PooledPostgresqlExtDatabase

    db = PooledPostgresqlExtDatabase(
        'my_app',
        max_connections=32,
        stale_timeout=300,  # 5 minutes.
        user='postgres')

    class BaseModel(Model):
        class Meta:
            database = db

That's it!
    N)
namedtuple)chain)TRANSACTION_STATUS_IDLE)TRANSACTION_STATUS_INERROR)TRANSACTION_STATUS_UNKNOWN)TransactionStatus)MySQLDatabase)PostgresqlDatabase)SqliteDatabasezpeewee.poolc                 C   s"   | d k	rt | ttfst| S | S N)
isinstanceintfloat)val r   2/tmp/pip-unpacked-wheel-5j60pwdk/playhouse/pool.pymake_int>   s    r   c                   @   s   e Zd ZdS )MaxConnectionsExceededN__name__
__module____qualname__r   r   r   r   r   D   s    r   PoolConnection)	timestamp
connectionchecked_outc                   @   s   e Zd Zdd ZdS )	_sentinelc                 C   s   dS NTr   )selfotherr   r   r   __lt__K   s    z_sentinel.__lt__N)r   r   r   r    r   r   r   r   r   J   s   r   c                    s   t   fdd}|S )Nc              
      s.   | j   | f||W  5 Q R  S Q R X d S r   )
_pool_lock)r   argskwargsfnr   r   innerP   s    zlocked.<locals>.inner)	functoolswraps)r%   r&   r   r$   r   lockedO   s    r)   c                       s   e Zd Zd fdd	Zd fdd	Zd fdd		Ze fd
dZdd Zdd Z	dd Z
ed  fdd	Zedd Zedd Zed!ddZedd Z  ZS )"PooledDatabase   Nc                    sf   t || _t || _t || _| jdkr2td| _t | _g | _i | _	t
| _tt| j|f| d S Nr   inf)r   _max_connections_stale_timeout_wait_timeoutr   	threadingRLockr!   _connections_in_useidconn_keysuperr*   __init__)r   databasemax_connectionsstale_timeouttimeoutr#   	__class__r   r   r8   X   s    





zPooledDatabase.__init__c                    sb   t t| j|f| |d k	r&t|| _|d k	r8t|| _|d k	r^t|| _| jdkr^td| _d S r,   )r7   r*   initr   r.   r/   r0   r   )r   r9   r:   r;   r<   Zconnect_kwargsr=   r   r   r?   r   s    



zPooledDatabase.initFc                    sv   | j stt| |S t | j  }|t krjztt| |}W n tk
rb   td Y q$X |S q$tdd S )Ng?z:Max connections exceeded, timed out attempting to connect.)r0   r7   r*   connecttimer   sleep)r   Zreuse_if_openexpiresretr=   r   r   r@   ~   s    zPooledDatabase.connectc                    s  z$t | j\}}}|}| |}W n* tk
rN   d  }}td Y qY q X | |rptd| d  }}q | jr| 	|rtd| | 
|d d  }}q qq |d kr| jrt| j| jkrtdtt|  }t }| |}td| t||t | j|< |S )Nz No connection available in pool.zConnection %s was closed.z!Connection %s was stale, closing.TzExceeded maximum connections.zCreated new connection %s.)heapqheappopr3   r6   
IndexErrorloggerdebug
_is_closedr/   	_is_stale_closer.   lenr4   r   r7   r*   _connectrA   r   )r   ts_Zc_connconnkeyr=   r   r   rN      s6    





zPooledDatabase._connectc                 C   s   t   | | jkS r   )rA   r/   )r   r   r   r   r   rK      s    zPooledDatabase._is_stalec                 C   s   dS )NFr   r   rQ   r   r   r   rJ      s    zPooledDatabase._is_closedc                 C   s   dS r   r   rS   r   r   r   
_can_reuse   s    zPooledDatabase._can_reusec                    s   |  |}|r tt| | n|| jkr| j|}| jrf| |jrft	
d| tt| | n<| |rt	
d| t| j|jt |f nt	
d| d S )NzClosing stale connection %s.zReturning %s to pool.z
Closed %s.)r6   r7   r*   rL   r4   popr/   rK   r   rH   rI   rT   rE   heappushr3   r   )r   rQ   
close_connrR   	pool_connr=   r   r   rL      s    


zPooledDatabase._closec                 C   sB   |   rdS |  }| j| |d |   | j|dd dS )zS
        Close the underlying connection without returning it to the pool.
        FNTrW   )Z	is_closedr   r4   rU   r6   closerL   rS   r   r   r   manual_close   s    zPooledDatabase.manual_closec                 C   s*   | j D ]\}}}| j|dd qg | _ d S NTrY   )r3   rL   )r   rP   rQ   r   r   r   
close_idle   s    zPooledDatabase.close_idleX  c                 C   s^   i }t   | }d}| j D ]4\}}|j|k rJ| j|jdd |d7 }q|||< q|| _|S )Nr   TrY      )rA   r4   itemsr   rL   r   )r   ZageZin_usecutoffnrR   rX   r   r   r   close_stale   s    


zPooledDatabase.close_stalec                 C   sX   |    | jD ]\}}}| j|dd q| j D ]}| j|jdd q2g | _i | _d S r\   )rZ   r3   rL   r4   valuesr   )r   rP   rQ   rX   r   r   r   	close_all   s    zPooledDatabase.close_all)r+   NN)NNN)F)F)r^   )r   r   r   r8   r?   r@   r)   rN   rK   rJ   rT   rL   r[   r]   rc   re   __classcell__r   r   r=   r   r*   W   s,       +

r*   c                   @   s   e Zd Zdd ZdS )PooledMySQLDatabasec                 C   s>   | j d dkrd}nd}z|j|  W n   Y dS X dS d S )Nr      r   )FTF)Zserver_versionZping)r   rQ   r"   r   r   r   rJ     s    zPooledMySQLDatabase._is_closedNr   r   r   rJ   r   r   r   r   rg     s   rg   c                   @   s   e Zd Zdd Zdd ZdS )_PooledPostgresqlDatabasec                 C   s2   |j r
dS | }|tkrdS |tkr.|  dS NTF)closedget_transaction_statusr   r   rollbackr   rQ   Z
txn_statusr   r   r   rJ     s    z$_PooledPostgresqlDatabase._is_closedc                 C   s:   |  }|tkrdS |tkr&|  n|tkr6|  dS NFT)rm   r   r   resetr   rn   ro   r   r   r   rT   '  s    
z$_PooledPostgresqlDatabase._can_reuseNr   r   r   rJ   rT   r   r   r   r   rj     s   rj   c                   @   s   e Zd ZdS )PooledPostgresqlDatabaseNr   r   r   r   r   rs   4  s   rs   )PostgresqlExtDatabasec                   @   s   e Zd ZdS )PooledPostgresqlExtDatabaseNr   r   r   r   r   ru   :  s   ru   )Psycopg3Databasec                   @   s   e Zd Zdd Zdd ZdS )PooledPsycopg3Databasec                 C   s6   |j r
dS |jj}|tjkr dS |tjkr2|  dS rk   )rl   pgconntransaction_statusr   UNKNOWNIDLErn   ro   r   r   r   rJ   D  s    

z!PooledPsycopg3Database._is_closedc                 C   s@   |j j}|tjkrdS |tjkr*|  n|tjkr<|  dS rp   )rx   ry   r   rz   ZINERRORrq   r{   rn   ro   r   r   r   rT   O  s    



z!PooledPsycopg3Database._can_reuseNrr   r   r   r   r   rw   C  s   rw   c                   @   s   e Zd Zdd ZdS )_PooledSqliteDatabasec                 C   s"   z
|j  W n   Y dS X dS d S rk   )Ztotal_changesrS   r   r   r   rJ   `  s
    
z _PooledSqliteDatabase._is_closedNri   r   r   r   r   r|   _  s   r|   c                   @   s   e Zd ZdS )PooledSqliteDatabaseNr   r   r   r   r   r}   h  s   r}   )SqliteExtDatabasec                   @   s   e Zd ZdS )PooledSqliteExtDatabaseNr   r   r   r   r   r   n  s   r   )CSqliteExtDatabasec                   @   s   e Zd ZdS )PooledCSqliteExtDatabaseNr   r   r   r   r   r   v  s   r   )/__doc__r'   rE   loggingr1   rA   collectionsr   	itertoolsr   Zpsycopg2.extensionsr   r   r   ImportErrorZ
psycopg.pqr   Zpeeweer   r	   r
   	getLoggerrH   r   
ValueErrorr   r   objectr   r)   r*   rg   rj   rs   Zplayhouse.postgres_extrt   ru   Zplayhouse.psycopg3_extrv   rw   r|   r}   Zplayhouse.sqlite_extr~   r   r   r   r   r   r   r   <module>   sr   !

 7

	
