Fix migration 275222846605

app/alembic/versions/275222846605_initial_ldap_schema.py

@@ -10,17 +10,16 @@
 
 import sqlalchemy as sa
 from alembic import op
-from dishka import AsyncContainer, Scope
+from dishka import AsyncContainer
 from ldap3.protocol.schemas.ad2012R2 import ad_2012_r2_schema
-from sqlalchemy import delete, or_
+from sqlalchemy import delete, or_, select
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, selectinload
 
-from entities import Attribute
+from entities import Attribute, AttributeType, ObjectClass
 from extra.alembic_utils import temporary_stub_entity_type_name
 from ldap_protocol.ldap_schema.attribute_type_dao import AttributeTypeDAO
 from ldap_protocol.ldap_schema.dto import AttributeTypeDTO
-from ldap_protocol.ldap_schema.object_class_dao import ObjectClassDAO
 from ldap_protocol.utils.raw_definition_parser import (
     RawDefinitionParser as RDParser,
 )
@@ -37,7 +36,7 @@
 
 
 @temporary_stub_entity_type_name
-def upgrade(container: AsyncContainer) -> None:
+def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
     """Upgrade."""
     bind = op.get_bind()
     session = Session(bind=bind)
@@ -370,21 +369,29 @@ async def _create_attribute_types(connection: AsyncConnection) -> None:
 
     async def _modify_object_classes(connection: AsyncConnection) -> None:
         session = AsyncSession(bind=connection)
-        await session.begin()
-
-        async with container(scope=Scope.REQUEST) as cnt:
-            at_dao = await cnt.get(AttributeTypeDAO)
-            oc_dao = await cnt.get(ObjectClassDAO)
 
         for oc_name, at_names in (
             ("user", ["nsAccountLock", "shadowExpire"]),
             ("computer", ["userAccountControl"]),
             ("posixAccount", ["posixEmail"]),
             ("organizationalUnit", ["title", "jpegPhoto"]),
         ):
-            object_class = await oc_dao.get(oc_name)
-            attribute_types_may = await at_dao.get_all_by_names(at_names)
-            object_class.attribute_types_may.extend(attribute_types_may)
+            object_class = await session.scalar(
+                select(ObjectClass)
+                .filter_by(name=oc_name)
+                .options(selectinload(qa(ObjectClass.attribute_types_may))),
+            )
+
+            if not object_class:
+                continue
+
+            attribute_types = await session.scalars(
+                select(AttributeType).where(
+                    qa(AttributeType.name).in_(at_names),
+                ),
+            )
+
+            object_class.attribute_types_may.extend(attribute_types.all())
 
         await session.commit()