U
/�Xp* � @ s� d dl Zd dlmZmZmZm Z
m Z m Z m
Z
mZmZmZmZmZ d dlmZ G dd� de�ZG dd� dee�ZG dd � d ee�ZG d
d � d ee
�Z d d
� ZG dd� de�ZdS )� N) �
MySQLMigrator�PostgresqlMigrator�SchemaMigrator�SqliteMigrator� Operation�SQL�Entity�Clause�PostgresqlDatabase� operation�SqliteDatabase�
MySQLDatabase)�LOGGERc sZ e Zd ZdZe� fdd��Zddd�Zedd� �Zd d
� Z ed d � �Z
ed
d� �Z � Z S )r zImplement migrations.c sF t |t�rt|�S t |t�r$t|�S t |t�r6t|�S tt| �� |�S )zInitialize migrator by db.)
�
isinstancer
r r r r
r �superr �
from_database)�cls�database�� __class__� � /migrator.pyr s
z SchemaMigrator.from_databaseTc s � �fdd�S )Nc s �j � d�S )N��cascade��
drop_tabler �r �modelr r �<lambda> � z+SchemaMigrator.drop_table.<locals>.<lambda>r ��selfr r r r r r s zSchemaMigrator.drop_tablec C s. | � |||�g}|js*|�| �||�g� |S )zChange column.)�alter_change_column�null�extend� add_not_null)r! �table� column_name�fieldZ
operationsr r r �
change_column s z SchemaMigrator.change_columnc C s> |j d }|_ | j�� �|�}||_ ttd�t|�td�|�S )�Support change columns.T� ALTER TABLEz ALTER COLUMN)r# r �compiler�field_definitionr r r )r! r& �columnr( �
field_null� field_clauser r r r" ' s z"SchemaMigrator.alter_change_columnc G s t t|f|�� �S )zExecute raw SQL.)r r �r! �sqlZparamsr r r r2 . s zSchemaMigrator.sqlc C sn |j d }|_ ||_| j�� �|�}||_ td�t|�td�|g}t|tj �r`|�
| � |�� n||_ t
|� S )zKeep fieldname unchanged.Tr+ z
ADD COLUMN)r# � db_columnr r, r- r r r �pw�ForeignKeyFieldr$ Zget_inline_fk_sql�namer )r! r& r' r( r/ r0 �partsr r r �alter_add_column3 s � z SchemaMigrator.alter_add_column)T)
�__name__�
__module__� __qualname__�__doc__� classmethodr r r r) r" r2 r8 �
__classcell__r r r r r
s
r c s e Zd ZdZ� fdd�Z� ZS )r z Support the migrations in MySQL.c s4 t t| ��|||�}|jd }|j�dtd�� |S �r* ���� ZTYPE)r r r" �nodes�insertr �r! r& r' r( �clauser0 r r r r" K s
z!MySQLMigrator.alter_change_column�r9 r: r; r< r" r> r r r r r G s r c s e Zd ZdZ� fdd�Z� ZS )r z%Support the migrations in postgresql.c s4 t t| ��|||�}|jd }|j�dtd�� |S r? )r r r" rB rC r rD r r r r" W s
z&PostgresqlMigrator.alter_change_columnrF r r r r r S s r c @ s" e Zd ZdZddd�Zdd� ZdS ) r z!Support the migrations in sqlite.Tc s � fdd�S )z0SQLite doesnt support cascade syntax by default.c s � j dd�S )NFr r r �r r r r e r z+SqliteMigrator.drop_table.<locals>.<lambda>r r r rG r r c s zSqliteMigrator.drop_tablec s � �fdd�}�� |||�S )r* c s& �j �� }|�� �}|�|�\}}|S �N)r r, r- Z
parse_node)r' Z
column_defr, rE r2 �_�r( r! r r �_changei s
z3SqliteMigrator.alter_change_column.<locals>._change)Z_update_column)r! r& r. r( rK r rJ r r" g s z"SqliteMigrator.alter_change_columnN)T)r9 r: r; r< r r" r r r r r _ s
r c s � fdd�}|S )z Convert string to model class.c s4 t |t�r"� | | j| f|�|�S � | |f|�|�S rH )r �str�orm)�migratorr �args�kwargs��methodr r �wrappers s
zget_model.<locals>.wrapperr )rR rS r rQ r � get_modelq s rT c @ s� e Zd ZdZdd� Zdd� Zdd� Zdd � Zd
d � Zd d
� Z e Z
e d(dd��Z e Z
e dd� �ZeZe dd� �ZeZe dd� �ZeZdd� Ze dd� �ZeZe dd � �Ze d d � �Ze d d � �Ze d!d"� �Ze d#d$� �Ze d%d&� �Z d'S ))�MigratorzProvide migrations.c C s: t |tj�r|j}|| _t� | _t� | _t �
| j�| _ dS )zInitialize the migrator.N) r r4 ZProxy�objr �dictrM �list�opsr r rN )r! r r r r �__init__~ s zMigrator.__init__c C sD | j D ]0}t|t�r0t�d|j|j� |�� q|� q| �� dS )zRun operations.z%s %sN) rY r r r �inforR rO �run�clean)r! Zopnr r r r\ � s
z Migrator.runc s | j �� ��fdd�� dS )zRun python code.c s
�� ��S rH r r �rO �funcrP r r r � r z!Migrator.python.<locals>.<lambda>N)rY �append)r! r_ rO rP r r^ r �python� s zMigrator.pythonc G s | j �| jj|f|�� � dS )zExecure raw SQL.N)rY r` rN r2 r1 r r r r2 � s z Migrator.sqlc C s t � | _dS )zClean the operations.N)rX rY )r! r r r r] � s zMigrator.cleanc C s* || j |jj< | j|j_| j�|j� |S )zUCreate model and table in database.
>> migrator.create_table(model)
)rM �_meta�db_tabler rY r` � create_table)r! r r r r rd � s
zMigrator.create_tableTc C s& | j |jj= | j�| j�||�� dS )zaDrop model and table from database.
>> migrator.drop_table(model, cascade=True)
N)rM rb rc rY r` rN r r r r r r � s zMigrator.drop_tablec K sh |� � D ]Z\}}|�||� | j�| j�|jj|j|�� |j r| j�| jj
|jj|jfdd�� q|S )zCreate new fields.T��unique) �items� add_to_classrY r` rN Z
add_columnrb rc r3 rf � add_index�r! r �fieldsr6 r( r r r � add_columns� s �
zMigrator.add_columnsc K sh |� � D ]Z\}}|�||� | j�| j�|jj|j|�� |j r| j�| jj
|jj|jfdd�� q|S )zChange fields.Tre ) rg rh rY r` rN r) rb rc r3 rf ri rj r r r �change_columns� s � �
zMigrator.change_columnsc s� � fdd�|j j�� D �}|�dd�}|D ]n}| �||� |jrx| j�� }|�|j j |j
f�}| j � | j
�|j j |�� | j � | j
j|j j |j
|d�� q*|S )zRemove fields from model.c s g | ]}|j � kr|�qS r )r6 )�.0r( ��namesr r �
<listcomp>� s
z)Migrator.drop_columns.<locals>.<listcomp>r Tr )rb rk �values�pop�
__del_field__rf r r, �
index_namerc r3 rY r` rN �
drop_indexZ drop_column) r! r rp rP rk r r( r, ru r ro r � drop_columns� s
�zMigrator.drop_columnsc C sF |j �|j� t||j� t|tj�rBt|j|j� |jj j |j= dS )zDelete field from model.N)
rb Z remove_fieldr6 �delattrr r4 r5 Z rel_modelZ related_nameZ reverse_rel)r! r r( r r r rt � s
zMigrator.__del_field__c C s~ |j j| }t|tj�r |j}| �||� | |_|_|�||� t|tj�r^|jd |_}| j �
| j � |j j
||�� |S )zRename field in model.�_id)rb rk r r4 r5 r3 rt r6 rh rY r` rN �
rename_columnrc )r! r Zold_name�new_namer( r r r rz � s zMigrator.rename_columnc C s@ | j |jj= ||j_|| j |jj< | j�| j�|jj|�� |S )zRename table in database.)rM rb rc rY r` rN � rename_table)r! r r{ r r r r| � s
zMigrator.rename_tablec O sz |� dd�}|jj�||f� g }|D ]0}|jj�|�}t|tj�rL|d }|�|� q&| j �| j
j |jj ||d�� |S )zCreate indexes.rf Fry re )
rs rb �indexesr` rk �getr r4 r5 rY rN ri rc )r! r �columnsrP rf �columns_�colr( r r r ri s zMigrator.add_indexc s� g }� D ]0}|j j�|�}t|tj�r.|d }|�|� q| jj� � �
|j j |�}� fdd�|j j D �|j _ | j
�| j�|j j |�� |S )z
Drop indexes.ry c s g | ]\}}� |kr||f�qS r r )rn ZcolsrI �r r r rq s z'Migrator.drop_index.<locals>.<listcomp>)rb rk r~ r r4 r5 r` rN r r, ru rc r} rY rv )r! r r r� r� r( ru r r� r rv s zMigrator.drop_indexc G s< |D ]2}|j j| }d|_| j�| j�|j j|j�� q|S )z
Add not null.F) rb rk r# rY r` rN r% rc r3 �r! r rp r6 r( r r r r% s
zMigrator.add_not_nullc G s< |D ]2}|j j| }d|_| j�| j�|j j|j�� q|S )zDrop not null.T) rb rk r# rY r` rN �
drop_not_nullrc r3 r� r r r r� ' s
zMigrator.drop_not_nullc C s> |j j| }| |j j|< |_| j�| j�|j j||�� |S )z Add default.) rb rk �defaults�defaultrY r` rN Z
apply_defaultrc )r! r r6 r� r( r r r � add_default0 s zMigrator.add_defaultN)T) r9 r: r; r< rZ r\ ra r2 r] rd Z create_modelrT r Z remove_modelrl Z
add_fieldsrm Z
change_fieldsrw Z
remove_fieldsrt rz Z rename_fieldr| ri rv r% r� r� r r r r rU z sH
rU )Zpeeweer4 Zplayhouse.migrater ZMqMr ZPgMr ZScMr ZSqMr r r r r
r r r
Zpeewee_migrater rT �objectrU r r r r �<module> s 8 :