Python alembic.op.get_bind() Examples

def upgrade():
    agentstatus.create(op.get_bind())
    kernelstatus.create(op.get_bind())
    op.add_column('agents', sa.Column('lost_at', sa.DateTime(timezone=True), nullable=True))
    op.add_column('agents', sa.Column('status', sa.Enum('ALIVE', 'LOST', 'RESTARTING', 'TERMINATED', name='agentstatus'), nullable=False))
    op.create_index(op.f('ix_agents_status'), 'agents', ['status'], unique=False)
    op.add_column('kernels', sa.Column('agent_addr', sa.String(length=128), nullable=False))
    op.add_column('kernels', sa.Column('cpu_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('gpu_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('mem_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('repl_in_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('repl_out_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('stdin_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('stdout_port', sa.Integer(), nullable=False))
    op.drop_column('kernels', 'allocated_cores')
    op.add_column('kernels', sa.Column('cpu_set', sa.ARRAY(sa.Integer), nullable=True))
    op.add_column('kernels', sa.Column('gpu_set', sa.ARRAY(sa.Integer), nullable=True))
    op.alter_column('kernels', column_name='status', type_=sa.Enum(*kernelstatus_choices, name='kernelstatus'),
                    postgresql_using='status::kernelstatus') 

def upgrade():
    op.create_table(
        'project_locks',
        sa.Column('project_id',
                  sa.Integer,
                  sa.ForeignKey(
                    'projects.id', onupdate='CASCADE', ondelete='CASCADE'
                  ),
                  nullable=False,
                  primary_key=True
        ),
        sa.Column('lock_type',
                  sa.Enum(
                    'WORKER',
                    name='lock_type_enum'
                  ),
                  nullable=False,
                  primary_key=True
        )
    )

    # Add WORKER locks everywhere
    conn = op.get_bind()
    conn.execute("""INSERT INTO project_locks (project_id, lock_type)
                    SELECT id, 'WORKER' from projects""") 

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###

    userrole_choices = list(map(lambda v: v.value, UserRole))
    if 'superadmin' in userrole_choices:
        userrole_choices.remove('superadmin')
        conn = op.get_bind()

        # First, change all superadmin role to admin.
        query = textwrap.dedent("UPDATE users SET role = 'admin' WHERE role = 'superadmin';")
        conn.execute(query)

        # Remove superadmin from user role choices.
        conn.execute('ALTER TYPE userrole RENAME TO userrole___;')
        conn.execute('CREATE TYPE userrole as enum (%s)' % ("'" + "','".join(userrole_choices) + "'"))
        conn.execute(textwrap.dedent('''\
            ALTER TABLE users
                ALTER COLUMN role TYPE userrole USING role::text::userrole;
        '''))
        conn.execute('DROP TYPE userrole___;') 

def upgrade():
    conn = op.get_bind()

    insp = inspect(conn)
    fk_names = [fk['name'] for fk in insp.get_foreign_keys('host')]
    if ("fk_hypervisor_id_resource_id" not in fk_names and
            "fk_host_id_resource_id" in fk_names):
        # NOTE(sileht): we are already good, the BD have been created from
        # scratch after "a54c57ada3f5"
        return

    op.drop_constraint("fk_hypervisor_id_resource_id", "host",
                       type_="foreignkey")
    op.drop_constraint("fk_hypervisor_history_resource_history_revision",
                       "host_history", type_="foreignkey")
    op.create_foreign_key("fk_host_id_resource_id", "host", "resource",
                          ["id"], ["id"], ondelete="CASCADE")
    op.create_foreign_key("fk_host_history_resource_history_revision",
                          "host_history", "resource_history",
                          ["revision"], ["revision"], ondelete="CASCADE") 

def upgrade():
	op.add_column('campaigns', sqlalchemy.Column('credential_regex_username', sqlalchemy.String))
	op.add_column('campaigns', sqlalchemy.Column('credential_regex_password', sqlalchemy.String))
	op.add_column('campaigns', sqlalchemy.Column('credential_regex_mfa_token', sqlalchemy.String))

	op.add_column('credentials', sqlalchemy.Column('mfa_token', sqlalchemy.String))
	op.add_column('credentials', sqlalchemy.Column('regex_validated', sqlalchemy.Boolean))

	op.add_column('users', sqlalchemy.Column('access_level', sqlalchemy.Integer))
	op.execute('UPDATE users SET access_level = 1000')
	op.alter_column('users', 'access_level', nullable=False)

	# adjust the schema version metadata
	db_manager.Session.remove()
	db_manager.Session.configure(bind=op.get_bind())
	session = db_manager.Session()
	db_manager.set_metadata('schema_version', 9, session=session)
	session.commit() 

def upgrade():
    bind = op.get_bind()
    # NOTE(sileht): mysql can't delete an index on a foreign key
    # even this one is not the index used by the foreign key itself...
    # In our case we have two indexes fk_resource_history_id_resource_id and
    # and ix_resource_history_id, we want to delete only the second, but mysql
    # can't do that with a simple DROP INDEX ix_resource_history_id...
    # so we have to remove the constraint and put it back...
    if bind.engine.name == "mysql":
        op.drop_constraint("fk_resource_history_id_resource_id",
                           type_="foreignkey", table_name="resource_history")

    for table, colname in resource_tables + history_tables + other_tables:
        op.drop_index("ix_%s_%s" % (table, colname), table_name=table)

    if bind.engine.name == "mysql":
        op.create_foreign_key("fk_resource_history_id_resource_id",
                              "resource_history", "resource", ["id"], ["id"],
                              ondelete="CASCADE") 

def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.add_column('rss_parser_funcs', sa.Column('last_changed', sa.DateTime(), nullable=True))
    op.add_column('rss_parser_funcs_version', sa.Column('last_changed', sa.DateTime(), autoincrement=False, nullable=True))


    bind = op.get_bind()
    sess = Session(bind=bind)

    print("Updating date/time stamps for functions.")
    sess.query(RssFeedEntry).update({'last_changed' : datetime.datetime.now()})
    sess.commit()
    print("Update done.")

    op.alter_column('rss_parser_funcs', 'last_changed', nullable=False)

    ### end Alembic commands ### 

def upgrade():
    blockchain_status = postgresql.ENUM('PENDING', 'SUCCESS', 'FAILED', name='blockchainstatus')
    blockchain_status.create(op.get_bind())

    op.create_table('worker_messages',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('authorising_user_id', sa.Integer(), nullable=True),
    sa.Column('created', sa.DateTime(), nullable=True),
    sa.Column('updated', sa.DateTime(), nullable=True),
    sa.Column('message', sa.String(), nullable=True),
    sa.Column('error', sa.String(), nullable=True),
    sa.Column('worker_timestamp', sa.DateTime(), nullable=True),
    sa.Column('blockchain_task_uuid', sa.String(), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.add_column('credit_transfer', sa.Column('blockchain_hash', sa.String(), nullable=True))
    op.add_column('credit_transfer', sa.Column('blockchain_status', sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='blockchainstatus'), nullable=True))
    op.add_column('credit_transfer', sa.Column('last_worker_update', sa.DateTime(), nullable=True))
    op.add_column('exchange', sa.Column('blockchain_hash', sa.String(), nullable=True))
    op.add_column('exchange', sa.Column('blockchain_status', sa.Enum('PENDING', 'SUCCESS', 'FAILED', name='blockchainstatus'), nullable=True))
    op.add_column('exchange', sa.Column('last_worker_update', sa.DateTime(), nullable=True))
    op.create_index(op.f('ix_worker_messages_table_blockchain_task_uuid'), 'worker_messages', ['blockchain_task_uuid'], unique=False) 

def upgrade():
    op.add_column('compute_node',
                  sa.Column('rp_uuid', sa.String(length=36), nullable=True))
    op.create_unique_constraint('uniq_compute_node0rp_uuid',
                                'compute_node', ['rp_uuid'])

    # perform data migration between tables
    session = sa.orm.Session(bind=op.get_bind())
    with session.begin(subtransactions=True):
        for row in session.query(COMPUTE_NODE_TABLE):
            session.execute(
                COMPUTE_NODE_TABLE.update().values(
                    rp_uuid=row.uuid).where(
                        COMPUTE_NODE_TABLE.c.uuid == row.uuid)
            )
    # this commit is necessary to allow further operations
    session.commit()

    op.alter_column('compute_node', 'rp_uuid',
                    nullable=False,
                    existing_type=sa.String(length=36),
                    existing_nullable=True,
                    existing_server_default=False) 

def upgrade():
    bind = op.get_bind()
    sess = Session(bind=bind)

    op.add_column('nu_resolved_outbound', sa.Column('disabled', sa.Boolean(), nullable=True))

    print("Setting value for new column 1")
    sess.query(NuReleaseItem).filter(NuReleaseItem.release_date == None).update({"first_seen" : datetime.datetime.min})
    print("Setting value for new column 2")
    sess.query(NuResolvedOutbound).update({"disabled" : False})
    sess.commit()
    print("Setting nullability")
    op.alter_column('nu_resolved_outbound', 'disabled',
               existing_type=sa.Boolean(),
               nullable=False)

    ### end Alembic commands ### 

def downgrade():
    ''' Remove the column namespace from the table projects.
    '''
    # Update all the existing projects
    engine = op.get_bind()
    Session = sa.orm.scoped_session(sa.orm.sessionmaker())
    Session.configure(bind=engine)
    session = Session()
    for project in session.query(model.Project).filter(
            model.Project.namespace != None).all():
        if project.namespace.strip():
            project.name = '%s/%s' % (project.namespace, project.name)
            session.add(project)
    session.commit()

    op.drop_column('projects', 'namespace') 

def upgrade():
    ''' Add the column namespace to the table projects.
    '''
    op.add_column(
        'projects',
        sa.Column('namespace', sa.String(255), nullable=True, index=True)
    )

    # Update all the existing projects
    engine = op.get_bind()
    Session = sa.orm.scoped_session(sa.orm.sessionmaker())
    Session.configure(bind=engine)
    session = Session()
    for project in session.query(model.Project).filter(
            model.Project.name.ilike('%/%')).all():
        nspace, name = project.name.split('/', 1)
        project.name = name
        project.namespace = nspace
        session.add(project)
    session.commit() 

def downgrade():
    bind = op.get_bind()
    session = Session(bind=bind)
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('pull_mirror', sa.Column('gitlab_id', sa.INTEGER(), nullable=True))

    for pull_mirror in session.query(PullMirror):
        pull_mirror.gitlab_id = pull_mirror.project.gitlab_id
        session.add(pull_mirror)
    session.commit()

    op.drop_constraint(None, 'pull_mirror', type_='foreignkey')
    op.drop_index(op.f('ix_pull_mirror_project_id'), table_name='pull_mirror')
    op.drop_column('pull_mirror', 'project_id')
    op.drop_index(op.f('ix_push_mirror_user_id'), table_name='push_mirror')
    op.drop_index(op.f('ix_push_mirror_project_id'), table_name='push_mirror')
    op.drop_table('push_mirror')
    op.drop_table('project')
    # ### end Alembic commands ### 

def upgrade():
    userstatus.create(op.get_bind())
    op.add_column(
        'users',
        sa.Column('status', sa.Enum(*userstatus_choices, name='userstatus'), nullable=True)
    )
    op.add_column('users', sa.Column('status_info', sa.Unicode(), nullable=True))

    # Set user's status field.
    conn = op.get_bind()
    query = textwrap.dedent(
        "UPDATE users SET status = 'active', status_info = 'migrated' WHERE is_active = 't';"
    )
    conn.execute(query)
    query = textwrap.dedent(
        "UPDATE users SET status = 'inactive', status_info = 'migrated' WHERE is_active <> 't';"
    )
    conn.execute(query)

    op.alter_column('users', column_name='status', nullable=False)
    op.drop_column('users', 'is_active') 

def upgrade():
    inspector = reflection.Inspector.from_engine(op.get_bind())

    fks_to_cascade = {
        'sfc_flow_classifier_l7_parameters': 'classifier_id',
        'sfc_chain_group_associations': 'portchain_id',
        'sfc_port_chain_parameters': 'chain_id',
        'sfc_service_function_params': 'pair_id',
        'sfc_chain_classifier_associations': 'portchain_id'
    }

    for table, column in fks_to_cascade.items():
        fk_constraints = inspector.get_foreign_keys(table)
        for fk in fk_constraints:
            if column in fk['constrained_columns']:
                fk['options']['ondelete'] = 'CASCADE'
                migration.remove_foreign_keys(table, fk_constraints)
                migration.create_foreign_keys(table, fk_constraints) 

def downgrade():
    ''' Add the column _close_status to the table projects.
    '''
    engine = op.get_bind()
    Session = sa.orm.scoped_session(sa.orm.sessionmaker())
    Session.configure(bind=engine)
    session = Session()

    statuses = ['Invalid', 'Insufficient data', 'Fixed', 'Duplicate']
    for status in statuses:
        ticket_stat = model.StatusIssue(status=status)
        session.add(ticket_stat)
        session.commit()

    # Set the close_status for all the closed tickets
    op.execute('''UPDATE "issues" SET status=close_status where status != 'Open'; ''')

    # Remove the old status
    op.execute('''DELETE FROM "status_issue" WHERE status = 'Closed'; ''')

    op.drop_column('projects', '_close_status')
    op.drop_column('issues', 'close_status') 

def upgrade():
    op.add_column('watchers', sa.Column('watch_commits', sa.Boolean(),
                                        nullable=True))
    op.add_column('watchers', sa.Column('watch_issues', sa.Boolean(),
                                        nullable=True))
    # This section is to update the `watch_issues` and `watch_commits` columns
    # with the value of `watch`
    connection = op.get_bind()
    for watcher in connection.execute(watcher_helper.select()):
        connection.execute(
            watcher_helper.update().where(
                watcher_helper.c.id == watcher.id
            ).values(
                watch_issues=watcher.watch,
                watch_commits=False
            )
        )

    with op.batch_alter_table('watchers') as b:
        # Set nullable to False now that we've set values
        b.alter_column('watch_issues', nullable=False)
        b.alter_column('watch_commits', nullable=False)
        # Remove the watch column
        b.drop_column('watch') 

def downgrade():
    op.add_column('watchers', sa.Column('watch', sa.BOOLEAN(), nullable=True))

    # This section is to update the `watch` column with the value of
    # `watch_issues`
    connection = op.get_bind()
    for watcher in connection.execute(watcher_helper.select()):
        connection.execute(
            watcher_helper.update().where(
                watcher_helper.c.id == watcher.id
            ).values(
                watch=watcher.watch_issues
            )
        )

    with op.batch_alter_table('watchers') as b:
        # Set nullable to False now that we've set values
        b.alter_column('watch', nullable=False)
        # Drop the added columns
        b.drop_column('watch_issues')
        b.drop_column('watch_commits') 

def downgrade():
    engine = op.get_bind()
    Session = sa.orm.scoped_session(sa.orm.sessionmaker())
    Session.configure(bind=engine)
    session = Session()

    # Update all the existing projects
    for project in session.query(model.Project).all():
        # Install the default hook
        plugin = pagure.lib.plugins.get_plugin('default')
        dbobj = plugin.db_object()
        dbobj.active = False
        dbobj.project_id = project.id
        session.add(dbobj)
        session.flush()
        plugin.remove(project, dbobj)

        # Save the change
        session.commit() 

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('domains', sa.Column('integration_id', sa.String(length=512), nullable=True))
    op.alter_column('domains', 'total_resource_slots',
                    existing_type=postgresql.JSONB(astext_type=sa.Text()),
                    nullable=True)
    op.add_column('groups', sa.Column('integration_id', sa.String(length=512), nullable=True))
    op.add_column('groups', sa.Column('total_resource_slots',
                                      postgresql.JSONB(astext_type=sa.Text()), nullable=True))
    op.add_column('users', sa.Column('integration_id', sa.String(length=512), nullable=True))
    # ### end Alembic commandk ###

    print('\nSet group\'s total_resource_slots with empty dictionary.')
    query = textwrap.dedent('''\
        UPDATE groups SET total_resource_slots = '{}'::jsonb;
    ''')
    connection = op.get_bind()
    connection.execute(query) 

def downgrade():
    op.drop_index('ix_kernels_unique_sess_token', table_name='kernels')
    conn = op.get_bind()
    conn.execute('ALTER TYPE kernelstatus RENAME TO kernelstatus_new;')
    kernelstatus_old.create(conn)
    conn.execute(textwrap.dedent('''\
    ALTER TABLE kernels
        ALTER COLUMN "status" DROP DEFAULT,
        ALTER COLUMN "status" TYPE kernelstatus USING (
            CASE "status"::text
                WHEN 'CANCELLED' THEN 'TERMINATED'
                ELSE "status"::text
            END
        )::kernelstatus,
        ALTER COLUMN "status" SET DEFAULT 'PREPARING'::kernelstatus;
    DROP TYPE kernelstatus_new;
    '''))
    op.create_index(
        'ix_kernels_unique_sess_token', 'kernels', ['access_key', 'sess_id'],
        unique=True, postgresql_where=sa.text(
            "status != 'TERMINATED' and role = 'master'"
        )) 

def upgrade():
    conn = op.get_bind()
    conn.execute('DROP INDEX IF EXISTS ix_kernels_unique_sess_token;')
    conn.execute('ALTER TYPE kernelstatus RENAME TO kernelstatus_old;')
    kernelstatus_new.create(conn)
    conn.execute(textwrap.dedent('''\
    ALTER TABLE kernels
        ALTER COLUMN "status" DROP DEFAULT,
        ALTER COLUMN "status" TYPE kernelstatus USING "status"::text::kernelstatus,
        ALTER COLUMN "status" SET DEFAULT 'PENDING'::kernelstatus;
    DROP TYPE kernelstatus_old;
    '''))
    op.create_index(
        'ix_kernels_unique_sess_token', 'kernels', ['access_key', 'sess_id'],
        unique=True, postgresql_where=sa.text(
            "status NOT IN ('TERMINATED', 'CANCELLED') and role = 'master'"
        )) 

def downgrade():
    op.drop_column('kernels', 'stdout_port')
    op.drop_column('kernels', 'stdin_port')
    op.drop_column('kernels', 'repl_out_port')
    op.drop_column('kernels', 'repl_in_port')
    op.drop_column('kernels', 'mem_slot')
    op.drop_column('kernels', 'gpu_slot')
    op.drop_column('kernels', 'cpu_slot')
    op.drop_column('kernels', 'agent_addr')
    op.drop_index(op.f('ix_agents_status'), table_name='agents')
    op.drop_column('agents', 'status')
    op.drop_column('agents', 'lost_at')
    op.alter_column('kernels', column_name='status', type_=sa.String(length=64))
    op.add_column('kernels', sa.Column('allocated_cores', sa.ARRAY(sa.Integer), nullable=True))
    op.drop_column('kernels', 'cpu_set')
    op.drop_column('kernels', 'gpu_set')
    agentstatus.drop(op.get_bind())
    kernelstatus.drop(op.get_bind()) 

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('domains',
                  sa.Column('allowed_docker_registries',
                            postgresql.ARRAY(sa.String()), nullable=True))
    # ### end Alembic commands ###

    print('\nSet default allowed_docker_registries.')
    allowed_registries = os.environ.get('ALLOWED_DOCKER_REGISTRIES', None)
    if allowed_registries:
        allowed_registries = allowed_registries.replace(' ', '')
        allowed_registries = '{index.docker.io,' + allowed_registries + '}'
    else:
        allowed_registries = '{index.docker.io}'
    connection = op.get_bind()
    query = ("UPDATE domains SET allowed_docker_registries = '{}';".format(allowed_registries))
    connection.execute(query)

    op.alter_column('domains', column_name='allowed_docker_registries',
                    nullable=False) 

def downgrade():
	db_schema_migration.drop_columns('users', ('access_level',))
	db_schema_migration.drop_columns('credentials', ('regex_validated', 'mfa_token'))
	db_schema_migration.drop_columns('campaigns', ('credential_regex_mfa_token', 'credential_regex_password', 'credential_regex_username'))

	# adjust the schema version metadata
	db_manager.Session.remove()
	db_manager.Session.configure(bind=op.get_bind())
	session = db_manager.Session()
	db_manager.set_metadata('schema_version', 8, session=session)
	session.commit() 

def upgrade():
	op.drop_column('messages', 'company_name')

	db_manager.Session.remove()
	db_manager.Session.configure(bind=op.get_bind())
	session = db_manager.Session()
	db_manager.set_meta_data('schema_version', 5, session=session)
	session.commit() 

def upgrade():
	op.add_column('campaigns', sqlalchemy.Column('description', sqlalchemy.String))
	op.add_column('messages', sqlalchemy.Column('opener_ip', sqlalchemy.String))
	op.add_column('messages', sqlalchemy.Column('opener_user_agent', sqlalchemy.String))

	db_manager.Session.remove()
	db_manager.Session.configure(bind=op.get_bind())
	session = db_manager.Session()
	db_manager.set_meta_data('schema_version', 4, session=session)
	session.commit() 

def downgrade():
	op.drop_column('campaigns', 'description')
	op.drop_column('messages', 'opener_ip')
	op.drop_column('messages', 'opener_user_agent')

	db_manager.Session.remove()
	db_manager.Session.configure(bind=op.get_bind())
	session = db_manager.Session()
	db_manager.set_meta_data('schema_version', 3, session=session)
	session.commit() 

def get_session():
    engine = op.get_bind()
    Session = sa.orm.scoped_session(sa.orm.sessionmaker())
    Session.configure(bind=engine)
    return Session() 

def upgrade():
    connection = op.get_bind()

    ### commands auto generated by Alembic - please adjust! ###
    op.add_column('departments', sa.Column('short_name', sa.String(length=80), nullable=True))
    op.create_unique_constraint("dept_short_name", 'departments', ['short_name'])
    ### end Alembic commands ###