R6�W�] � @ sG d Z d d l m Z d d l Z d d l Z d d l Td 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 Gd d � d e � Z
d
d � Z Gd d � d e � Z Gd d � d e � Z d" Z Gd d � d e d e � � Z Gd d � d e � Z Gd d � d e � Z d d! � Z d S)#a�
Lightweight schema migrations.
NOTE: Currently tested with SQLite and Postgresql. MySQL may be missing some
features.
Example Usage
-------------
Instantiate a migrator:
# Postgres example:
my_db = PostgresqlDatabase(...)
migrator = PostgresqlMigrator(my_db)
# SQLite example:
my_db = SqliteDatabase('my_database.db')
migrator = SqliteMigrator(my_db)
Then you will use the `migrate` function to run various `Operation`s which
are generated by the migrator:
migrate(
migrator.add_column('some_table', 'column_name', CharField(default=''))
)
Migrations are not run inside a transaction, so if you wish the migration to
run in a transaction you will need to wrap the call to `migrate` in a
transaction block, e.g.:
with my_db.transaction():
migrate(...)
Supported Operations
--------------------
Add new field(s) to an existing model:
# Create your field instances. For non-null fields you must specify a
# default value.
pubdate_field = DateTimeField(null=True)
comment_field = TextField(default='')
# Run the migration, specifying the database table, field name and field.
migrate(
migrator.add_column('comment_tbl', 'pub_date', pubdate_field),
migrator.add_column('comment_tbl', 'comment', comment_field),
)
Renaming a field:
# Specify the table, original name of the column, and its new name.
migrate(
migrator.rename_column('story', 'pub_date', 'publish_date'),
migrator.rename_column('story', 'mod_date', 'modified_date'),
)
Dropping a field:
migrate(
migrator.drop_column('story', 'some_old_field'),
)
Making a field nullable or not nullable:
# Note that when making a field not null that field must not have any
# NULL values present.
migrate(
# Make `pub_date` allow NULL values.
migrator.drop_not_null('story', 'pub_date'),
# Prevent `modified_date` from containing NULL values.
migrator.add_not_null('story', 'modified_date'),
)
Renaming a table:
migrate(
migrator.rename_table('story', 'stories_tbl'),
)
Adding an index:
# Specify the table, column names, and whether the index should be
# UNIQUE or not.
migrate(
# Create an index on the `pub_date` column.
migrator.add_index('story', ('pub_date',), False),
# Create a multi-column index on the `pub_date` and `status` fields.
migrator.add_index('story', ('pub_date', 'status'), False),
# Create a unique index on the category and title fields.
migrator.add_index('story', ('category_id', 'title'), True),
)
Dropping an index:
# Specify the index name.
migrate(migrator.drop_index('story', 'story_pub_date_status'))
� )�
namedtupleN)�*)� CommaClause)�EnclosedClause)�Entity)�
Expression)�Node)�OPc @ sR e Z d Z d Z d d � Z d d � Z d d � Z d d � Z d
d � Z d S)
� Operationz/Encapsulate a single schema altering operation.c O s( | | _ | | _ | | _ | | _ d S)N)�migrator�method�args�kwargs)�selfr r r
r � r � /migrate.py�__init__u s zOperation.__init__c C s | j j j � } | j | � S)N)r �database�compilerZ
parse_node)r �noder r r r � _parse_node{ s zOperation._parse_nodec C s/ | j | � \ } } | j j j | | � d S)N)r r r � execute_sql)r r �sqlZparamsr r r �execute s zOperation.executec C sr t | t � r | j | � nO t | t � r; | j � n3 t | t t f � rn x | D] } | j | � qW Wd S)N)�
isinstancer r r
�run�list�tuple�_handle_result)r �result�itemr r r r � s
zOperation._handle_resultc C sE | j j � } d | d <| j t | j | j � | j | � � d S)NT�generate)r �copyr �getattrr r r
)r r r r r r � s
z
Operation.runN) �__name__�
__module__� __qualname__�__doc__r r r r r r r r r r
s s r
c s% t j � � � f d d � � } | S)Nc s> | j d d � } | r( � | | | � St | � j | | � S)Nr! F)�popr
r$ )r r
r r! )�fnr r �inner� s zoperation.<locals>.inner)� functools�wraps)r) r* r )r) r � operation� s !r- c @ s0 e Z d Z d Z d Z d d � Z e d d � � Z e d d � � Z e d d � � Z
d
d � Z e d d
� � Z e d d � � Z
e d d � � Z e d d d � � Z e d d � � Z d d � Z e d d � � Z e d d � � Z e d d � � Z e d d d � � Z e d! d" � � Z d# S)$�SchemaMigratorFc C s
| | _ d S)N)r )r r r r r r � s zSchemaMigrator.__init__c C s@ t | t � r t | � St | t � r2 t | � St | � Sd S)N)r ZPostgresqlDatabase�PostgresqlMigratorZ
MySQLDatabase�
MySQLMigrator�SqliteMigrator)�clsr r r r �
from_database� s
z SchemaMigrator.from_databasec
C sm | j } t | � r | � } t t d � t | � t d � t t | � t j t | j | � � d d �� S)NZUPDATEZSETZflatT)
�default�callable�Clause�SQLr r r ZEQZParamZdb_value)r �table� column_name�fieldr4 r r r �
apply_default� s z SchemaMigrator.apply_defaultc C s� | j d } | _ | | _ | _ | j j � j | � } | | _ t d � t | � t d � | g } t | t � r� | j
| j | � � t | � S)NTz ALTER TABLEz
ADD COLUMN)
�null�name� db_columnr r Zfield_definitionr7 r r �ForeignKeyField�extend�get_inline_fk_sqlr6 )r r8 r9 r: Z
field_nullZ field_clause�partsr r r �alter_add_column� s z SchemaMigrator.alter_add_columnc C s4 t d � t | j j j � t t | j j � � g S)N�
REFERENCES)r7 r � rel_model�_meta�db_tabler �to_fieldr> )r r: r r r rA � s z SchemaMigrator.get_inline_fk_sqlc C s
t � d S)N)�NotImplementedError)r r8 r9 r: r r r �add_foreign_key_constraint� s z)SchemaMigrator.add_foreign_key_constraintc C s� | j r) | j d k r) t d | � � t | t � } | rT | j rT t d � � | j | | | � g } | j s� | j | j | | | � | j | | � g � | r� | j
r� | j | j | | | j
j j | j j � � | S)Nz!%s is not null but has no defaultz'Foreign keys must specify a `to_field`.)r< r4 �
ValueErrorr r? rH rC r@ r; � add_not_null�explicit_create_foreign_key�appendrJ rE rF rG r> )r r8 r9 r: �is_foreign_key�
operationsr r r �
add_column� s$ zSchemaMigrator.add_columnc C s
t � d S)N)rI )r r8 r9 r r r �drop_foreign_key_constraint� s z*SchemaMigrator.drop_foreign_key_constraintTc C s� t d � t | � t d � t | � g } | rC | j t d � � t | � } d d � | j j | � D� } | | k r� | j r� | j | | � | g S| Sd S)Nz ALTER TABLEz DROP COLUMNZCASCADEc S s g | ] } | j � q Sr )�column)�.0Z foreign_keyr r r �
<listcomp> s z.SchemaMigrator.drop_column.<locals>.<listcomp>)r7 r rN r6 r �get_foreign_keys�explicit_delete_foreign_keyrR )r r8 r9 �cascade�nodesZdrop_column_nodeZ
fk_columnsr r r � drop_column� s zSchemaMigrator.drop_columnc C s= t t d � t | � t d � t | � t d � t | � � S)Nz ALTER TABLEz
RENAME COLUMN�TO)r6 r7 r )r r8 �old_name�new_namer r r �
rename_column s z SchemaMigrator.rename_columnc C s( t d � t | � t d � t | � g S)Nz ALTER TABLEz ALTER COLUMN)r7 r )r r8 rS r r r �
_alter_column s z SchemaMigrator._alter_columnc C s/ | j | | � } | j t d � � t | � S)Nz SET NOT NULL)r_ rN r7 r6 )r r8 rS rY r r r rL " s zSchemaMigrator.add_not_nullc C s/ | j | | � } | j t d � � t | � S)Nz
DROP NOT NULL)r_ rN r7 r6 )r r8 rS rY r r r �
drop_not_null( s z SchemaMigrator.drop_not_nullc C s+ t t d � t | � t d � t | � � S)Nz ALTER TABLEz RENAME TO)r6 r7 r )r r\ r] r r r � rename_table. s
zSchemaMigrator.rename_tablec C sn | j j � } | r d n d } t t | � t | j | | � � t d � t | � t d d � | D� � � S)NzCREATE UNIQUE INDEXz CREATE INDEX�ONc S s g | ] } t | � � q Sr )r )rT rS r r r rU ? s z,SchemaMigrator.add_index.<locals>.<listcomp>)r r r6 r7 r �
index_namer )r r8 �columns�uniquer Z statementr r r � add_index6 s zSchemaMigrator.add_indexc C s t t d � t | � � S)Nz
DROP INDEX)r6 r7 r )r r8 rc r r r �
drop_indexA s zSchemaMigrator.drop_indexN)r$ r% r& rM rW r � classmethodr3 r- r; rC rA rJ rQ rR rZ r^ r_ rL r` ra rf rg r r r r r. � s( "
r. c s4 e Z d Z d d � Z e � f d d � � Z � S)r/ c C s3 d } | j j | | � } d d � | j � D� S)Nai
SELECT pg_attribute.attname
FROM pg_index, pg_class, pg_attribute
WHERE
pg_class.oid = '%s'::regclass AND
indrelid = pg_class.oid AND
pg_attribute.attrelid = pg_class.oid AND
pg_attribute.attnum = any(pg_index.indkey) AND
indisprimary;
c S s g | ] } | d � q S)r r )rT �rowr r r rU U s z;PostgresqlMigrator._primary_key_columns.<locals>.<listcomp>)r r �fetchall)r Ztbl�query�cursorr r r �_primary_key_columnsI s
z'PostgresqlMigrator._primary_key_columnsc
s� | j | � } t t | � } | j | | d d �g } t | � d k r� d | | d f } d } | j j | | f � } t | j � � r� d | | d f } | j | j | | d d �� | S)Nr! T� z %s_%s_seqr z�
SELECT 1
FROM information_schema.sequences
WHERE LOWER(sequence_name) = LOWER(%s)
)
rm �superr/ ra �lenr r �bool�fetchonerN )
r r\ r] Zpk_namesZ ParentClassrP Zseq_namerk rl Z new_seq_name)� __class__r r ra W s z PostgresqlMigrator.rename_table)r$ r% r& rm r- ra r r )rs r r/ H s r/ r= �
definitionr< �pkr4 �extrac @ sX e Z d Z e d d � � Z e d d � � Z e d d � � Z d d d d � Z d S)
� MySQLColumnc C s
| j d k S)NZPRI)ru )r r r r �is_pkr s zMySQLColumn.is_pkc C s
| j d k S)NZUNI)ru )r r r r � is_uniquev s zMySQLColumn.is_uniquec C s
| j d k S)NZYES)r< )r r r r �is_nullz s zMySQLColumn.is_nullNc C s� | d k r | j } | d k r* | j } t | � t | j � g } | j ra | j t d � � | r} | j t d � � n | j t d � � | j r� | j t d � � | j r� | j t | j � � t | � S)NZUNIQUEZNULLzNOT NULLz PRIMARY KEY)
rz r= r r7 rt ry rN rx rv r6 )r r9 rz rB r r r r ~ s zMySQLColumn.sql)r$ r% r& �propertyrx ry rz r r r r r rw q s rw Z_Columnc @ s� e Z d Z d Z d Z e d d � � Z d d � Z e d d � � Z d d � Z e d
d � � Z
d d
� Z e d d � � Z e d d � � Z
e d d � � Z e d d � � Z d S)r0 Tc C s+ t t d � t | � t d � t | � � S)Nz RENAME TABLEr[ )r6 r7 r )r r\ r] r r r ra � s
zMySQLMigrator.rename_tablec C sV | j j d | � } | j � } x- | D]% } t | � } | j | k r) | Sq) Wd S)Nz DESCRIBE %s;F)r r rj rw r= )r r8 r9 rl Zrowsri rS r r r �_get_column_definition� s
z$MySQLMigrator._get_column_definitionc C sw d | | | f } t t d � t | � t d � t | � t d � t t | � � t d � t | � t t | � � � S)Nzfk_%s_%s_refs_%sz ALTER TABLEzADD CONSTRAINTz FOREIGN KEYrD )r6 r7 r r )r r8 r9 ZrelZ
rel_columnZ
constraintr r r rJ � s z(MySQLMigrator.add_foreign_key_constraintc C sK | j j d | | f � } | j � } | sC t d | | f � � | d S)Nz�SELECT constraint_name FROM information_schema.key_column_usage WHERE table_schema = DATABASE() AND table_name = %s AND column_name = %s;z=Unable to find foreign key constraint for "%s" on table "%s".r )r r rr �AttributeError)r r8 r9 rl r r r r �get_foreign_key_constraint� s z(MySQLMigrator.get_foreign_key_constraintc C s7 t t d � t | � t d � t | j | | � � � S)Nz ALTER TABLEzDROP FOREIGN KEY)r6 r7 r r~ )r r8 r9 r r r rR � s
z)MySQLMigrator.drop_foreign_key_constraintc C s g S)Nr )r r: r r r rA � s z MySQLMigrator.get_inline_fk_sqlc C sC | j | | � } t t d � t | � t d � | j d d � � S)Nz ALTER TABLE�MODIFYrz F)r| r6 r7 r r )r r8 rS r r r rL � s zMySQLMigrator.add_not_nullc C sX | j | | � } | j r' t d � � t t d � t | � t d � | j d d � � S)Nz Primary keys can not be nullz ALTER TABLEr rz T)r| rx rK r6 r7 r r )r r8 rS r r r r` � s zMySQLMigrator.drop_not_nullc C s� t d d � | j j | � D� � } | | k } | j | | � } t t d � t | � t d � t | � | j d | � � } | r� | | } | j | | � | | j | | | j
| j � g S| Sd S)Nc s s | ] } | j | f Vq d S)N)rS )rT Zfkr r r � <genexpr>� s z.MySQLMigrator.rename_column.<locals>.<genexpr>z ALTER TABLEZCHANGEr9 ) �dictr rV r| r6 r7 r r rR rJ Z
dest_tableZ dest_column) r r8 r\ r] Z
fk_objectsrO rS Z
rename_clauseZ fk_metadatar r r r^ � s*
zMySQLMigrator.rename_columnc C s+ t t d � t | � t d � t | � � S)Nz
DROP INDEXrb )r6 r7 r )r r8 rc r r r rg s
zMySQLMigrator.drop_indexN)r$ r% r& rM rW r- ra r| rJ r~ rR rA rL r` r^ rg r r r r r0 � s r0 c @ s� e Z d Z d Z e j d � Z e j d � Z e j d � Z e j d e j � Z
d d � Z d d � Z e
d
d � � Z d d
� Z e
d d d � � Z e
d d � � Z e
d d � � Z e
d d � � Z d S)r1 z�
SQLite supports a subset of ALTER TABLE queries, view the docs for the
full details http://sqlite.org/lang_altertable.html
z
(.+?)\((.+)\)z(?:[^,(]|\([^)]*\))+z
["`']?([\w]+)z FOREIGN KEY\s+\("?([\w]+)"?\)\s+c C s* | j j d | � } d d � | j D� S)Nzselect * from "%s" limit 1c S s g | ] } | d � q S)r r )rT r r r r rU s z4SqliteMigrator._get_column_names.<locals>.<listcomp>)r r � description)r r8 �resr r r �_get_column_names s z SqliteMigrator._get_column_namesc C s+ | j j d d | j � g � } | j � S)NzBselect name, sql from sqlite_master where type=? and LOWER(name)=?r8 )r r �lowerrr )r r8 r� r r r �_get_create_table s z SqliteMigrator._get_create_tablec s> t d d � � j j | � D� � } | j � | k rM t d | | f � � � j | � \ } } � j j | � } � j j | � t j d d | � } � j
j | � j � \ } } � j
j | � } d d � | D� }
g } g } g }
x� |
D]� } � j j | � j � \ } | | k r|| | | � } | r�| j | � |
j | � � j j | � j � \ } | j | � q� | j | � | j � j d � s� | j | � |
j | � q� Wt t |
| � � } | j | � � d
d � } � sd d � } n! � | k r"� � f d
d � } g } xa | D]Y } � j j | � } | d k ru| j � d | k ru| | � } | r/| j | � q/W| d } t j d | t j � } | j d | | � } d j | � } t t d � t | � � t d | j � | f � g } t t d � t | � t d d � | D� � t d � t d d � |
D� � t d � t | � � } | j | � | j t t d � t | � � � | j � j | | � � x� t! d d � | � D]k } | | j" k r�| j t | j# � � q�� r�� j$ | j# | � � } | d k r�| j t | � � q�W| S) Nc s s | ] } | j j � Vq d S)N)r= r� )rT rS r r r r� s z0SqliteMigrator._update_column.<locals>.<genexpr>z"Column "%s" does not exist on "%s"z\s+� c S s g | ] } | j � � q Sr )�strip)rT �colr r r rU 8 s z1SqliteMigrator._update_column.<locals>.<listcomp>�foreign�primaryc S s | S)Nr )�
column_defr r r �<lambda>S s z/SqliteMigrator._update_column.<locals>.<lambda>c S s d S)Nr )r� r r r r� V s c s � j j d � | � S)NzFOREIGN KEY ("%s") )�fk_re�sub)r� )�
new_columnr r r r� Y s r Z__tmp__z
("?)%s("?)z\1%s\2z, zDROP TABLE IF EXISTSz%s (%s)z INSERT INTOc S s g | ] } t | � � q Sr )r )rT r� r r r rU v s ZSELECTc S s g | ] } t | � � q Sr )r )rT r� r r r rU x s ZFROMz
DROP TABLEc S s | j S)N)r )�idxr r r r� � s )r� r� )%�setr Z get_columnsr� rK r� Z get_indexesrV �rer� � column_re�search�groups�column_split_re�findall�column_name_re�matchrN �
startswithr� �zip�getr� �compile�I�joinr6 r7 r r� r r ra �filterrd r �
_fix_index) r r8 �column_to_updater) rd Z create_tableZindexesZ
raw_createZ raw_columnsZ
split_columnsZ column_defsZnew_column_defsZnew_column_namesZoriginal_column_namesr� r9 Znew_column_defZoriginal_to_newZ fk_filter_fnZcleaned_columnsr� Z
temp_tableZrgxZcreateZqueriesZpopulate_table�indexr r )r� r r �_update_column s�
"
z SqliteMigrator._update_columnc
C s( | j | � } t | � d k r1 | j | | � S| j d d � \ } } t | j | � � d k r~ d | | j | | � f S| j d d � d j d � } d d � | D� } g } xK | D]C } t j d
| | � r� t | t | � d � } | j | � q� Wd | d j d
d � | D� � f S)N� �(rn z%s(%s�)r �,c S s g | ] } | j d � � q S)z"`[]' )r� )rT �partr r r rU � s z-SqliteMigrator._fix_index.<locals>.<listcomp>z%s(?:['"`\]]?\s|$)z%s(%s)z, c s s | ] } d | Vq d S)z"%s"Nr )rT �cr r r r� � s z,SqliteMigrator._fix_index.<locals>.<genexpr>) �splitrp �replace�rsplitr� r� Z new_columnerN r� )
r r r� r� rB ZlhsZrhsrd ZcleanrS r r r r� � s
zSqliteMigrator._fix_indexTc C s | j | | d d � � S)Nc S s d S)Nr )�a�br r r r� � s z,SqliteMigrator.drop_column.<locals>.<lambda>)r� )r r8 r9 rX r r r rZ � s zSqliteMigrator.drop_columnc s% � f d d � } | j | | | � S)Nc s | j | � � S)N)r� )r9 r� )r] r r �_rename� s z-SqliteMigrator.rename_column.<locals>._rename)r� )r r8 r\ r] r� r )r] r r^ � s z SqliteMigrator.rename_columnc C s d d � } | j | | | � S)Nc S s | d S)Nz NOT NULLr )r9 r� r r r �
_add_not_null� s z2SqliteMigrator.add_not_null.<locals>._add_not_null)r� )r r8 rS r� r r r rL � s zSqliteMigrator.add_not_nullc C s d d � } | j | | | � S)Nc S s | j d d � S)NzNOT NULL� )r� )r9 r� r r r �_drop_not_null� s z4SqliteMigrator.drop_not_null.<locals>._drop_not_null)r� )r r8 rS r� r r r r` � s z SqliteMigrator.drop_not_nullN)r$ r% r& r' r� r� r� r� r� r� r� r� r� r- r� r� rZ r^ rL r` r r r r r1 s q r1 c O s x | D] } | j � q Wd S)N)r )rP r r- r r r �migrate� s
r� )r= rt r< ru r4 rv )r' � collectionsr r+ r� Zpeeweer r r r r r �objectr
r- r. r/ Z_column_attributesrw r0 r1 r� r r r r �<module>e s&
�' "v