Polymorphic Associations in SQLAlchemy

Table of Contents

• Polymorphic Associations in SQLAlchemy
• Table of Contents
• What this guide contains
• First Example: images belonging to either posts or comments (One-to-Many)
  • Database Diagram
  • The models
  • Downgrade fix
  • Setting up relationships
  • Relationship from Image-side
  • Seeding Images and a Route
• Second Example: A likes table connecting users to both posts and comments (Many-To-Many)
  • Database Diagram
  • The models
  • Setting up relationships through likes
  • Seeding Likes and a Route
• Epilogue

What this guide contains

There are two examples here of polymorphic associations in Flask. The first will be a One-to-Many style association where instances on one table could belong to one of two other tables. In this case, it will be an Image that could belong to either a Post or a Comment.

The second example is a bit more complex, as it involves a joins table for Likes. This will entail a Many-to-Many relationship from a User model, through a Like model, to either a Comment model or a Post model on a case-by-case basis.

I will include a database diagram, code snippets, and descriptions for both.

First Example: images belonging to either posts or comments (One-to-Many)

---

Database Diagram

This first example is in many ways easier to implement, though it still has some interesting code in it! Here's what we'll be working towards:

This may seem a bit daunting, but we've worked with most of it before. We have a users table, a posts table with a user_id that indicates who created the post, and a comments table with a user_id that indicates who created the comment as well as a post_id which indicates which post this comment is on.

The new addition is the images table. The url attribute is meant to hold the url via which we will be able to display the image on the frontend of our app. The important bit for the polymorphic association is that images will be associated with either a post or a comment via the combination of the imageable_id and the imageable_type, as we learned in the first part of this guide. That's why I've added a little box around those two attributes.

Now let's get into a bit of the code.

The models

The User, Post, and Comment models won't be groundbreaking for this to work. In fact, prior to setting up the associations, they won't change at all! So we'll skip those in the guide and assume that they're already in place. Starting with the Image model, it will look like this:

from .db import db
from sqlalchemy.dialects.postgresql import ENUM


class Image(db.Model):
  __tablename__ = 'images'

  id = db.Column(db.Integer, primary_key = True)
  url = db.Column(db.String, nullable = False)
  imageable_id = db.Column(db.Integer, nullable=False)
  imageable_type = db.Column(ENUM('post', 'comment', name='imageable_types'), nullable=False)
  
  def to_dict(self):
    return {
      "id": self.id,
      "url": self.url,
      "imageable_id": self.imageable_id,
      "imageable_type": self.imageable_type,
    }

A lot of this code is the same as what you're used to in our python projects. The first key difference is the imageable_id. It's not marked as a foreign key, even though it will act like one. That's because it's only half of the equation. We cannot mark it as the foreign key because then it will try to fill the whole role, and we need it to rely on imageable_type to tell it which model it should even be looking towards to match up the integer that it holds with the id of the other table.

The second difference is we're importing Enum from sqlalchemy.dialects.postgresql and using it for the imageable_type attribute. This could technically be replaced by a string if you want simpler functionality, but with an enum we'll ensure that we never get an image that has an imageable_type that we don't expect. In this case we only allow the values to be 'post' or 'comment'.

With this in place we're ready for our migration! You're ready to migrate and upgrade your database, though we'll want to jump over to that new migration quickly to ensure that the downgrade won't throw errors later.

Downgrade fix

Over in the new migration that alembic created for you, let's take a look at the generated code. The upgrade method looks perfect, creating all of our attributes as well as an enum that is set in your database for this project to use. However, there's a problem. The downgrade method drops the images table, but there's no mention of that imageable_types enum being dropped. Types like this enum aren't interwoven into the table like a lot of other attributes are. Instead, they sit in the same database, but outside of the table(s) they're used in. So we need to add some code again.

Start by adding

#...
from sqlalchemy.orm import sessionmaker

Session = sessionmaker()
#...

underneath from sqlalchemy.dialects import postgresql. Then, add the necessary lines below op.drop_table('images') to make your downgrade method look similar to this:

def downgrade():
  # ### commands auto generated by Alembic - please adjust! ###
  op.drop_table('images')
  bind = op.get_bind()
  session = Session(bind=bind)
  session.execute('DROP TYPE imageable_types')

