<s> import contextlib <EOL> import functools <EOL> import alembic <EOL> from alembic import context <EOL> from alembic import op <EOL> import sqlalchemy as sa <EOL> from sqlalchemy . engine import reflection <EOL> from neutron . _i18n import _ <EOL> LIBERTY = '<STR_LIT>' <EOL> MITAKA = '<STR_LIT>' <EOL> NEWTON = '<STR_LIT>' <EOL> NEUTRON_MILESTONES = [ <EOL> LIBERTY , <EOL> MITAKA , <EOL> NEWTON , <EOL> ] <EOL> CREATION_OPERATIONS = ( sa . sql . ddl . CreateIndex , <EOL> sa . sql . ddl . CreateTable , <EOL> sa . sql . ddl . CreateColumn , <EOL> ) <EOL> DROP_OPERATIONS = ( sa . sql . ddl . DropConstraint , <EOL> sa . sql . ddl . DropIndex , <EOL> sa . sql . ddl . DropTable , <EOL> alembic . ddl . base . DropColumn ) <EOL> def skip_if_offline ( func ) : <EOL> """<STR_LIT>""" <EOL> @ functools . wraps ( func ) <EOL> def decorator ( * args , ** kwargs ) : <EOL> if context . is_offline_mode ( ) : <EOL> return <EOL> return func ( * args , ** kwargs ) <EOL> return decorator <EOL> def raise_if_offline ( func ) : <EOL> """<STR_LIT>""" <EOL> @ functools . wraps ( func ) <EOL> def decorator ( * args , ** kwargs ) : <EOL> if context . is_offline_mode ( ) : <EOL> raise RuntimeError ( _ ( "<STR_LIT>" ) % <EOL> func . __name__ ) <EOL> return func ( * args , ** kwargs ) <EOL> return decorator <EOL> @ raise_if_offline <EOL> def schema_has_table ( table_name ) : <EOL> """<STR_LIT>""" <EOL> bind = op . get_bind ( ) <EOL> insp = sa . engine . reflection . Inspector . from_engine ( bind ) <EOL> return table_name in insp . get_table_names ( ) <EOL> @ raise_if_offline <EOL> def schema_has_column ( table_name , column_name ) : <EOL> """<STR_LIT>""" <EOL> bind = op . get_bind ( ) <EOL> insp = sa . engine . reflection . Inspector . from_engine ( bind ) <EOL> if not schema_has_table ( table_name ) : <EOL> return <EOL> return column_name in [ column [ '<STR_LIT:name>' ] for column in <EOL> insp . get_columns ( table_name ) ] <EOL> @ raise_if_offline <EOL> def alter_column_if_exists ( table_name , column_name , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if schema_has_column ( table_name , column_name ) : <EOL> op . alter_column ( table_name , column_name , ** kwargs ) <EOL> @ raise_if_offline <EOL> def drop_table_if_exists ( table_name ) : <EOL> if schema_has_table ( table_name ) : <EOL> op . drop_table ( table_name ) <EOL> @ raise_if_offline <EOL> def rename_table_if_exists ( old_table_name , new_table_name ) : <EOL> if schema_has_table ( old_table_name ) : <EOL> op . rename_table ( old_table_name , new_table_name ) <EOL> def alter_enum ( table , column , enum_type , nullable , do_drop = True , <EOL> do_rename = True , do_create = True ) : <EOL> """<STR_LIT>""" <EOL> bind = op . get_bind ( ) <EOL> engine = bind . engine <EOL> if engine . name == '<STR_LIT>' : <EOL> values = { '<STR_LIT>' : table , <EOL> '<STR_LIT>' : column , <EOL> '<STR_LIT:name>' : enum_type . name } <EOL> if do_rename : <EOL> op . execute ( "<STR_LIT>" % values ) <EOL> if do_create : <EOL> enum_type . create ( bind , checkfirst = False ) <EOL> op . execute ( "<STR_LIT>" <EOL> "<STR_LIT>" % values ) <EOL> op . add_column ( table , sa . Column ( column , enum_type , nullable = nullable ) ) <EOL> op . execute ( "<STR_LIT>" <EOL> "<STR_LIT>" % values ) <EOL> op . execute ( "<STR_LIT>" % values ) <EOL> if do_drop : <EOL> op . execute ( "<STR_LIT>" % values ) <EOL> else : <EOL> op . alter_column ( table , column , type_ = enum_type , <EOL> existing_nullable = nullable ) <EOL> def create_table_if_not_exist_psql ( table_name , values ) : <EOL> if op . get_bind ( ) . engine . dialect . server_version_info < ( <NUM_LIT:9> , <NUM_LIT:1> , <NUM_LIT:0> ) : <EOL> op . execute ( "<STR_LIT>" ) <EOL> op . execute ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> op . execute ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> op . execute ( "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> { '<STR_LIT:name>' : table_name , <EOL> '<STR_LIT>' : values } ) <EOL> def remove_foreign_keys ( table , foreign_keys ) : <EOL> for fk in foreign_keys : <EOL> op . drop_constraint ( <EOL> constraint_name = fk [ '<STR_LIT:name>' ] , <EOL> table_name = table , <EOL> type_ = '<STR_LIT>' <EOL> ) <EOL> def create_foreign_keys ( table , foreign_keys ) : <EOL> for fk in foreign_keys : <EOL> op . create_foreign_key ( <EOL> constraint_name = fk [ '<STR_LIT:name>' ] , <EOL> source_table = table , <EOL> referent_table = fk [ '<STR_LIT>' ] , <EOL> local_cols = fk [ '<STR_LIT>' ] , <EOL> remote_cols = fk [ '<STR_LIT>' ] , <EOL> ondelete = fk [ '<STR_LIT>' ] . get ( '<STR_LIT>' ) <EOL> ) <EOL> @ contextlib . contextmanager <EOL> def remove_fks_from_table ( table ) : <EOL> try : <EOL> inspector = reflection . Inspector . from_engine ( op . get_bind ( ) ) <EOL> foreign_keys = inspector . get_foreign_keys ( table ) <EOL> remove_foreign_keys ( table , foreign_keys ) <EOL> yield <EOL> finally : <EOL> create_foreign_keys ( table , foreign_keys ) </s>
