
    & j#                        S SK r S SKrS SKrS SK7  S SKJr  S SKJr  S SKJr  S SKJr  S SKJr  S SKJ	r	  S S	KJ
r
  S S
KJr  S SKJr   S SKJr  S SKJr  S SKJr  \rSrSr " S S\5      r " S S\5      r " S S\5      r " S S\5      r " S S\5      r  SS jr " S S\\5      r g! \ a    S=r=r=rr N_f = f)    N)*)_atomic)_manual)ColumnMetadata)EnclosedNodeList)Entity)ForeignKeyMetadata)IndexMetadata)NodeList)_PooledPostgresqlDatabase)
ArrayField)BinaryJSONField)IntervalFieldi a  CockroachDB does not support nested transactions. You may alternatively use the @transaction context-manager/decorator, which only wraps the outer-most block in transactional logic. To run a transaction with automatic retries, use the run_transaction() helper.c                       \ rS rSrSrg)ExceededMaxAttempts     N__name__
__module____qualname____firstlineno____static_attributes__r       f/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/playhouse/cockroachdb.pyr   r       s    Tr   r   c                   ,   ^  \ rS rSrSrU 4S jrSrU =r$ )UUIDKeyField#   Tc                    > UR                  S5      (       a  [        S[        U 5      -  5      e[        S5      /US'   UR	                  SS5        [
        [        U ]  " U0 UD6  g )Nconstraints%s cannot specify constraints.zDEFAULT gen_random_uuid()primary_keyT)get
