Recursive CTEs results as ORM model instances in SQLModel or SQLAlchemy

17 Nov, 2024

Given an ORM model like this:

class Venue(SQLModel, table=True):
    id: UUID = Field(primary_key=True, default_factory=uuid7)
    ...
    parent_id: UUID | None = Field(foreign_key="venue.id", default=None)
    parent: "Venue" | None = Relationship(back_populates="children")
    children: list["Venue"] = Relationship(
        back_populates="parent",
        sa_relationship_kwargs={"remote_side": "Venue.id"},
    )

We want to write a query selecting a given root element, and all of its children. In PostgreSQL, with recursive Common Table Expressions, we can write this:

-- Recurses for each row
WITH RECURSIVE all_venues AS (
    -- Get root row
    SELECT * FROM venue WHERE venue.id = :venue_id

    UNION ALL

    -- Get all of the root's direct descendants
    SELECT child.*
    FROM venue AS child
    INNER JOIN all_venues ON child.parent_id = all_venues.id
)
SELECT * FROM all_venues

However, on our application side, we want to get a set of Venue instances, rather than tuple/dict structures. Here's an approach using SQLAlchemy SQL Core (based on Sanjaya's post):

# The query to get the root row, that is then converted to a recursive CTE
cte = (
    select(Venue)
    .where(Venue.id == venue_id)
    .cte('all_venues', recursive=True)
)
# Then we build a query for the child venues, and create that union inside the CTE
child_venues = (
    select(Venue)
    .join(cte, Venue.parent_id == cte.c.id)
)
union_cte = cte.union_all(child_venues)
# Finally we select from the CTE
qry = select(union_cte)
results = (await sess.exec(qry)).fetchall()

The resulting query from this matches up with our original raw query:

WITH RECURSIVE all_venues(id, ..., parent_id) AS
  (
    SELECT venue.id AS id, ..., venue.parent_id AS parent_id
    FROM venue
    WHERE venue.id = :id_1
  UNION ALL
    SELECT venue.id AS id, ..., venue.parent_id AS parent_id
    FROM venue JOIN all_venues ON venue.parent_id = all_venues.id
  )
SELECT all_venues.id, ..., all_venues.parent_id
FROM all_venues

Unfortunately, we'll get an array of Venue.id rather than the instance. I believe this is taking the first column of each row of union_cte. We could make another query to fetch rows for all these IDs, but it's nicer if we can do this all in one query.

To coerce these into Venue instances, we need to change the way we query from the CTE:

# As above
cte = (...)
child_venues = (...)
# Instruct the ORM to select Venues from the CTE, 
qry = (
    select(Venue)
    .join(
        cte.union_all(child_venues),
        Venue.id == cte.c.id
    )
)

I expected the query to be wrapped with yet another select, but in practice it's just a table join, which to my understanding acts pretty much as a where in:

WITH RECURSIVE all_venues(id, slug, name, latlng, parent_id) AS
  (
    SELECT venue.id AS id, ..., venue.parent_id AS parent_id
    FROM venue
    WHERE venue.id = :id_1
  UNION ALL
    SELECT venue.id AS id, ..., venue.parent_id AS parent_id
    FROM venue
    JOIN all_venues ON venue.parent_id = all_venues.id
  )
SELECT venue.id, ..., venue.parent_id
FROM venue
JOIN all_venues ON venue.id = all_venues.id

Et voilà!

[
  Venue(id=UUID('067...301'), parent_venue_id=None, ...),
  Venue(id=UUID('067...b83'), parent_venue_id=UUID('067...301'), ...)
]

Side note, after an hour of fruitless searching, Claude 3.5 Sonnet answered this correctly on the first try, with minimal prompting - impressive!