What we're doing here is adding the ability to write raw SQL in the downgrade so that we can explicitly tell our database to drop the imageable_types type when this method is being run. If we hadn't made this fix, the next time we ran this downgrade and then ran the associated upgrade, it would've tried to create the same Type again, and thrown an error. With that fixed, future downgrades and upgrades shouldn't present a problem. Now let's head back to the models.

If you returned here from the Many-to-Many Likes portion of the guide, click here to return to where you left off.

Setting up relationships

We'll look at the Post model first, because it's as good a place as any. This relationship will be mostly similar to other relationships you've seen in SQLAlchemy in the past, with one key difference. We need to ensure that we're specifying that the relationship depends on both the imageable_id and the imageable_type as the "foreign key". It's that combination that is important. Here's what it will look like:

images = db.relationship(
  'Image',
  primaryjoin="and_(Image.imageable_type=='post', foreign(Image.imageable_id)==Post.id)",
  lazy="dynamic",
)

As you can see, it's a relationship from Post (the model we're currently looking at) to Image. We're specifying with that primaryjoin line, though, that the connection is on the combination of the imageable_id matching the id of the Post instance and the imageable_type being 'post'. This right here is what is preventing a comment with an id of 5 from being connected to images meant for a post with an id of 5, and vice versa.

With that, we're pretty much set! Throw a list comprehension in the Post model's to_dict() method like "images": [image.url for image in self.images] and you should have easy access to all of the images that are associated with the post.

The relationship for the comment model will look almost identical, except that you'll have comment instead of post and Comment instead of Comment. Make sure you change these things carefully, though, or you'll have some very confusing data to work with!

The next section will be dealing with how to setup a relationship from Image to Post and Comment, depending on which is needed. This doesn't make a whole lot of sense with with the current app, as normally I wouldn't want to query for images to get their associated post/comment information. However, I thought I'd include a workaround because it may be necessary on a similar model in a different application. Feel free to skip ahead for now if you're not concerned about that.

Relationship from Image-side

So this is a bit of a weird one. I'll readily admit that this is not in any sort of documentation that I read. It is very much a work-around, because I could not find anything that doesn't require a whole lot of other setup that wasn't necessary for anything else in this polymorphic relationship. On the plus side, it turns out it's pretty similar to how Sequelize handles these things (which really just makes me kind of disappointed with Sequelize).

Anyway, it works, so that's something. Let's look at the code and then we can discuss a few points.

  #...top of the Image model
  post = db.relationship(
      'Post',
      primaryjoin="and_(Image.imageable_type=='post', foreign(Image.imageable_id)==Post.id)",
      uselist=False
  )
  comment = db.relationship(
      'Comment',
      primaryjoin="and_(Image.imageable_type=='comment', foreign(Image.imageable_id)==Comment.id)",
      uselist=False
  )
  def parent(self):
    if self.imageable_type == 'post':
      return self.post.to_dict()
    elif self.imageable_type == 'comment':
      return self.comment.to_dict()
    else:
      return "Invalid imageable_type"

  def to_dict(self):
    return {
      "id": self.id,
      "url": self.url,
      "imageable_id": self.imageable_id,
      "imageable_type": self.imageable_type,
      "parent": self.parent()
    }

What we have here a lot of code that won't be used every time, but is necessary because it could potentially be used. So we start off with two relationships, one for Post and one for Comment. Only one of these relationships will be used for each instance of an Image, but both could be used.

Next we've got that parent method. The sole purpose of the method is to check what the imageable_type is and return a dictionary based on the appropriate relationship.

Finally we call self.parent() in the to_dict() method of the Image, and we've got the necessary information from either a Post or a Comment, depending on which model this image is tied to.

Seeding Images and a Route

Why even set this all up if we're not going to test it? This will just be a small sample to show how you can test to make sure everything is working. This is not good seed data and a near useless route, but you should have the tools and know-how to build something better once you're ready.

For the seeds, let's set up a new seed file for images like this:

from app.models import db, Image

# Adding seeded images to the database
def seed_images():

    image1 = Image(url="image_1_for_post_1", imageable_id=1, imageable_type="post")

    image2 = Image(url="image_2_for_post_1", imageable_id=1, imageable_type="post")

    image3 = Image(url="image_1_for_comment_1", imageable_id=1, imageable_type="comment")

    image4 = Image(url="image_1_for_comment_2", imageable_id=2, imageable_type="comment")

    db.session.add(image1)
    db.session.add(image2)
    db.session.add(image3)
    db.session.add(image4)

    db.session.commit()


