From f8dbb7696520193a8fdb6ff609c0ef9d032e7f19 Mon Sep 17 00:00:00 2001 From: collerek Date: Thu, 26 Nov 2020 06:33:24 +0100 Subject: [PATCH] add aliases to test prefetch_related --- ormar/models/modelproxy.py | 4 +--- ormar/queryset/prefetch_query.py | 3 ++- tests/test_prefetch_related.py | 15 ++++++++------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/ormar/models/modelproxy.py b/ormar/models/modelproxy.py index 1ae257e..22db178 100644 --- a/ormar/models/modelproxy.py +++ b/ormar/models/modelproxy.py @@ -236,9 +236,7 @@ class ModelTableProxy: @staticmethod def _populate_pk_column( - model: Type["Model"], - columns: List[str], - use_alias: bool = False, + model: Type["Model"], columns: List[str], use_alias: bool = False, ) -> List[str]: pk_alias = ( model.get_column_alias(model.Meta.pkname) diff --git a/ormar/queryset/prefetch_query.py b/ormar/queryset/prefetch_query.py index 9f1930c..ed4de9d 100644 --- a/ormar/queryset/prefetch_query.py +++ b/ormar/queryset/prefetch_query.py @@ -385,6 +385,7 @@ class PrefetchQuery: instance = self._populate_nested_related( model=instance, prefetch_dict=prefetch_dict, ) + field_db_name = target_model.get_column_alias(field_name) self.already_extracted[target_model.get_name()].setdefault( field_name, dict() - ).setdefault(row[field_name], []).append(instance) + ).setdefault(row[field_db_name], []).append(instance) diff --git a/tests/test_prefetch_related.py b/tests/test_prefetch_related.py index 378fdb4..cd0fb2e 100644 --- a/tests/test_prefetch_related.py +++ b/tests/test_prefetch_related.py @@ -17,7 +17,7 @@ class RandomSet(ormar.Model): metadata = metadata database = database - id: int = ormar.Integer(primary_key=True) + id: int = ormar.Integer(name='random_id', primary_key=True) name: str = ormar.String(max_length=100) @@ -28,7 +28,7 @@ class Tonation(ormar.Model): database = database id: int = ormar.Integer(primary_key=True) - name: str = ormar.String(max_length=100) + name: str = ormar.String(name='tonation_name', max_length=100) rand_set: Optional[RandomSet] = ormar.ForeignKey(RandomSet) @@ -38,7 +38,7 @@ class Division(ormar.Model): metadata = metadata database = database - id: int = ormar.Integer(primary_key=True) + id: int = ormar.Integer(name='division_id', primary_key=True) name: str = ormar.String(max_length=100) @@ -77,11 +77,11 @@ class Track(ormar.Model): metadata = metadata database = database - id: int = ormar.Integer(primary_key=True) + id: int = ormar.Integer(name='track_id', primary_key=True) album: Optional[Album] = ormar.ForeignKey(Album) title: str = ormar.String(max_length=100) position: int = ormar.Integer() - tonation: Optional[Tonation] = ormar.ForeignKey(Tonation) + tonation: Optional[Tonation] = ormar.ForeignKey(Tonation, name='tonation_id') class Cover(ormar.Model): @@ -91,7 +91,7 @@ class Cover(ormar.Model): database = database id: int = ormar.Integer(primary_key=True) - album: Optional[Album] = ormar.ForeignKey(Album, related_name="cover_pictures") + album: Optional[Album] = ormar.ForeignKey(Album, related_name="cover_pictures", name='album_id') title: str = ormar.String(max_length=100) artist: str = ormar.String(max_length=200, nullable=True) @@ -221,7 +221,8 @@ async def test_prefetch_related_with_select_related(): await Track.objects.create(album=album, title="Heart don't stand a chance", position=2, tonation=ton1) await Track.objects.create(album=album, title="The Waters", position=3, tonation=ton1) - album = await Album.objects.select_related('tracks__tonation__rand_set').filter(name='Malibu').prefetch_related( + album = await Album.objects.select_related('tracks__tonation__rand_set').filter( + name='Malibu').prefetch_related( ['cover_pictures', 'shops__division']).get() assert len(album.tracks) == 3 assert album.tracks[0].tonation == album.tracks[2].tonation == ton1