a
    g	g?                     @  s~  d dl m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 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 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 ddlmZ ddlmZ ddlm Z  er^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/eed*d+eed*d+eed*d+eed*d+d,d- Z0ee-d*d+d%d.d/d0d1d2Z1ee/d*d+d)d.d/d0d3d4Z2ee.d*d+d'd.d/d0d5d6Z3d.d7d8d9d7d:d/d;d<d=Z4ee	j%d*d+d>d.d/d0d?d@Z5dS )A    )annotationsN)Any)Optional)TYPE_CHECKING)Unionschema)types)compiles   )alter_table)AlterColumn)ColumnDefault)
ColumnName)ColumnNullable)
ColumnType)format_column_name)format_server_default)DefaultImpl   )util)compare)sqla_compat)_is_mariadb)_is_type_bound)Literal)MySQLDDLCompiler)DropConstraint)
Constraint)
TypeEngine)_ServerDefaultc                      s   e Zd ZdZdZejddhddhf ZddgZd$d
d
ddddddddddddddd fddZddd fddZ	ddddddZ
dd Zd d! Zd"d# Z  ZS )%	MySQLImplmysqlFZBOOLZTINYINTZJSONZLONGTEXTzcharacter set ([\w\-_]+)zcollate ([\w\-_]+)NstrOptional[bool]z%Union[_ServerDefault, Literal[False]]Optional[str]Optional[TypeEngine]zOptional[_ServerDefault]$Optional[Union[str, Literal[False]]]r   None)
table_namecolumn_namenullableserver_defaultnametype_r   existing_typeexisting_server_defaultexisting_nullableautoincrementexisting_autoincrementcommentexisting_commentkwreturnc                   s  t ||	st ||	r>t j||f|||||
||	d| |d us^| |d urV|n||r| t||||d urv|n||d ur|n|
d ur|
nd|d ur|n||dur|n|	|d ur|n||dur|n|d	 n|d us|d us|d us|durz| t||||d ur|n||d ur |n|
d ur.|
nd|d ur>|n||durN|n|	|d ur^|n||durn|n|d	 n |dur| t	||||d d S )N)r+   r.   r   r/   r1   r,   r0   TF)r   newnamer+   r.   defaultr2   r4   r   )
r   Z_server_default_is_identityZ_server_default_is_computedsuperalter_column$_is_mysql_allowed_functional_default_execMySQLChangeColumnMySQLModifyColumnMySQLAlterDefault)selfr)   r*   r+   r,   r-   r.   r   r/   r0   r1   r2   r3   r4   r5   r6   	__class__ T/var/www/html/llm_bihealth/app/venv/lib/python3.9/site-packages/alembic/ddl/mysql.pyr;   1   s    

zMySQLImpl.alter_columnr   )constr7   c                   s(   t |tjrt|rd S t | d S N)
isinstancer   CheckConstraintr   r:   drop_constraint)rA   rF   rB   rD   rE   rJ      s    zMySQLImpl.drop_constraintbool)r.   r,   r7   c                 C  s   |d uo|j tju o|d uS rG   )_type_affinitysqltypesDateTime)rA   r.   r,   rD   rD   rE   r<      s    z.MySQLImpl._is_mysql_allowed_functional_defaultc                 C  s   |j jtju r*|jr*|js*|s*|dkr*dS |j jtju rZ|d urNtdd|nd }||kS |r|rtd|	 }td|	 }|r|sdS |
d|
dkrdS |
d}|
d}td	d
|	 td	d
|	 kS ||kS d S )Nz'0'Fz^'|'$ z(.*) (on update.*?)(?:\(\))?$Tr   r   z(.*?)(?:\(\))?$z\1)typerL   rM   ZIntegerZprimary_keyr2   resubmatchlowergroup)rA   Zinspector_columnZmetadata_columnZrendered_metadata_defaultZrendered_inspector_defaultZonupdate_insZonupdate_metrD   rD   rE   compare_server_default   sR    



z MySQLImpl.compare_server_defaultc           	      C  s   t  }t|D ]~}|jrq|jD ]j}|j|jkrJ|| ||j  q|jD ]*}|j|jkrP|| ||j  q|qP|j|v r  qq qt|D ]}|j|v r|| qd S rG   )setlistuniquecolumnsr-   removeaddZforeign_keys)	rA   Zconn_unique_constraintsZconn_indexesZmetadata_unique_constraintsZmetadata_indexesremovedidxcolfkrD   rD   rE   correct_for_autogen_constraints   s&    






