-
Notifications
You must be signed in to change notification settings - Fork 226
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix N+1 problem for one-to-many and many-to-many relationships (#254)
This optimization batches what used to be multiple SQL statements into a single SQL statement. For now, you'll have to enable the optimization via the `SQLAlchemyObjectType.Meta.connection_field_factory` (see `test_batching.py`).
- Loading branch information
Showing
8 changed files
with
458 additions
and
120 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,69 @@ | ||
import sqlalchemy | ||
from promise import dataloader, promise | ||
from sqlalchemy.orm import Session, strategies | ||
from sqlalchemy.orm.query import QueryContext | ||
|
||
|
||
def get_batch_resolver(relationship_prop): | ||
class RelationshipLoader(dataloader.DataLoader): | ||
cache = False | ||
|
||
def batch_load_fn(self, parents): # pylint: disable=method-hidden | ||
""" | ||
Batch loads the relationships of all the parents as one SQL statement. | ||
There is no way to do this out-of-the-box with SQLAlchemy but | ||
we can piggyback on some internal APIs of the `selectin` | ||
eager loading strategy. It's a bit hacky but it's preferable | ||
than re-implementing and maintainnig a big chunk of the `selectin` | ||
loader logic ourselves. | ||
The approach here is to build a regular query that | ||
selects the parent and `selectin` load the relationship. | ||
But instead of having the query emits 2 `SELECT` statements | ||
when callling `all()`, we skip the first `SELECT` statement | ||
and jump right before the `selectin` loader is called. | ||
To accomplish this, we have to construct objects that are | ||
normally built in the first part of the query in order | ||
to call directly `SelectInLoader._load_for_path`. | ||
TODO Move this logic to a util in the SQLAlchemy repo as per | ||
SQLAlchemy's main maitainer suggestion. | ||
See https://git.io/JewQ7 | ||
""" | ||
child_mapper = relationship_prop.mapper | ||
parent_mapper = relationship_prop.parent | ||
session = Session.object_session(parents[0]) | ||
|
||
# These issues are very unlikely to happen in practice... | ||
for parent in parents: | ||
# assert parent.__mapper__ is parent_mapper | ||
# All instances must share the same session | ||
assert session is Session.object_session(parent) | ||
# The behavior of `selectin` is undefined if the parent is dirty | ||
assert parent not in session.dirty | ||
|
||
loader = strategies.SelectInLoader(relationship_prop, (('lazy', 'selectin'),)) | ||
|
||
# Should the boolean be set to False? Does it matter for our purposes? | ||
states = [(sqlalchemy.inspect(parent), True) for parent in parents] | ||
|
||
# For our purposes, the query_context will only used to get the session | ||
query_context = QueryContext(session.query(parent_mapper.entity)) | ||
|
||
loader._load_for_path( | ||
query_context, | ||
parent_mapper._path_registry, | ||
states, | ||
None, | ||
child_mapper, | ||
) | ||
|
||
return promise.Promise.resolve([getattr(parent, relationship_prop.key) for parent in parents]) | ||
|
||
loader = RelationshipLoader() | ||
|
||
def resolve(root, info, **args): | ||
return loader.load(root) | ||
|
||
return resolve |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.