Multiple joins with two association tables in SQLAlchemy?

Viewed 105

I've this (simplified, but working) piece of code to handle the concept of workspaces, where every workspace can have members (users assigned to a workspace) and teams (also belonging to workspaces).

import json

from datetime import datetime
from sqlalchemy import create_engine  
from sqlalchemy import Column, String, DateTime, Boolean, Integer, ForeignKey
from sqlalchemy.ext.declarative import declarative_base  
from sqlalchemy.orm import sessionmaker, relationship, backref


db_uri = "postgresql://postgres:postgres@127.0.0.1/test_db"

db = create_engine(db_uri)  
Base = declarative_base()


class UserWorkspaceRole(Base):
    __tablename__ = "user_workspace_role"
    user_id = Column(Integer, ForeignKey('user.id'), primary_key=True)
    workspace_id = Column(Integer, ForeignKey('workspace.id'), primary_key=True)
    role_id = Column(Integer, ForeignKey('role.id'), primary_key=True)
    

class UserWorkspaceTeam(Base):
    __tablename__ = "user_workspace_team"
    user_id = Column(Integer, ForeignKey('user.id'), primary_key=True)
    workspace_id = Column(Integer, ForeignKey('workspace.id'), primary_key=True)
    team_id = Column(Integer, ForeignKey('team.id'), primary_key=True)


class User(Base):
    __tablename__ = "user"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String)
    team_workspaces = relationship(UserWorkspaceTeam, cascade="all, delete-orphan", backref="user")
    role_workspaces = relationship(UserWorkspaceRole, cascade="all, delete-orphan", backref="user")


class Workspace(Base):
    __tablename__ = "workspace"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String)
    team_users = relationship(UserWorkspaceRole, cascade="all, delete-orphan", backref="workspace")
    role_users = relationship(UserWorkspaceTeam, cascade="all, delete-orphan", backref="workspace")


class Role(Base):
    __tablename__ = "role"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String)
    role = relationship(UserWorkspaceRole, cascade="all, delete-orphan", backref="role")


class Team(Base):
    __tablename__ = "team"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String)
    team = relationship(UserWorkspaceTeam, cascade="all, delete-orphan", backref="team")


Session = sessionmaker(db)  
session = Session()

Base.metadata.create_all(db)
#Base.metadata.drop_all(db)

# Add roles
role_owner = Role(name="owner")
role_member = Role(name="member")
roles = [role_owner, role_member]
session.add_all(roles)
session.commit()

# Add teams
team_1 = Team(name="Team_1")
team_2 = Team(name="Team_2")
teams = [team_1, team_2]
session.add_all(teams)
session.commit()

# Add users
user1 = User(name="User_1")
user2 = User(name="User_2")
user3 = User(name="User_3")
users = [user1, user2, user3]
session.add_all(users)
session.commit()


workspace1 = Workspace(name="Apple")

test1 = UserWorkspaceRole(user=user1,workspace=workspace1, role_id="1")
test2 = UserWorkspaceRole(user=user2,workspace=workspace1, role_id="2")
test3 = UserWorkspaceRole(user=user3,workspace=workspace1, role_id="2")
test4 = UserWorkspaceTeam(user=user1,workspace=workspace1, team_id="1")
test5 = UserWorkspaceTeam(user=user1,workspace=workspace1, team_id="2")

records = [test1, test2, test3, test4, test5]
session.add_all(records)

session.commit()

x = session.query(UserWorkspaceRole).join(User).join(Workspace).join(Role).filter(Workspace.id == '1').all()

members = []

for n in range(len(x)):
    members.append(
            {
            "workspace": x[n].workspace.id,
            "name": x[n].user.name,
            "role": x[n].role.name
            }
        )
print(json.dumps(members, indent=4))

Which gives the following output:

[
    {
        "workspace": 1,
        "name": "User_1",
        "role": "owner"
    },
    {
        "workspace": 1,
        "name": "User_3",
        "role": "member"
    },
    {
        "workspace": 1,
        "name": "User_2",
        "role": "member"
    }
]

What I'd like to achieve is to add to every dictionary in the list, the teams a user belongs to in a certain workspace. Can I achieve this with one query? I tried to join also the UserWorkspaceTeam, but probably I'm missing something (I'm pretty newbie with db stuff).

The output I'd like should look like this:

[
    {
        "workspace": 1,
        "name": "User_1",
        "role": "owner",
        "teams": [
            {"id": "1", "name": "Team_1"},
            {"id": "2", "name": "Team_2"}
        ]
    },
    {
        "workspace": 1,
        "name": "User_3",
        "role": "member"
    },
    {
        "workspace": 1,
        "name": "User_2",
        "role": "member"
    }
]

How can I achieve that?

1 Answers

When using the range() function in your for loop the upper bound that you specified falls short by one. Try joining the UserWorkspaceTeam like you did before but this time add to the upper bound of the range function allowing you to reach the last index. Hopefully this helps.

x = session.query(UserWorkspaceRole).join(User).join(Workspace).join(Role).join(UserWorkspaceTeam).filter(Workspace.id == '1').all()

members = []

# The range() function includes values from 0 until a bound you specify but not including it.
for n in range(len(x)+1):
    members.append(
        {
        "workspace": x[n].workspace.id,
        "name": x[n].user.name,
        "role": x[n].role.name,
        }
    )
print(json.dumps(members, indent=4))
Related