z)MySQLImpl.correct_for_autogen_constraintsc                 C  s   dd |D }dd |D }t ||D ]d}|| }|| }|jd urf|j dkrf|jd u rfd|_|jd ur*|j dkr*|jd u r*d|_q*d S )Nc                 S  s   i | ]}t |j|qS rD   r   Z_fk_constraint_sigsig.0r`   rD   rD   rE   
<dictcomp>
  s   z=MySQLImpl.correct_for_autogen_foreignkeys.<locals>.<dictcomp>c                 S  s   i | ]}t |j|qS rD   rb   rd   rD   rD   rE   rf     s   ZrestrictZRESTRICT)rW   intersectionZondeleterT   Zonupdate)rA   Zconn_fksZmetadata_fksZconn_fk_by_sigZmetadata_fk_by_sigrc   ZmdfkZcnfkrD   rD   rE   correct_for_autogen_foreignkeys	  s.    z)MySQLImpl.correct_for_autogen_foreignkeys)NFNNNNNNNNFN)__name__
__module____qualname____dialect__Ztransactional_ddlr   Ztype_synonymsZtype_arg_extractr;   rJ   r<   rV   ra   rh   __classcell__rD   rD   rB   rE   r!   '   s2               2c	:&r!   c                   @  s   e Zd ZdZdS )MariaDBImplmariadbN)ri   rj   rk   rl   rD   rD   rD   rE   rn   %  s   rn   c                      s,   e Zd Zd	dddddd fddZ  ZS )
r@   Nr#   r    r%   r(   )r-   r*   r9   r   r7   c                   s$   t t| j||d || _|| _d S )Nr   )r:   r   __init__r*   r9   )rA   r-   r*   r9   r   rB   rD   rE   rp   *  s    zMySQLAlterDefault.__init__)Nri   rj   rk   rp   rm   rD   rD   rB   rE   r@   )  s    r@   c                      s6   e Zd Zddddddddddd	d

 fddZ  ZS )r>   NFr#   r%   r&   r$   /Optional[Union[_ServerDefault, Literal[False]]]r'   r(   )
r-   r*   r   r8   r.   r+   r9   r2   r4   r7   c
           
        sZ   t t| j||d || _|| _|| _|| _|| _|	| _|d u rJt	
dt|| _d S )Nr   zDAll MySQL CHANGE/MODIFY COLUMN operations require the existing type.)r:   r   rp   r*   r+   r8   r9   r2   r4   r   CommandErrorrM   Zto_instancer.   )
rA   r-   r*   r   r8   r.   r+   r9   r2   r4   rB   rD   rE   rp   7  s    zMySQLChangeColumn.__init__)NNNNFNFrq   rD   rD   rB   rE   r>   6  s          r>   c                   @  s   e Zd ZdS )r?   N)ri   rj   rk   rD   rD   rD   rE   r?   S  s   r?   r"   ro   c                 K  s   t dd S )Nz9Individual alter column constructs not supported by MySQL)NotImplementedErrorelementcompilerr6   rD   rD   rE    _mysql_doesnt_support_individualW  s    rx   r   r#   )rv   rw   r7   c                 K  s<   dt || j| jt|| j| jd ur4dt|| j ndf S )Nz%s ALTER COLUMN %s %szSET DEFAULT %szDROP DEFAULT)r   r)   r   r   r*   r9   r   ru   rD   rD   rE   _mysql_alter_defaulta  s    
ry   c                 K  s<   dt || j| jt|| jt|| j| j| j| j	| j
df S )Nz%s MODIFY %s %sr+   r,   r.   r2   r4   )r   r)   r   r   r*   _mysql_colspecr+   r9   r.   r2   r4   ru   rD   rD   rE   _mysql_modify_columnn  s    
r|   c                 K  sF   dt || j| jt|| jt|| jt|| j| j| j	| j
| jdf S )Nz%s CHANGE %s %s %srz   )r   r)   r   r   r*   r8   r{   r+   r9   r.   r2   r4   ru   rD   rD   rE   _mysql_change_column  s    

r}   r$   rr   r   r'   )rw   r+   r,   r.   r2   r4   r7   c                 C  sn   d| j j||rdndf }|r*|d7 }|durL|d urL|dt| | 7 }|rj|d| j|t  7 }|S )Nz%s %sZNULLzNOT NULLz AUTO_INCREMENTFz DEFAULT %sz COMMENT %s)dialectZtype_compilerprocessr   Zsql_compilerZrender_literal_valuerM   String)rw   r+   r,   r.   r2   r4   specrD   rD   rE   r{     s    

r{   r   c                 K  s   | j }t|tjtjtjfr.|r$J || S t|tjrt|j	rbd|j
|j|j
|f S d|j
|j|j
|f S ntddS )zVRedefine SQLAlchemy's drop constraint to
    raise errors for invalid constraint type.z!ALTER TABLE %s DROP CONSTRAINT %szALTER TABLE %s DROP CHECK %szFNo generic 'DROP CONSTRAINT' in MySQL - please specify constraint typeN)rv   rH   r   ZForeignKeyConstraintZPrimaryKeyConstraintZUniqueConstraintZvisit_drop_constraintrI   r   r~   preparerZformat_tabletableZformat_constraintrt   )rv   rw   r6   
constraintrD   rD   rE   _mysql_drop_constraint  s.    



r   )6
__future__r   rQ   typingr   r   r   r   Z
sqlalchemyr   r	   rM   Zsqlalchemy.ext.compilerr
   baser   r   r   r   r   r   r   r   implr   rO   r   Zautogenerater   r   Zutil.sqla_compatr   r   r   Zsqlalchemy.dialects.mysql.baser   Zsqlalchemy.sql.ddlr   Zsqlalchemy.sql.schemar   Zsqlalchemy.sql.type_apir   r    r!   rn   r@   r>   r?   rx   ry   r|   r}   r{   r   rD   rD   rD   rE   <module>   sd      