def undo_images():
    db.session.execute('TRUNCATE images RESTART IDENTITY CASCADE;')
    db.session.commit()

This is a very simple example. It's just creating two images for posts and two images for comments. The urls obviously won't display actual images, but what's most important here anyway is seeding the imageable_id with an actual id and the imageable_type with either 'post' or 'comment', depending on your previous seed data and what you'd like to connect them to.

Please note that this does require you to set up other seed methods as well, so that there will be posts and comments in your database ready to be connected to images.
Then just add them to the session, commit it and you're done!

For a testing route, a good quick option is to set up a blueprint or use one that you already have, and then create a route that will take in an id param so you can toss it in a query or two and look at the results.

@test_routes.route('/<int:id>')
def test(id):
    image = Image.query.get(id)
    post = Post.query.get(id)
    comment = Post.query.get(id)
    return { 
        "image": image.to_dict(),
        "post": post.to_dict(),
        "comment": comment.to_dict()
    }
    

Just make sure that you import all of the models in the appropriate files and that you've seeded enough data to get results!

Second Example: A likes table connecting users to both posts and comments (Many-To-Many)

---

Database Diagram

The second example may be a lot to take in at first, but once it's set up it's very useful! Here's what we'll be working towards:

Similar to what we saw in the first example, here we have three tables you've likely seen before. We have a users table, a posts table with a user_id that indicates who created the post, and a comments table with a user_id that indicates who created the comment as well as a post_id which indicates which post this comment is on.

This time, the new addition takes the form of the likes table, which is a joins table between users and both posts and comments, though each like will be connected to either a post or a comment. One connection of the likes table could be considered normal. It's just a user_id that connects to the id of the users table. The other side of things will be a connection to either a post or a comment through the combination of likeable_id and likeable_type. That's what that little red box is for. Now that we've covered the thought process behind it, let's look at some code.

The models

The User, Post, and Comment models are again pretty standard. We'll skip those and assume that they already exist before the likes table is ever added (because they should anyway). Frankly, as long as you set them up with an id, they'll be good to go for this code.

Now for the Like model:

from .db import db
from sqlalchemy.dialects.postgresql import ENUM

class Like(db.Model):
  __tablename__ = 'likes'

  id = db.Column(db.Integer, primary_key = True)
  user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
  likeable_id = db.Column(db.Integer, nullable=False)
  likeable_type = db.Column(ENUM('post', 'comment', name='likeable_types'), nullable=False)

  def to_dict(self):
    return {
        "id": self.id,
        "user_id": self.user_id,
        "likeable_id": self.likeable_id,
        "likeable_type": self.likeable_type
    }

You may notice (if you read the first part of this guide), that this is very much like the Image model from earlier. One difference is that we have a user_id on this model, but that just acts as a normal foreign key. I'll recap a bit of what was said before just in case.

A lot of this code is the same as what you're used to in our python projects. The first key difference is the likeable_id. It's not marked as a foreign key, even though it will act like one. That's because it's only half of the equation. We cannot mark it as the foreign key because then it will try to fulfill the whole role, and we need it to rely on likeable_type to tell it which model it should even be looking towards to match up the integer that it holds with the id of the other table.

The second difference is we're importing Enum from sqlalchemy.dialects.postgresql and using it for the likeable_type attribute. This could technically be replaced by a string if you want simpler functionality, but with an enum we'll ensure that we never get a like that has a likeableable_type that we don't expect. In this case we only allow the values to be 'post' or 'comment'.

With this in place we’re ready to create our migration and upgrade the database! As with the Image model from earlier, once you've created the new migration, there's a downgrade fix that we'll have to add for the enum here. For the sake of relative brevity, you can click here to jump up to that. Just make sure to change "imageable" to "likeable" everywhere in this new migration.

Setting up relationships through likes

We're going to start the relationships on the User side of things, and then we're going to stop there. There's no reason in our app to bother with starting relationships in likes or connecting relationships from either side of the Many-to-Many relationship to likes, but if you're interested in doing so you can consult the first part of this walkthrough. With a little creativity you can see how the parts of that setup map to the parts of this one.

Here's the relationship from User to Post through Like:

posts_liked = db.relationship(
        "Post",
        secondary="likes",
        primaryjoin="foreign(Like.user_id)==User.id",
        secondaryjoin="and_(Like.likeable_type=='post', foreign(Like.likeable_id)==Post.id)",
        backref=db.backref("users_liked", lazy="dynamic"),
        lazy="dynamic"
    )