ValueErrortypeSQL
setdefaultsuperr   __init__selfargskwargs	__class__s      r   r)   UUIDKeyField.__init__&   s^    ::m$$=T
JKK!$%@!A B}-.lD*D;F;r   r   )r   r   r   r   auto_incrementr)   r   __classcell__r.   s   @r   r   r   #   s    N< <r   r   c                   ,   ^  \ rS rSrSrU 4S jrSrU =r$ )
RowIDField.   INTc                    > UR                  S5      (       a  [        S[        U 5      -  5      e[        S5      /US'   [        [
        U ]  " U0 UD6  g )Nr    r!   zDEFAULT unique_rowid())r#   r$   r%   r&   r(   r4   r)   r*   s      r   r)   RowIDField.__init__1   sN    ::m$$=T
JKK!$%=!> ?}j$($9&9r   r   )r   r   r   r   
field_typer)   r   r1   r2   s   @r   r4   r4   .   s    J: :r   r4   c                     ^  \ rS rSr\R
                  R                  5       r\R                  SS05        SrU 4S jr	U 4S jr
SS jrSU 4S jjrS	 rU 4S
 jrS rS rSU 4S jjrSU 4S jjrU 4S jr  SS jr  SS jrSrU =r$ )CockroachDatabase8   BLOBBYTESTc                    > SU;  aA  U(       a:  UR                  S5      (       d$  UR                  SS5        UR                  SS5        [        [        U ]  " U/UQ70 UD6  g )Ndsnzpostgresql://userrootportif  )
startswithr'   r(   r;   r)   )r+   databaser,   r-   r.   s       r   r)   CockroachDatabase.__init__@   sZ     H(0(;(;O(L(Lff-fe,/J4J6Jr   c                 .  > UR                  5       nUR                  S5        UR                  5       u  n[        R                  " SU5      nUb4  S[        S UR                  5        5       5      -  n[        U5      U l        g [        [        U ]/  U5        g )Nzselect version()z#^CockroachDB.+?v(\d+)\.(\d+)\.(\d+)z
%d%02d%02dc              3   8   #    U  H  n[        U5      v   M     g 7fN)int).0is     r   	<genexpr>8CockroachDatabase._set_server_version.<locals>.<genexpr>O   s     (L9KAQ9Ks   )cursorexecutefetchonerematchtuplegroupsrJ   server_versionr(   r;   _set_server_version)r+   conncursraw	match_objcleanr.   s         r   rW   %CockroachDatabase._set_server_versionI   s{    {{}'(}}HHCSI	  5(L9I9I9K(L#LLE"%e*D #T>tDr   c                     SnU R                  X1U=(       d    SS45      nUR                  5       nU=(       a    US   =(       d    S $ )NzSELECT constraint_name FROM information_schema.table_constraints WHERE table_name = %s AND table_schema = %s AND constraint_type = %spubliczPRIMARY KEYr   )execute_sqlrQ   )r+   tableschemaqueryrO   rows         r   _get_pk_constraint$CockroachDatabase._get_pk_constraintU   sK    , !!%1C8*7*9 :oo~s1v%%r   c                    > [         [        U ]  X5      nU R                  X5      nU Vs/ s H  oT(       a  UR                  U:w  d  M  UPM     sn$ s  snf rI   )r(   r;   get_indexesre   name)r+   ra   rb   indexespkcidxr.   s         r   rh   CockroachDatabase.get_indexes_   sK     )4<UK%%e4&IwsCwIIIs   AAc                     UR                   (       d  g UR                   R                  5       nUS;   a  [        S5      $ US;  a  [        S5      eg )NreplaceupsertUPSERT)ignorenothingupdatezfUn-supported action for conflict resolution. CockroachDB supports REPLACE (UPSERT), IGNORE and UPDATE.)_actionlowerr&   r$   )r+   on_conflictrc   actions       r   conflict_statement$CockroachDatabase.conflict_statementf   sT    ""F$$**,**x= :: + , , ;r   c                   > UR                   (       a  UR                   R                  5       OSnUS;   a  [        S5      /nUR                  (       aU  UR	                  [        UR                   Vs/ s H%  n[        U[        5      (       a  [        U5      OUPM'     sn5      5        UR	                  [        S5      5        [        U5      $ US;   a  g UR                  (       a  [        S5      e[        [        U ];  X5      $ s  snf )N )rs   rt   zON CONFLICTz
DO NOTHINGro   zWCockroachDB does not support the usage of a constraint name. Use the column(s) instead.)rv   rw   r&   _conflict_targetappendr   
isinstancestrr   r   _conflict_constraintr$   r(   r;   conflict_update)r+   ocrc   ry   partscolr.   s         r   r   !CockroachDatabase.conflict_updateq   s    ')zz!!#r**'(E""-!22/42 $.c3#7#7F3KS@2/4 5 6 LL\*+E?",, $$ K L L &=bHH/4s   0,Dc                 ,    [         R                  X5      $ rI   )fnextract)r+   	date_part
date_fields      r   extract_dateCockroachDatabase.extract_date   s    zz)00r   c                 B    UR                  S5      R                  S5      $ )NrJ   timestamptz)cast)r+   r   s     r   from_timestamp CockroachDatabase.from_timestamp   s     u%**=99r   c                    > [         [        U ]  5         Ub!  U R                  5       R	                  SU45        UbD  UR                  5       nUS;  a  [        S5      eU R                  5       R	                  SU-  5        g g )Nz$SET TRANSACTION AS OF SYSTEM TIME %s)lownormalhighz$priority must be low, normal or highzSET TRANSACTION PRIORITY %s)r(   r;   beginrO   rP   rw   r$   r+   system_timepriorityr.   s      r   r   CockroachDatabase.begin   sy    ,."KKM!!"H#..2~~'H88 !GHHKKM!!"?("JK	  r   c                    > U R                  5       (       a  U R                  5         U R                  [        :  a  [	        XU5      $ [
        [        U ]  X5      $ rI   )	is_closedconnectrV   NESTED_TX_MIN_VERSION_crdb_atomicr(   r;   atomicr   s      r   r   CockroachDatabase.atomic   sG    >>T\\^!668<<&4[KKr   c                    > U R                  5       (       a  U R                  5         U R                  [        :  a  [	        [
        5      e[        [        U ]#  5       $ rI   )	r   r   rV   r   NotImplementedErrorTXN_ERR_MSGr(   r;   	savepointr+   r.   s    r   r   CockroachDatabase.savepoint   sC    >>T\\^!66%k22&799r   c                     ^ ^^^ UUU U4S jnU$ )Nc                 P   >^  [         R                  " T 5      U UUUU4S j5       nU$ )Nc                  "   > [        TT TTT5      $ rI   run_transaction)cbmax_attemptsr   r+   r   s   r   new_fnACockroachDatabase.retry_transaction.<locals>.deco.<locals>.new_fn   s    &tR{'/1 1r   )	functoolswraps)r   r   r   r   r+   r   s   ` r   deco1CockroachDatabase.retry_transaction.<locals>.deco   s*    __R 1 1 !1 Mr   r   )r+   r   r   r   r   s   ```` r   retry_transaction#CockroachDatabase.retry_transaction   s    	 	 r   c                     [        XX#U5      $ rI   r   )r+   r   r   r   r   s        r   r   !CockroachDatabase.run_transaction   s    tHMMr   )rV   rI   )NNNNN)r   r   r   r   PostgresqlDatabasefield_typescopyru   release_after_rollbackr)   rW   re   rh   rz   r   r   r   r   r   r   r   r   r   r1   r2   s   @r   r;   r;   8   s    $00557K  "K
E&J	,I(1:
	LL: @D#' BF!%N Nr   r;   c                   (   ^  \ rS rSrU 4S jrSrU =r$ )r      c                    > U R                   R                  5       S:  a<  [        U R                   R                  5       [        5      (       d  [        [        5      e[        [        U ]'  5       $ )Nr   )
dbtransaction_depthr   top_transactionr   r   r   r(   r   	__enter__r   s    r   r   _crdb_atomic.__enter__   sM    77$$&*dgg557AA)+66\4244r   r   )r   r   r   r   r   r   r1   r2   s   @r   r   r      s    5 5r   r   c                    U=(       d    SnU R                  X4S9 nU R                  S5        US:w  a%   U" U 5      nU R                  S5        UsSSS5        $  SSS5        [        SS
5      e! [         a<  nUR                  R                  S:X  a  US-  nU R                  S	5         SnAM  e SnAff = f! , (       d  f       Nc= f)a  
Run transactional SQL in a transaction with automatic retries.

User-provided `callback`:
* Must accept one parameter, the `db` instance representing the connection
  the transaction is running under.
* Must not attempt to commit, rollback or otherwise manage transactions.
* May be called more than once.
* Should ideally only contain SQL operations.

Additionally, the database must not have any open transaction at the time
this function is called, as CRDB does not support nested transactions.
)r   r   zSAVEPOINT cockroach_restartr   z#RELEASE SAVEPOINT cockroach_restartN40001   z'ROLLBACK TO SAVEPOINT cockroach_restartzunable to commit transaction)r   r`   OperationalErrororigpgcoder   )r   callbackr   r   r   txnresultexcs           r   r   r      s      %2L	{	>#
45a	!"DE 
?	> 
? d$B
CC $ 88??g- A%LNN#LM 
?	>s4   B6A--
B370B.'B6-B..B33B66
Cc                       \ rS rSrSrg)PooledCockroachDatabase   r   Nr   r   r   r   r   r      s    r   r   r   )!r   rR   syspeeweer   r   r   r   r   r	   r
   r   playhouse.poolr   playhouse.postgres_extr   r   r   	JSONFieldImportErrorr   r   r   r   	UUIDFieldr   	AutoFieldr4   r   r;   r   r   r   r   r   r   <module>r      s     	 
    ! #  %    4D164I
  + 2* 1<9 <: :xN* xNv57 5 BF!D@	79J 	O  D?CCJCC=9Ds   	B, ,B=<B=