Rework dependencies initialization in migrations

app/alembic/versions/01f3f05a5b11_add_primary_group_id.py

@@ -114,8 +114,9 @@ async def _add_domain_computers_group(connection: AsyncConnection) -> None:  # n
 
     op.run_async(_add_domain_computers_group)
 
-    async def _add_primary_group_id(connection: AsyncConnection) -> None:
-        session = AsyncSession(connection)
+    async def _add_primary_group_id(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         base_dn_list = await get_base_directories(session)
         if not base_dn_list:
@@ -164,15 +165,17 @@ async def _add_primary_group_id(connection: AsyncConnection) -> None:
     op.run_async(_add_primary_group_id)
 
 
-def downgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def downgrade(container: AsyncContainer) -> None:
     """Downgrade."""
     bind = op.get_bind()
     session = Session(bind=bind)
 
     async def _delete_domain_computers_group(
-        connection: AsyncConnection,
+        connection: AsyncConnection,  # noqa: ARG001
     ) -> None:
-        session = AsyncSession(connection)
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+
         base_dn_list = await get_base_directories(session)
         if not base_dn_list:
             return

app/alembic/versions/05ddc0bd562a_add_roles.py

@@ -7,13 +7,11 @@
 
 import sqlalchemy as sa
 from alembic import op
-from dishka import AsyncContainer
+from dishka import AsyncContainer, Scope
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
 from entities import Directory, Group
-from ldap_protocol.roles.ace_dao import AccessControlEntryDAO
-from ldap_protocol.roles.role_dao import RoleDAO
 from ldap_protocol.roles.role_use_case import RoleUseCase
 from ldap_protocol.utils.queries import get_base_directories
 from repo.pg.tables import queryable_attr as qa
@@ -25,7 +23,7 @@
 depends_on: None = None
 
 
-def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
     op.create_table(
         "Roles",
@@ -153,17 +151,15 @@ def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
     op.drop_table("AccessPolicyMemberships")
     op.drop_table("AccessPolicies")
 
-    async def _create_system_roles(connection: AsyncConnection) -> None:
-        session = AsyncSession(connection)
-        await session.begin()
+    async def _create_system_roles(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            role_use_case = await cnt.get(RoleUseCase)
 
         base_dn_list = await get_base_directories(session)
         if not base_dn_list:
             return
 
-        role_dao = RoleDAO(session)
-        ace_dao = AccessControlEntryDAO(session)
-        role_use_case = RoleUseCase(role_dao, ace_dao)
         await role_use_case.create_domain_admins_role()
         await role_use_case.create_read_only_role()
 

app/alembic/versions/275222846605_initial_ldap_schema.py

@@ -268,9 +268,9 @@ def upgrade(container: AsyncContainer) -> None:
     session.commit()
 
     # NOTE: Load objectClasses into the database
-    async def _create_object_classes(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _create_object_classes(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         oc_already_created_oids = set()
         oc_first_priority_raw_definitions = (
@@ -343,11 +343,11 @@ async def _create_object_classes(connection: AsyncConnection) -> None:
 
     op.run_async(_create_object_classes)
 
-    async def _create_attribute_types(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _create_attribute_types(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            attribute_type_dao = await cnt.get(AttributeTypeDAO)
 
-        attribute_type_dao = AttributeTypeDAO(session)
         for oid, name in (
             ("2.16.840.1.113730.3.1.610", "nsAccountLock"),
             ("1.3.6.1.4.1.99999.1.1", "posixEmail"),
@@ -368,11 +368,9 @@ async def _create_attribute_types(connection: AsyncConnection) -> None:
 
     op.run_async(_create_attribute_types)
 
-    async def _modify_object_classes(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
-
+    async def _modify_object_classes(connection: AsyncConnection) -> None:  # noqa: ARG001
         async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
             at_dao = await cnt.get(AttributeTypeDAO)
             oc_dao = await cnt.get(ObjectClassDAO)
 

app/alembic/versions/4e8772277cfe_add_web_permissions.py

@@ -7,7 +7,7 @@
 """
 
 from alembic import op
-from dishka import AsyncContainer
+from dishka import AsyncContainer, Scope
 from sqlalchemy import Column, select, text
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
@@ -22,12 +22,12 @@
 depends_on: None | list[str] = None
 
 
-def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
 
-    async def _add_api_permissions(connection: AsyncConnection) -> None:
-        session = AsyncSession(connection)
-        await session.begin()
+    async def _add_api_permissions(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
         query = (
             select(Role)
             .filter_by(name=RoleConstants.DOMAIN_ADMINS_ROLE_NAME)

app/alembic/versions/6303f5c706ec_update_krbadmin_useraccountcontrol_.py

@@ -7,7 +7,7 @@
 """
 
 from alembic import op
-from dishka import AsyncContainer
+from dishka import AsyncContainer, Scope
 from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 from sqlalchemy.orm import joinedload
@@ -24,12 +24,12 @@
 depends_on: None | list[str] = None
 
 
-def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
 
-    async def _update_krbadmin_uac(connection: AsyncConnection) -> None:
-        session = AsyncSession(connection)
-        await session.begin()
+    async def _update_krbadmin_uac(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         krbadmin_user_dir = await session.scalar(
             select(Directory)
@@ -52,9 +52,9 @@ async def _update_krbadmin_uac(connection: AsyncConnection) -> None:
                 ),
             )
 
-    async def _change_uid_admin(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _change_uid_admin(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         directory = await session.scalar(
             select(Directory)
@@ -90,12 +90,12 @@ async def _change_uid_admin(connection: AsyncConnection) -> None:
     op.run_async(_change_uid_admin)
 
 
-def downgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def downgrade(container: AsyncContainer) -> None:
     """Downgrade."""
 
-    async def _downgrade_krbadmin_uac(connection: AsyncConnection) -> None:
-        session = AsyncSession(connection)
-        await session.begin()
+    async def _downgrade_krbadmin_uac(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         krbadmin_user_dir = await session.scalar(
             select(Directory)

app/alembic/versions/8164b4a9e1f1_add_ou_computers.py

@@ -35,11 +35,9 @@ def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
     from ldap_protocol.auth.setup_gateway import SetupGateway
 
-    async def _create_ou_computers(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
-
+    async def _create_ou_computers(connection: AsyncConnection) -> None:  # noqa: ARG001
         async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
             setup_gateway = await cnt.get(SetupGateway)
             role_use_case = await cnt.get(RoleUseCase)
 
@@ -80,12 +78,12 @@ async def _create_ou_computers(connection: AsyncConnection) -> None:
     op.run_async(_create_ou_computers)
 
 
-def downgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def downgrade(container: AsyncContainer) -> None:
     """Downgrade."""
 
-    async def _delete_ou_computers(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _delete_ou_computers(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         base_dn_list = await get_base_directories(session)
         if not base_dn_list:

app/alembic/versions/ba78cef9700a_initial_entity_type.py

@@ -96,16 +96,14 @@ def upgrade(container: AsyncContainer) -> None:
         ["oid"],
     )
 
-    async def _create_entity_types(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _create_entity_types(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            entity_type_use_case = await cnt.get(EntityTypeUseCase)
 
         if not await get_base_directories(session):
             return
 
-        async with container(scope=Scope.REQUEST) as cnt:
-            entity_type_use_case = await cnt.get(EntityTypeUseCase)
-
         for entity_type_data in ENTITY_TYPE_DATAS:
             await entity_type_use_case.create(
                 EntityTypeDTO(
@@ -118,10 +116,10 @@ async def _create_entity_types(connection: AsyncConnection) -> None:
         await session.commit()
 
     async def _append_object_class_to_user_dirs(
-        connection: AsyncConnection,
+        connection: AsyncConnection,  # noqa: ARG001
     ) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         if not await get_base_directories(session):
             return
@@ -156,17 +154,15 @@ async def _append_object_class_to_user_dirs(
         await session.commit()
 
     async def _attach_entity_type_to_directories(
-        connection: AsyncConnection,
+        connection: AsyncConnection,  # noqa: ARG001
     ) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            entity_type_dao = await cnt.get(EntityTypeDAO)
 
         if not await get_base_directories(session):
             return
 
-        async with container(scope=Scope.REQUEST) as cnt:
-            entity_type_dao = await cnt.get(EntityTypeDAO)
-
         await entity_type_dao.attach_entity_type_to_directories()
 
         await session.commit()

app/alembic/versions/c4888c68e221_fix_admin_attr_and_policy.py

@@ -8,16 +8,12 @@
 
 import sqlalchemy as sa
 from alembic import op
-from dishka import AsyncContainer
+from dishka import AsyncContainer, Scope
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 from sqlalchemy.orm import joinedload
 
 from entities import Attribute, Directory, NetworkPolicy
-from ldap_protocol.ldap_schema.attribute_value_validator import (
-    AttributeValueValidator,
-)
 from ldap_protocol.ldap_schema.entity_type_dao import EntityTypeDAO
-from ldap_protocol.ldap_schema.object_class_dao import ObjectClassDAO
 from ldap_protocol.utils.helpers import create_integer_hash
 from ldap_protocol.utils.queries import get_base_directories
 from repo.pg.tables import queryable_attr as qa
@@ -29,32 +25,25 @@
 depends_on: None | list[str] = None
 
 
-def upgrade(container: AsyncContainer) -> None:  # noqa: ARG001
+def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
 
     async def _attach_entity_type_to_directories(
-        connection: AsyncConnection,
+        connection: AsyncConnection,  # noqa: ARG001
     ) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            entity_type_dao = await cnt.get(EntityTypeDAO)
 
         if not await get_base_directories(session):
             return
 
-        object_class_dao = ObjectClassDAO(
-            session,
-        )
-        entity_type_dao = EntityTypeDAO(
-            session,
-            object_class_dao=object_class_dao,
-            attribute_value_validator=AttributeValueValidator(),
-        )
         await entity_type_dao.attach_entity_type_to_directories()
         await session.commit()
 
-    async def _change_uid_admin(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _change_uid_admin(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         directory = await session.scalar(
             sa.select(Directory)
@@ -86,9 +75,9 @@ async def _change_uid_admin(connection: AsyncConnection) -> None:
         )
         await session.commit()
 
-    async def _change_ldap_session_ttl(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
+    async def _change_ldap_session_ttl(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
 
         await session.execute(
             sa.update(NetworkPolicy)

app/alembic/versions/df4c52a613e5_migrate_password_prop_from_E.py

@@ -46,11 +46,9 @@ def upgrade(container: AsyncContainer) -> None:
         ),
     )
 
-    async def _create_common_passwords(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
-        await session.begin()
-
+    async def _create_common_passwords(connection: AsyncConnection) -> None:  # noqa: ARG001
         async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
             password_ban_word_repo = await cnt.get(PasswordBanWordRepository)
 
         await password_ban_word_repo.replace(_BAN_WORDS)

app/alembic/versions/e4d6d99d32bd_add_audit_policies.py

@@ -30,15 +30,15 @@
 def upgrade(container: AsyncContainer) -> None:
     """Upgrade."""
 
-    async def _create_audit_policies(connection: AsyncConnection) -> None:
-        session = AsyncSession(bind=connection)
+    async def _create_audit_policies(connection: AsyncConnection) -> None:  # noqa: ARG001
+        async with container(scope=Scope.REQUEST) as cnt:
+            session = await cnt.get(AsyncSession)
+            audit_dao = await cnt.get(AuditPoliciesDAO)
+            dest_dao = await cnt.get(AuditDestinationDAO)
 
         if not await get_base_directories(session):
             return
 
-        async with container(scope=Scope.REQUEST) as cnt:
-            audit_dao = await cnt.get(AuditPoliciesDAO)
-            dest_dao = await cnt.get(AuditDestinationDAO)
         manager = Mock(spec=RawAuditManager)
         use_case = AuditUseCase(audit_dao, dest_dao, manager)
         await use_case.create_policies()