There's a bunch going on here, so let's break it down a bit. First of all, the relationship is to the Post model. This means that when we use this relationship, we'll end up with a list of instances of the Post model. It gives us easy access to all of the information for every Post that a User likes.

Of course, we'll need to go through the likes table to get here. Please note that the value for secondary is the likes table rather than the Like model.

The primaryjoin value is easy, just a standard foreign key pointing at the id of the model we're on.

The secondaryjoin is a bit more complicated. That connection is facilitated via the combination of a likeable_id matching the id of the current instance of a Post or Comment and a likeable_type dictating whether it's a Post or Comment for that Like.

The backref is there to make our lives easier by automatically setting up a relationship on Post to User via the same connections. This means that you can use the users_liked relationship in the to_dict() method of Post, as well as in routes or other parts of the backend.

The relationships relating to Comment are almost identical, just with 'comment' instead of 'post' and 'Comment' instead of 'Post'. With all of this set, we're able to at least set up a couple of list comprehensions in the to_dict() method of our User model.

 def to_dict(self):
    return {
      "id": self.id,
      "username": self.username,
      "email": self.email,
      "posts_liked": [post.to_dict() for post in self.posts_liked],
      "comments_liked": [comment.to_dict() for comment in self.comments_liked]
    }

Here we have two list comprehensions that will turn the return of our relationships to give us lists of objects built by the appropriate to_dict() methods of our other models. And we're all set!

Important Note: You should choose very carefully when you use to_dict() methods in another to_dict() method. It's very easy to start an endless loop between two models that are associated with each other, especially in a many-to-many relationship like this one. Often a better idea is to just use a list comprehension like [post.id for post in self.liked_posts], because with some Redux setup you can easily grab all the information you need from a slice of state with only an id. Use to_dict() methods like this very sparingly and figure out where you get the most benefit from using one.

Seeding Likes and a Route

Seeding through a Many-to-Many relationship can be difficult. You may have seen something similar to this when seeding through a helper table that joins two other tables:

# We can't use this
user = User.query.get(2)
comment = Comment.query.get(3)

user.comments_liked.append(comment)
db.session.commit()

Seeding via an existing association between two models works great for that, but unfortunately cannot help us here. Luckily, we created a model for our joins table, so we can do things the old fashioned way. Here's a sample seed_likes() method:

like1 = Like(user_id=2, likeable_id=1, likeable_type='post')
like2 = Like(user_id=3, likeable_id=2, likeable_type='post')
like3 = Like(user_id=2, likeable_id=1, likeable_type='comment')
like4 = Like(user_id=1, likeable_id=2, likeable_type='comment')

db.session.add(like1)
db.session.add(like2)
db.session.add(like3)
db.session.add(like4)
db.session.commit()

Here we have two new likes for some posts, and two new likes for some comments. Please note that while two of these likes have a likeable_id of 1 and two have a likeable_id of 2, the likeable_types dictate that they are actually connecting to different entities.

Now to the testing. We can use a very similar setup to what we had in the first part of this walkthrough by adding a test route to a new or existing blueprint that merely takes in an integer and uses it to query for some of our seed data.

@test_routes.route('/<int:id>')
def test(id):
    post = Post.query.get(id)
    comment = Comment.query.get(id)
    user = User.query.get(id)
    return {
        "post": post.to_dict(),
        "comment": comment.to_dict(),
        "user": user.to_dict(),
    }

This kind of testing is not only useful for seeing if the proper connections are in place, but also checking to make sure that our to_dict() methods won't cause an infinite loop.

Epilogue

---

That's it! You're done! You now know enough to build out some glorious polymorphic associations using SQLAlchemy. This is by no means a comprehensive guide to Polymorphism in the ORM. In fact, there is a lot of really cool code related to model class inheritance that you could learn and use. It's not necessary for this type of functionality, however, so if it interests you go check the documentation!

I hope this guide was helpful to you. If you'd like to see the Sequelize style of doing these very same setups, that walkthrough is here. I'm also working on a guide for polymorphic associations in a Ruby on Rails project, though the timeline for that is iffy.

If you have any feedback or questions, please do not hesitate to reach out to me over Slack. Have a wonderful day and best of luck on whichever project you're working on!

tags: SQLAlchemy Polymorphic Associations