root/galaxy-central/lib/galaxy/model/migrate/versions/0047_job_table_user_id_column.py

"""
Add a user_id column to the job table.
"""

from sqlalchemy import *
from sqlalchemy.orm import *
from sqlalchemy.exc import *
from migrate import *
from migrate.changeset import *
from galaxy.model.custom_types import *

import datetime
now = datetime.datetime.utcnow

import sys, logging
log = logging.getLogger( __name__ )
log.setLevel(logging.DEBUG)
handler = logging.StreamHandler( sys.stdout )
format = "%(name)s %(levelname)s %(asctime)s %(message)s"
formatter = logging.Formatter( format )
handler.setFormatter( formatter )
log.addHandler( handler )

metadata = MetaData( migrate_engine )
db_session = scoped_session( sessionmaker( bind=migrate_engine, autoflush=False, autocommit=True ) )

def upgrade():
    print __doc__
    metadata.reflect()
    try:
        Job_table = Table( "job", metadata, autoload=True )
    except NoSuchTableError:
        Job_table = None
        log.debug( "Failed loading table job" )
    if Job_table:
        try:
            col = Column( "user_id", Integer, ForeignKey( "galaxy_user.id" ), index=True, nullable=True )
            col.create( Job_table )
            assert col is Job_table.c.user_id
        except Exception, e:
            log.debug( "Adding column 'user_id' to job table failed: %s" % ( str( e ) ) )
        try:
            i = Index( "ix_job_user_id", Job_table.c.user_id )
            i.create()
        except Exception, e:
            log.debug( "Adding index 'ix_job_user_id' failed: %s" % str( e ) )
        try:
            cmd = "SELECT job.id AS galaxy_job_id, " \
                + "galaxy_session.user_id AS galaxy_user_id " \
                + "FROM job " \
                + "JOIN galaxy_session ON job.session_id = galaxy_session.id;"
            job_users = db_session.execute( cmd ).fetchall()
            print "Updating user_id column in job table for ", len( job_users ), " rows..."
            print ""
            update_count = 0
            for row in job_users:
                if row.galaxy_user_id:
                    cmd = "UPDATE job SET user_id = %d WHERE id = %d" % ( int( row.galaxy_user_id ), int( row.galaxy_job_id ) )
                    update_count += 1
                db_session.execute( cmd )
            print "Updated the user_id column for ", update_count, " rows in the job table.  "
            print len( job_users ) - update_count, " rows have no user_id since the value was NULL in the galaxy_session table."
            print ""
        except Exception, e:
            log.debug( "Updating job.user_id column failed: %s" % str( e ) )
def downgrade():
    metadata.reflect()
    try:
        Job_table = Table( "job", metadata, autoload=True )
    except NoSuchTableError:
        Job_table = None
        log.debug( "Failed loading table job" )
    if Job_table:
        try:
            col = Job_table.c.user_id
            col.drop()
        except Exception, e:
            log.debug( "Dropping column 'user_id' from job table failed: %s" % ( str( e ) ) )