shell

Custom Shell

Working with database models needs a lot of use of the flask shell. You can access it with:

python manage.py shell

The default shell is quite unintuitive. It doesn’t pretty print the outputs and has no support for auto-completion. I’ve installed IPython that takes care of that. But still working with models means writing a lot of import statements. Plus if there was some change in the code related to the app, then the shell had to restarted again so the changes could be loaded. Meaning writing the import statements again.

We were using Flask-Script and I wanted to run a custom shell that imports all the required modules, models and helper functions, so I don’t have to write them over and over. Some of them were as long as:

from open_event.models.users_events_roles import UsersEventsRoles

So I created a custom shell command with different context that overrides the default shell provided by the Flask-Script Manager. It was pretty easy with Flask-Script. One thing I had to keep in mind is that it needed to be in a different file than manage.py. Since manage.py was committed to source repo, changes to it would be tracked. So I needed a different file that could be excluded from the source repo. I created an smg.py that imported the Manager from the open_event module and overrides the shell command.

from flask_script import Shell

from open_event import manager
from open_event.models.user import User
from open_event.models.event import Event
from open_event.helpers.data import save_to_db, delete_from_db

def _make_context():
    return dict(
        uq=User.query,
        eq=Event.query,
        su=User.query.get(1),
        User=User,
        Event=Event,
        savetodb=save_to_db,
        deletefromdb=delete_from_db
    )


if __name__ == "__main__":
    manager.add_command('shell', Shell(make_context=_make_context))
    manager.run()

Place this smg.py file in the same directory as manage.py, so you can access it with python smg.py shell.

The code is pretty simple to understand. We import the Shell class from flask_script, create its object with our context and then add it to the manager as a command. _make_context contains what I usually like to have in my shell. It must always return a dictionary. The keys of this dictionary would be available as statements inside the shell with their values specified here.

This helps a lot. Most of the time I would be working with the super_admin user, and I would need its User object from time to time. The super_admin user is always going to be the first user (User object with id 1). So instead of from open_event.models.user import User; su = User.query.get(1) I could just use the su variable. Models like User and Event are also readily available (so are their base queries). This is the default context that I always keep, but many times you need more models than the ones specified here. Like when I was working with the permissions system.

from open_event.models.users_events_roles import UsersEventsRoles
from open_event.models.service import Service
from open_event.models.role import Role

def _make_context():
    return dict(
        # usual stuff
        UER=UsersEventsRoles,
        Service=Service,
        Role=Role
    )

You can even write a script that fetches all the database models (instance of sqlalchemy Model class) and then add them to the _make_context dictionary. I like to keep it minimum and differently named, so there are no conflicts of classes when try to auto-complete.

One more thing, you need to exclude the smg.py file so that git doesn’t track it. You can simply add it in the .git/info/exclude file.

I also wrote some useful one-liner bash commands.

Revision History

Every time someone updates the database models, he needs to migrate and provide the migrations file to others by committing it to the source. These files help us upgrade our already defined database. We work with Alembic. Regarding alembic revisions (migration files) you can keep two things in mind. One is the Head, that keeps track of the latest revision, and another is Current, that specifies what revision your tables in the database are based on. If your current revision is not a head, it means your database tables are not up-to-date. You need to upgrade (python manage.py db upgrade). The can be multiple heads in the revision history. python manage.py heads displays all of them. The current revision can be fetched with python manage.py current. I wanted something that automatically checks if current is at a head.

python manage.py db history | grep --color "$(python manage.py db current 2> /dev/null)|$(python manage.py db heads)|$"

You can specify it as an alias in .bashrc. This command displays the revision history with the head and current revision being colored.

You can identify the head as it is always appended by “(head)”. The other colored revision would be the current.

You can also reverse the output, by printing bottom to up. This is helpful when the revision history gets large and you have to scroll up to see the head.

alias rev='python manage.py db history | tac | grep --color "$(python manage.py db current 2> /dev/null)|$(python manage.py db heads)|$"'

Current Revision changer

The current revision is maintained at the database in the one-column alembic_version table. If for some reason the migration file for current revision no longer exists (like when changing to a branch at git that doesn’t have the file), alembic would raise an error. So to change the revision at the database I wrote a bash function.

change_db_current () {
    echo "Current Revision: "
    read current_rev
    echo "New Revision: "
    read new_rev

    echo 'c test \ update alembic_version set version_num='"'"$new_rev"'"' where version_num='"'"$current_rev"'"';' | sudo su - postgres -c psql
}

Enter the current revision hash (displayed in the alembic error) and the revision you need to point the current to, and it will update the table with the new revision. Sick right! It’s pretty stupid actually. It was when I didn’t know how to stamp in alembic. Anyways, it was useful some time ago.