Reusable SQLAlchemy Models

Recently looking at django, I took a liking to their idea of reusable apps. One thing that was common with most of these apps were some model that they would include and which you would call a function to automatically set up a relation to one of your models, thus having a "reusable model". An example would be any object that you'd like to have comments on automatically having a Comment model generated and relationships to your custom-built object automatically set up.

Out of mere curiosity, I wondered how difficult it would be to create a "reusable model" in SQLAlchemy. My end goal was to be able to do something like this...

@commentable
class Post(Base):
    __tablename__ = 'posts'
    id = sa.Column(sa.Integer, primary_key=True)
    text = sa.Column(sa.String)

A class that was decorated as commentable would automatically have a relation defined to contain multiple comments. After trying out a few ideas, I wrote a test for what I would want the end result to look like...

class TestModels(SATestCase):
    def test_make_comment(self):
        p = Post()
        p.text = 'SQLAlchemy is amazing!'

        text = 'I agree!'
        name = 'Mark'
        url = 'http://www.sqlalchemy.org/'
        
        c = Post.Comment()
        c.text = text
        c.name = name
        c.url = url 
        p.add_comment(c)
        Session.add(p)

        p = self.reload(p)

        self.assertEquals(len(p.comments), 1)
        c = p.comments[0]
        self.assertEquals(c.text, text)
        self.assertEquals(c.name, name)
        self.assertEquals(c.url, url)

(As a little note, self.reload() is a helper function that will force a complete reload of the session and return the objects you passed in after the session is opened again)

First off, a little bit about how the table structure would work. AFAIK, django stores comments in one large table and has what you might call a discriminator field to determine the object type that the row relates to. In my case, every type of commentable object (Post, NewItem, etc.) would instead get it's own comment class (PostComment, NewsItemComment) as well as their own table (post_comments, news_items_comments). No real reason to do it this way, I just thought it'd be easier.

In the end, it was actually pretty easy. Here is what the initial results look like...

class BaseComment(object):
    pass

def build_comment_model(clazz):
    class_table_name = str(class_mapper(clazz).local_table)
    metadata = clazz.metadata

    comment_class_name = clazz.__name__ + 'Comment'
    comment_class = type(comment_class_name, (BaseComment,), {}) 

    comment_table_name = class_table_name + '_comments'
    comment_table = sa.Table(comment_table_name, metadata,
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column(class_table_name + '_id',
            sa.Integer,
            sa.ForeignKey(class_table_name + '.id')),

        sa.Column('text', sa.String),
        sa.Column('name', sa.String(100)),
        sa.Column('url', sa.String(255)),
    )   

    mapper(comment_class, comment_table)

    return comment_class, comment_table

def commentable(clazz):
    comment_class, comment_table = build_comment_model(clazz)

    clazz.Comment = comment_class
    setattr(clazz, 'comments', relation(comment_class))

    def add_comment(self, comment):
        self.comments.append(comment)

    setattr(clazz, 'add_comment', add_comment)

    return clazz

First off, we have the BaseComment class. This is just an empty class, but you could easily see there being logic on here.

In the build_comment_model() function, you see where the creation of the table's metadata takes place. It'll first create a new class that will be used for the new comment model (in the case of a Post model being commentable, it's creating a PostComment class that inherits from BaseComment). Using SQLAlchemy information that we find on the mapper of the class we're decorating, we can determine things such as the name of the table so that we can create our own table (the Post model's table name is "posts", so we'll create a "posts_comments" table).

Finally, the commentable() function finishes by setting the Comment attribute on the model class we're decorating, as well as adding the relation and a class method. The Comment attribute allows us to easily get at the new class that we created, the relation allows us to easily work with the items, and the add_comment is an example of how we can also add extra methods to the model.

Currently, there are some pitfalls:

1.) It assumes that the class has been mapped before running. Because mapping of a class must happen after the class is instantiated, you couldn't use the commentable() function as a decorator for a non-declarative style mapping, since it wouldn't be able to find the mapper. The solution would be to just call the commentable function after the class is mapped. I guess it's not really a pitfall per se, although I haven't actually tried it, so I'm just guessing that it works :)

2.) Right now, it assumes that the foreign key for the Comment object would be a single column that corresponds to the column named "id" on whatever model it's decorating. I still need to work on this, since you might have a different name, different primary key type, or even a multi-column primary key. Basically when the new comment table's metadata is being laid out, it needs to look at the decorated model's primary keys to determine how it should construct it's foreign keys, rather than just jumping to the conclusions that it does now.