From 6e67b69385f39a9e0a99b93fb45bb132bef41528 Mon Sep 17 00:00:00 2001 From: collerek Date: Mon, 14 Dec 2020 15:36:04 +0100 Subject: [PATCH] fix for issue 73 --- docs/releases.md | 7 + ormar/fields/base.py | 7 +- ormar/fields/foreign_key.py | 39 ++++-- ormar/models/metaclass.py | 53 +++++-- ormar/models/model.py | 41 +++--- ormar/models/modelproxy.py | 19 ++- ormar/models/newbasemodel.py | 8 +- ormar/queryset/clause.py | 16 ++- ormar/queryset/join.py | 8 +- ormar/queryset/prefetch_query.py | 5 +- ormar/relations/alias_manager.py | 34 +++-- ormar/relations/relation_manager.py | 10 +- ormar/relations/utils.py | 7 +- tests/test_same_table_joins.py | 2 +- ...select_related_with_m2m_and_pk_name_set.py | 130 ++++++++++++++++++ 15 files changed, 306 insertions(+), 80 deletions(-) create mode 100644 tests/test_select_related_with_m2m_and_pk_name_set.py diff --git a/docs/releases.md b/docs/releases.md index b566795..ea17b2f 100644 --- a/docs/releases.md +++ b/docs/releases.md @@ -1,3 +1,10 @@ +# 0.7.4 + +* Allow multiple relations to the same related model/table. +* Fix for wrong relation column used in many_to_many relation joins (fix [#71][#71]) +* Fix for wrong relation population for m2m relations when also fk relation present for same model. +* Add check if user provide related_name if there are multiple relations to same table on one model + # 0.7.3 * Fix for setting fetching related model with UUDI pk, which is a string in raw (fix [#71][#71]) diff --git a/ormar/fields/base.py b/ormar/fields/base.py index 343e7f9..db9ae57 100644 --- a/ormar/fields/base.py +++ b/ormar/fields/base.py @@ -25,6 +25,7 @@ class BaseField(FieldInfo): """ __type__ = None + related_name = None column_type: sqlalchemy.Column constraints: List = [] @@ -222,7 +223,11 @@ class BaseField(FieldInfo): @classmethod def expand_relationship( - cls, value: Any, child: Union["Model", "NewBaseModel"], to_register: bool = True + cls, + value: Any, + child: Union["Model", "NewBaseModel"], + to_register: bool = True, + relation_name: str = None, ) -> Any: """ Function overwritten for relations, in basic field the value is returned as is. diff --git a/ormar/fields/foreign_key.py b/ormar/fields/foreign_key.py index 88e1928..daf9137 100644 --- a/ormar/fields/foreign_key.py +++ b/ormar/fields/foreign_key.py @@ -94,35 +94,40 @@ class ForeignKeyField(BaseField): @classmethod def _extract_model_from_sequence( - cls, value: List, child: "Model", to_register: bool + cls, value: List, child: "Model", to_register: bool, relation_name: str ) -> List["Model"]: return [ - cls.expand_relationship(val, child, to_register) # type: ignore + cls.expand_relationship( + value=val, + child=child, + to_register=to_register, + relation_name=relation_name, + ) # type: ignore for val in value ] @classmethod def _register_existing_model( - cls, value: "Model", child: "Model", to_register: bool + cls, value: "Model", child: "Model", to_register: bool, relation_name: str ) -> "Model": if to_register: - cls.register_relation(value, child) + cls.register_relation(model=value, child=child, relation_name=relation_name) return value @classmethod def _construct_model_from_dict( - cls, value: dict, child: "Model", to_register: bool + cls, value: dict, child: "Model", to_register: bool, relation_name: str ) -> "Model": if len(value.keys()) == 1 and list(value.keys())[0] == cls.to.Meta.pkname: value["__pk_only__"] = True model = cls.to(**value) if to_register: - cls.register_relation(model, child) + cls.register_relation(model=model, child=child, relation_name=relation_name) return model @classmethod def _construct_model_from_pk( - cls, value: Any, child: "Model", to_register: bool + cls, value: Any, child: "Model", to_register: bool, relation_name: str ) -> "Model": if cls.to.pk_type() == uuid.UUID and isinstance(value, str): value = uuid.UUID(value) @@ -134,18 +139,28 @@ class ForeignKeyField(BaseField): ) model = create_dummy_instance(fk=cls.to, pk=value) if to_register: - cls.register_relation(model, child) + cls.register_relation(model=model, child=child, relation_name=relation_name) return model @classmethod - def register_relation(cls, model: "Model", child: "Model") -> None: + def register_relation( + cls, model: "Model", child: "Model", relation_name: str + ) -> None: model._orm.add( - parent=model, child=child, child_name=cls.related_name, virtual=cls.virtual + parent=model, + child=child, + child_name=cls.related_name or child.get_name() + "s", + virtual=cls.virtual, + relation_name=relation_name, ) @classmethod def expand_relationship( - cls, value: Any, child: Union["Model", "NewBaseModel"], to_register: bool = True + cls, + value: Any, + child: Union["Model", "NewBaseModel"], + to_register: bool = True, + relation_name: str = None, ) -> Optional[Union["Model", List["Model"]]]: if value is None: return None if not cls.virtual else [] @@ -158,5 +173,5 @@ class ForeignKeyField(BaseField): model = constructors.get( # type: ignore value.__class__.__name__, cls._construct_model_from_pk - )(value, child, to_register) + )(value, child, to_register, relation_name) return model diff --git a/ormar/models/metaclass.py b/ormar/models/metaclass.py index 5649d22..90c5e11 100644 --- a/ormar/models/metaclass.py +++ b/ormar/models/metaclass.py @@ -42,16 +42,19 @@ class ModelMeta: signals: SignalEmitter -def register_relation_on_build(table_name: str, field: Type[ForeignKeyField]) -> None: - alias_manager.add_relation_type(field.to.Meta.tablename, table_name) +def register_relation_on_build_new(new_model: Type["Model"], field_name: str) -> None: + alias_manager.add_relation_type_new(new_model, field_name) -def register_many_to_many_relation_on_build( - table_name: str, field: Type[ManyToManyField] +def register_many_to_many_relation_on_build_new( + new_model: Type["Model"], field: Type[ManyToManyField] ) -> None: - alias_manager.add_relation_type(field.through.Meta.tablename, table_name) - alias_manager.add_relation_type( - field.through.Meta.tablename, field.to.Meta.tablename + + alias_manager.add_relation_type_new( + field.through, new_model.get_name(), is_multi=True + ) + alias_manager.add_relation_type_new( + field.through, field.to.get_name(), is_multi=True ) @@ -161,8 +164,27 @@ def check_pk_column_validity( return field_name +def validate_related_names_in_relations( + model_fields: Dict, new_model: Type["Model"] +) -> None: + already_registered: Dict[str, List[Optional[str]]] = dict() + for field in model_fields.values(): + if issubclass(field, ForeignKeyField): + previous_related_names = already_registered.setdefault(field.to, []) + if field.related_name in previous_related_names: + raise ModelDefinitionError( + f"Multiple fields declared on {new_model.get_name(lower=False)} " + f"model leading to {field.to.get_name(lower=False)} model without " + f"related_name property set. \nThere can be only one relation with " + f"default/empty name: '{new_model.get_name() + 's'}'" + f"\nTip: provide different related_name for FK and/or M2M fields" + ) + else: + previous_related_names.append(field.related_name) + + def sqlalchemy_columns_from_model_fields( - model_fields: Dict, table_name: str + model_fields: Dict, table_name: str, new_model: Type["Model"] ) -> Tuple[Optional[str], List[sqlalchemy.Column]]: columns = [] pkname = None @@ -172,6 +194,7 @@ def sqlalchemy_columns_from_model_fields( "Table {table_name} had no fields so auto " "Integer primary key named `id` created." ) + validate_related_names_in_relations(model_fields, new_model) for field_name, field in model_fields.items(): if field.primary_key: pkname = check_pk_column_validity(field_name, field, pkname) @@ -181,17 +204,16 @@ def sqlalchemy_columns_from_model_fields( and not issubclass(field, ManyToManyField) ): columns.append(field.get_column(field.get_alias())) - register_relation_in_alias_manager(table_name, field) return pkname, columns -def register_relation_in_alias_manager( - table_name: str, field: Type[ForeignKeyField] +def register_relation_in_alias_manager_new( + new_model: Type["Model"], field: Type[ForeignKeyField], field_name: str ) -> None: if issubclass(field, ManyToManyField): - register_many_to_many_relation_on_build(table_name, field) + register_many_to_many_relation_on_build_new(new_model=new_model, field=field) elif issubclass(field, ForeignKeyField): - register_relation_on_build(table_name, field) + register_relation_on_build_new(new_model=new_model, field_name=field_name) def populate_default_pydantic_field_value( @@ -255,7 +277,7 @@ def populate_meta_tablename_columns_and_pk( pkname = new_model.Meta.pkname else: pkname, columns = sqlalchemy_columns_from_model_fields( - new_model.Meta.model_fields, new_model.Meta.tablename + new_model.Meta.model_fields, new_model.Meta.tablename, new_model ) if pkname is None: @@ -263,7 +285,6 @@ def populate_meta_tablename_columns_and_pk( new_model.Meta.columns = columns new_model.Meta.pkname = pkname - return new_model @@ -379,6 +400,8 @@ class ModelMetaclass(pydantic.main.ModelMetaclass): new_model = populate_meta_tablename_columns_and_pk(name, new_model) new_model = populate_meta_sqlalchemy_table_if_required(new_model) expand_reverse_relationships(new_model) + for field_name, field in new_model.Meta.model_fields.items(): + register_relation_in_alias_manager_new(new_model, field, field_name) populate_choices_validators(new_model) if new_model.Meta.pkname not in attrs["__annotations__"]: field_name = new_model.Meta.pkname diff --git a/ormar/models/model.py b/ormar/models/model.py index 8e7a8f5..12b5d80 100644 --- a/ormar/models/model.py +++ b/ormar/models/model.py @@ -58,7 +58,8 @@ class Model(NewBaseModel): row: sqlalchemy.engine.ResultProxy, select_related: List = None, related_models: Any = None, - previous_table: str = None, + previous_model: Type[T] = None, + related_name: str = None, fields: Optional[Union[Dict, Set]] = None, exclude_fields: Optional[Union[Dict, Set]] = None, ) -> Optional[T]: @@ -69,28 +70,32 @@ class Model(NewBaseModel): if select_related: related_models = group_related_list(select_related) - if ( - previous_table - and previous_table in cls.Meta.model_fields - and issubclass(cls.Meta.model_fields[previous_table], ManyToManyField) - ): - previous_table = cls.Meta.model_fields[ - previous_table - ].through.Meta.tablename + rel_name2 = related_name - if previous_table: - table_prefix = cls.Meta.alias_manager.resolve_relation_join( - previous_table, cls.Meta.table.name + if ( + previous_model + and related_name + and issubclass( + previous_model.Meta.model_fields[related_name], ManyToManyField + ) + ): + through_field = previous_model.Meta.model_fields[related_name] + rel_name2 = previous_model.resolve_relation_name( + through_field.through, through_field.to, explicit_multi=True + ) + previous_model = through_field.through # type: ignore + + if previous_model and rel_name2: + table_prefix = cls.Meta.alias_manager.resolve_relation_join_new( + previous_model, rel_name2 ) else: table_prefix = "" - previous_table = cls.Meta.table.name item = cls.populate_nested_models_from_row( item=item, row=row, related_models=related_models, - previous_table=previous_table, fields=fields, exclude_fields=exclude_fields, ) @@ -111,7 +116,6 @@ class Model(NewBaseModel): instance.set_save_status(True) else: instance = None - return instance @classmethod @@ -120,7 +124,6 @@ class Model(NewBaseModel): item: dict, row: sqlalchemy.engine.ResultProxy, related_models: Any, - previous_table: sqlalchemy.Table, fields: Optional[Union[Dict, Set]] = None, exclude_fields: Optional[Union[Dict, Set]] = None, ) -> dict: @@ -135,7 +138,8 @@ class Model(NewBaseModel): child = model_cls.from_row( row, related_models=remainder, - previous_table=previous_table, + previous_model=cls, + related_name=related, fields=fields, exclude_fields=exclude_fields, ) @@ -146,7 +150,8 @@ class Model(NewBaseModel): exclude_fields = cls.get_excluded(exclude_fields, related) child = model_cls.from_row( row, - previous_table=previous_table, + previous_model=cls, + related_name=related, fields=fields, exclude_fields=exclude_fields, ) diff --git a/ormar/models/modelproxy.py b/ormar/models/modelproxy.py index b545ecf..d3e7435 100644 --- a/ormar/models/modelproxy.py +++ b/ormar/models/modelproxy.py @@ -21,7 +21,7 @@ from ormar.exceptions import ModelPersistenceError, RelationshipInstanceError from ormar.queryset.utils import translate_list_to_dict, update import ormar # noqa: I100 -from ormar.fields import BaseField +from ormar.fields import BaseField, ManyToManyField from ormar.fields.foreign_key import ForeignKeyField from ormar.models.metaclass import ModelMeta @@ -278,12 +278,21 @@ class ModelTableProxy: "ModelTableProxy", Type["ModelTableProxy"], ], + explicit_multi: bool = False, ) -> str: for name, field in item.Meta.model_fields.items(): - if issubclass(field, ForeignKeyField): - # fastapi is creating clones of response model - # that's why it can be a subclass of the original model - # so we need to compare Meta too as this one is copied as is + # fastapi is creating clones of response model + # that's why it can be a subclass of the original model + # so we need to compare Meta too as this one is copied as is + if issubclass(field, ManyToManyField): + attrib = "to" if not explicit_multi else "through" + if ( + getattr(field, attrib) == related.__class__ + or getattr(field, attrib).Meta == related.Meta + ): + return name + + elif issubclass(field, ForeignKeyField): if field.to == related.__class__ or field.to.Meta == related.Meta: return name diff --git a/ormar/models/newbasemodel.py b/ormar/models/newbasemodel.py index ff3f8b0..abe726e 100644 --- a/ormar/models/newbasemodel.py +++ b/ormar/models/newbasemodel.py @@ -99,7 +99,7 @@ class NewBaseModel( k: self._convert_json( k, self.Meta.model_fields[k].expand_relationship( - v, self, to_register=False + v, self, to_register=False, relation_name=k ), "dumps", ) @@ -128,7 +128,7 @@ class NewBaseModel( # register the columns models after initialization for related in self.extract_related_names(): self.Meta.model_fields[related].expand_relationship( - new_kwargs.get(related), self, to_register=True + new_kwargs.get(related), self, to_register=True, relation_name=related ) def __setattr__(self, name: str, value: Any) -> None: # noqa CCR001 @@ -138,7 +138,9 @@ class NewBaseModel( object.__setattr__(self, self.Meta.pkname, value) self.set_save_status(False) elif name in self._orm: - model = self.Meta.model_fields[name].expand_relationship(value, self) + model = self.Meta.model_fields[name].expand_relationship( + value=value, child=self, relation_name=name + ) if isinstance(self.__dict__.get(name), list): # virtual foreign key or many to many self.__dict__[name].append(model) diff --git a/ormar/queryset/clause.py b/ormar/queryset/clause.py index 362ba85..e5f84f7 100644 --- a/ormar/queryset/clause.py +++ b/ormar/queryset/clause.py @@ -131,17 +131,19 @@ class QueryClause: # Walk the relationships to the actual model class # against which the comparison is being made. - previous_table = model_cls.Meta.tablename + previous_model = model_cls for part in related_parts: + part2 = part if issubclass(model_cls.Meta.model_fields[part], ManyToManyField): - previous_table = model_cls.Meta.model_fields[ - part - ].through.Meta.tablename - current_table = model_cls.Meta.model_fields[part].to.Meta.tablename + through_field = model_cls.Meta.model_fields[part] + previous_model = through_field.through + part2 = model_cls.resolve_relation_name( + through_field.through, through_field.to, explicit_multi=True + ) manager = model_cls.Meta.alias_manager - table_prefix = manager.resolve_relation_join(previous_table, current_table) + table_prefix = manager.resolve_relation_join_new(previous_model, part2) model_cls = model_cls.Meta.model_fields[part].to - previous_table = current_table + previous_model = model_cls return select_related, table_prefix, model_cls def _compile_clause( diff --git a/ormar/queryset/join.py b/ormar/queryset/join.py index 1628017..fa32fd6 100644 --- a/ormar/queryset/join.py +++ b/ormar/queryset/join.py @@ -135,8 +135,8 @@ class SqlJoin: model_cls = join_params.model_cls.Meta.model_fields[part].to to_table = model_cls.Meta.table.name - alias = model_cls.Meta.alias_manager.resolve_relation_join( - join_params.from_table, to_table + alias = model_cls.Meta.alias_manager.resolve_relation_join_new( + join_params.prev_model, part ) if alias not in self.used_aliases: self._process_join( @@ -267,7 +267,9 @@ class SqlJoin: model_cls, join_params.prev_model ) to_key = model_cls.get_column_alias(to_field) - from_key = join_params.prev_model.get_column_alias(model_cls.Meta.pkname) + from_key = join_params.prev_model.get_column_alias( + join_params.prev_model.Meta.pkname + ) else: to_key = model_cls.get_column_alias(model_cls.Meta.pkname) from_key = join_params.prev_model.get_column_alias(part) diff --git a/ormar/queryset/prefetch_query.py b/ormar/queryset/prefetch_query.py index 13ad785..1a06d31 100644 --- a/ormar/queryset/prefetch_query.py +++ b/ormar/queryset/prefetch_query.py @@ -318,9 +318,8 @@ class PrefetchQuery: if issubclass(target_field, ManyToManyField): query_target = target_field.through select_related = [target_name] - table_prefix = target_field.to.Meta.alias_manager.resolve_relation_join( - from_table=query_target.Meta.tablename, - to_table=target_field.to.Meta.tablename, + table_prefix = target_field.to.Meta.alias_manager.resolve_relation_join_new( + query_target, target_name ) self.already_extracted.setdefault(target_name, {})["prefix"] = table_prefix diff --git a/ormar/relations/alias_manager.py b/ormar/relations/alias_manager.py index 2eead1e..477217f 100644 --- a/ormar/relations/alias_manager.py +++ b/ormar/relations/alias_manager.py @@ -1,11 +1,14 @@ import string import uuid from random import choices -from typing import Dict, List +from typing import Dict, List, TYPE_CHECKING, Type import sqlalchemy from sqlalchemy import text +if TYPE_CHECKING: # pragma: no cover + from ormar import Model + def get_table_alias() -> str: alias = "".join(choices(string.ascii_uppercase, k=2)) + uuid.uuid4().hex[:4] @@ -15,6 +18,7 @@ def get_table_alias() -> str: class AliasManager: def __init__(self) -> None: self._aliases: Dict[str, str] = dict() + self._aliases_new: Dict[str, str] = dict() @staticmethod def prefixed_columns( @@ -35,11 +39,25 @@ class AliasManager: def prefixed_table_name(alias: str, name: str) -> text: return text(f"{name} {alias}_{name}") - def add_relation_type(self, to_table_name: str, table_name: str,) -> None: - if f"{table_name}_{to_table_name}" not in self._aliases: - self._aliases[f"{table_name}_{to_table_name}"] = get_table_alias() - if f"{to_table_name}_{table_name}" not in self._aliases: - self._aliases[f"{to_table_name}_{table_name}"] = get_table_alias() + def add_relation_type_new( + self, source_model: Type["Model"], relation_name: str, is_multi: bool = False + ) -> None: + parent_key = f"{source_model.get_name()}_{relation_name}" + if parent_key not in self._aliases_new: + self._aliases_new[parent_key] = get_table_alias() + to_field = source_model.Meta.model_fields[relation_name] + child_model = to_field.to + related_name = to_field.related_name + if not related_name: + related_name = child_model.resolve_relation_name( + child_model, source_model, explicit_multi=is_multi + ) + child_key = f"{child_model.get_name()}_{related_name}" + if child_key not in self._aliases_new: + self._aliases_new[child_key] = get_table_alias() - def resolve_relation_join(self, from_table: str, to_table: str) -> str: - return self._aliases.get(f"{from_table}_{to_table}", "") + def resolve_relation_join_new( + self, from_model: Type["Model"], relation_name: str + ) -> str: + alias = self._aliases_new.get(f"{from_model.get_name()}_{relation_name}", "") + return alias diff --git a/ormar/relations/relation_manager.py b/ormar/relations/relation_manager.py index 81183f4..6462e48 100644 --- a/ormar/relations/relation_manager.py +++ b/ormar/relations/relation_manager.py @@ -56,8 +56,14 @@ class RelationsManager: return None @staticmethod - def add(parent: "Model", child: "Model", child_name: str, virtual: bool) -> None: - to_field: Type[BaseField] = child.resolve_relation_field(child, parent) + def add( + parent: "Model", + child: "Model", + child_name: str, + virtual: bool, + relation_name: str, + ) -> None: + to_field: Type[BaseField] = child.Meta.model_fields[relation_name] (parent, child, child_name, to_name,) = get_relations_sides_and_names( to_field, parent, child, child_name, virtual diff --git a/ormar/relations/utils.py b/ormar/relations/utils.py index c7a3fff..9fa09b0 100644 --- a/ormar/relations/utils.py +++ b/ormar/relations/utils.py @@ -18,8 +18,11 @@ def get_relations_sides_and_names( to_name = to_field.name if issubclass(to_field, ManyToManyField): child_name, to_name = ( - child.resolve_relation_name(parent, child), - child.resolve_relation_name(child, parent), + to_field.related_name + or child.resolve_relation_name( + parent, to_field.through, explicit_multi=True + ), + to_name, ) child = proxy(child) elif virtual: diff --git a/tests/test_same_table_joins.py b/tests/test_same_table_joins.py index 2375613..afba731 100644 --- a/tests/test_same_table_joins.py +++ b/tests/test_same_table_joins.py @@ -102,7 +102,7 @@ async def test_model_multiple_instances_of_same_table_in_schema(): async with database.transaction(force_rollback=True): await create_data() classes = await SchoolClass.objects.select_related( - ["teachers__category", "students"] + ["teachers__category", "students__schoolclass"] ).all() assert classes[0].name == "Math" assert classes[0].students[0].name == "Jane" diff --git a/tests/test_select_related_with_m2m_and_pk_name_set.py b/tests/test_select_related_with_m2m_and_pk_name_set.py new file mode 100644 index 0000000..36aba19 --- /dev/null +++ b/tests/test_select_related_with_m2m_and_pk_name_set.py @@ -0,0 +1,130 @@ +# type: ignore +from datetime import date +from typing import List, Optional, Union + +import databases +import pytest +import sqlalchemy +from sqlalchemy import create_engine + +import ormar +from ormar import ModelDefinitionError +from tests.settings import DATABASE_URL + +database = databases.Database(DATABASE_URL) +metadata = sqlalchemy.MetaData() + + +class MainMeta(ormar.ModelMeta): + metadata = metadata + database = database + + +class Role(ormar.Model): + class Meta(MainMeta): + pass + + name: str = ormar.Text(primary_key=True) + order: int = ormar.Integer(default=0, name="sort_order") + description: str = ormar.Text() + + +class Company(ormar.Model): + class Meta(MainMeta): + pass + + name: str = ormar.Text(primary_key=True) + + +class UserRoleCompany(ormar.Model): + class Meta(MainMeta): + pass + + +class User(ormar.Model): + class Meta(MainMeta): + pass + + registrationnumber: str = ormar.Text(primary_key=True) + company: Company = ormar.ForeignKey(Company) + company2: Company = ormar.ForeignKey(Company, related_name="secondary_users") + name: str = ormar.Text() + role: Optional[Role] = ormar.ForeignKey(Role) + roleforcompanies: Optional[Union[Company, List[Company]]] = ormar.ManyToMany( + Company, through=UserRoleCompany, related_name="role_users" + ) + lastupdate: date = ormar.DateTime(server_default=sqlalchemy.func.now()) + + +@pytest.fixture(autouse=True, scope="module") +def create_test_database(): + engine = create_engine(DATABASE_URL) + metadata.create_all(engine) + yield + metadata.drop_all(engine) + + +def test_wrong_model(): + with pytest.raises(ModelDefinitionError): + + class User(ormar.Model): + class Meta(MainMeta): + pass + + registrationnumber: str = ormar.Text(primary_key=True) + company: Company = ormar.ForeignKey(Company) + company2: Company = ormar.ForeignKey(Company) + + +@pytest.mark.asyncio +async def test_create_primary_models(): + async with database: + await Role.objects.create( + name="user", order=0, description="no administration right" + ) + role_1 = await Role.objects.create( + name="admin", order=1, description="standard administration right" + ) + await Role.objects.create( + name="super_admin", order=2, description="super administration right" + ) + assert await Role.objects.count() == 3 + + company_0 = await Company.objects.create(name="Company") + company_1 = await Company.objects.create(name="Subsidiary Company 1") + company_2 = await Company.objects.create(name="Subsidiary Company 2") + company_3 = await Company.objects.create(name="Subsidiary Company 3") + assert await Company.objects.count() == 4 + + user = await User.objects.create( + registrationnumber="00-00000", company=company_0, name="admin", role=role_1 + ) + assert await User.objects.count() == 1 + + await user.delete() + assert await User.objects.count() == 0 + + user = await User.objects.create( + registrationnumber="00-00000", + company=company_0, + company2=company_3, + name="admin", + role=role_1, + ) + await user.roleforcompanies.add(company_1) + await user.roleforcompanies.add(company_2) + + users = await User.objects.select_related( + ["company", "company2", "roleforcompanies"] + ).all() + assert len(users) == 1 + assert len(users[0].roleforcompanies) == 2 + assert len(users[0].roleforcompanies[0].role_users) == 1 + assert users[0].company.name == "Company" + assert len(users[0].company.users) == 1 + assert users[0].company2.name == "Subsidiary Company 3" + assert len(users[0].company2.secondary_users) == 1 + + users = await User.objects.select_related("roleforcompanies").all() + assert len(users) == 1 + assert len(users[0].roleforcompanies) == 2