amsdal 0.3.5__cp311-cp311-macosx_10_9_universal2.whl → 0.4.0__cp311-cp311-macosx_10_9_universal2.whl

amsdal/utils/rollback/__init__.py

@@ -4,7 +4,7 @@ from amsdal_data.application import DataApplication
 from amsdal_data.transactions.decorators import async_transaction
 from amsdal_data.transactions.decorators import transaction
 from amsdal_data.transactions.errors import AmsdalTransactionError
-from amsdal_models.classes.manager import ClassManager
+from amsdal_models.classes.class_manager import ClassManager
 from amsdal_models.querysets.executor import LAKEHOUSE_DB_ALIAS
 
 
@@ -28,14 +28,21 @@ def rollback_to_timestamp(timestamp: float) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
                     lookup=glue.FieldLookup.GT,
-                    value=glue.Value(timestamp),
+                    right=glue.Value(timestamp),
                 ),
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='prior_version'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
                     lookup=glue.FieldLookup.ISNULL,
-                    value=glue.Value(True),
+                    right=glue.Value(True),
                 ),
             ),
         )
@@ -48,14 +55,21 @@ def rollback_to_timestamp(timestamp: float) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
                     lookup=glue.FieldLookup.GT,
-                    value=glue.Value(timestamp),
+                    right=glue.Value(timestamp),
                 ),
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='prior_version'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
                     lookup=glue.FieldLookup.ISNULL,
-                    value=glue.Value(False),
+                    right=glue.Value(False),
                 ),
             ),
         )
@@ -83,9 +97,12 @@ def rollback_to_timestamp(timestamp: float) -> None:
             _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
             _conditions.append(
                 glue.Condition(
-                    field=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
                     lookup=glue.FieldLookup.EQ,
-                    value=glue.Value(transaction_id),
+                    right=glue.Value(transaction_id, output_type=str),
                 )
             )
 
@@ -94,9 +111,14 @@ def rollback_to_timestamp(timestamp: float) -> None:
                 table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
                 where=glue.Conditions(
                     glue.Condition(
-                        field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                        left=glue.FieldReferenceExpression(
+                            field_reference=glue.FieldReference(
+                                field=glue.Field(name='updated_at'),
+                                table_name='Metadata',
+                            ),
+                        ),
                         lookup=glue.FieldLookup.LTE,
-                        value=glue.Value(timestamp),
+                        right=glue.Value(timestamp),
                     ),
                     glue.Conditions(*_conditions, connector=glue.FilterConnector.OR),
                 ),
@@ -108,9 +130,7 @@ def rollback_to_timestamp(timestamp: float) -> None:
 
     for m in metadatas_to_delete:
         class_name = m.data['class_schema_reference']['ref']['object_id']
-        schema_type = class_manager.resolve_schema_type(class_name)
-
-        model_class = class_manager.import_model_class(class_name, schema_type)
+        model_class = class_manager.import_class(class_name)
         obj = (
             model_class.objects.filter(_address__object_id=m.data['object_id'])
             .using(LAKEHOUSE_DB_ALIAS)
@@ -123,8 +143,7 @@ def rollback_to_timestamp(timestamp: float) -> None:
             obj.delete()
 
     for object_id, class_name in ids_to_revert:
-        schema_type = class_manager.resolve_schema_type(class_name)
-        model_class = class_manager.import_model_class(class_name, schema_type)
+        model_class = class_manager.import_class(class_name)
 
         obj = (
             model_class.objects.filter(_address__object_id=object_id)
@@ -133,13 +152,14 @@ def rollback_to_timestamp(timestamp: float) -> None:
             .first()
             .execute()
         )
-        old_obj = (
-            model_class.objects.filter(_address__object_id=object_id, _metadata__updated_at__lte=timestamp)
-            .using(LAKEHOUSE_DB_ALIAS)
-            .order_by('-_metadata__updated_at')
-            .first()
-            .execute()
-        )
+        old_obj = obj.previous_version()  # type: ignore[union-attr]
+        # old_obj = (
+        #     model_class.objects.filter(_address__object_id=object_id, _metadata__updated_at__lte=timestamp)
+        #     .using(LAKEHOUSE_DB_ALIAS)
+        #     .order_by('-_metadata__updated_at')
+        #     .first()
+        #     .execute()
+        # )
 
         if obj and old_obj:
             for field, value in old_obj.model_dump().items():
@@ -180,9 +200,12 @@ def rollback_transaction(transaction_id: str) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
                     lookup=glue.FieldLookup.EQ,
-                    value=glue.Value(transaction_id),
+                    right=glue.Value(transaction_id, output_type=str),
                 )
             ),
             order_by=[
@@ -222,14 +245,21 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
                     lookup=glue.FieldLookup.GT,
-                    value=glue.Value(timestamp),
+                    right=glue.Value(timestamp),
                 ),
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='prior_version'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
                     lookup=glue.FieldLookup.ISNULL,
-                    value=glue.Value(True),
+                    right=glue.Value(True),
                 ),
             ),
         )
@@ -242,14 +272,20 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
                     lookup=glue.FieldLookup.GT,
-                    value=glue.Value(timestamp),
+                    right=glue.Value(timestamp),
                 ),
                 glue.Condition(
-                    field=glue.FieldReference(field=glue.Field(name='prior_version'), table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'), table_name='Metadata'
+                        ),
+                    ),
                     lookup=glue.FieldLookup.ISNULL,
-                    value=glue.Value(False),
+                    right=glue.Value(False),
                 ),
             ),
         )
@@ -277,9 +313,12 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
             _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
             _conditions.append(
                 glue.Condition(
-                    field=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
                     lookup=glue.FieldLookup.EQ,
-                    value=glue.Value(transaction_id),
+                    right=glue.Value(transaction_id, output_type=str),
                 )
             )
 
@@ -288,9 +327,14 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
                 table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
                 where=glue.Conditions(
                     glue.Condition(
-                        field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                        left=glue.FieldReferenceExpression(
+                            field_reference=glue.FieldReference(
+                                field=glue.Field(name='updated_at'),
+                                table_name='Metadata',
+                            ),
+                        ),
                         lookup=glue.FieldLookup.LTE,
-                        value=glue.Value(timestamp),
+                        right=glue.Value(timestamp),
                     ),
                     glue.Conditions(*_conditions, connector=glue.FilterConnector.OR),
                 ),
@@ -302,9 +346,7 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
 
     for m in metadatas_to_delete:
         class_name = m.data['class_schema_reference']['ref']['object_id']
-        schema_type = class_manager.resolve_schema_type(class_name)
-
-        model_class = class_manager.import_model_class(class_name, schema_type)
+        model_class = class_manager.import_class(class_name)
         obj = await (
             model_class.objects.filter(_address__object_id=m.data['object_id'])
             .using(LAKEHOUSE_DB_ALIAS)
@@ -317,8 +359,7 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
             await obj.adelete()
 
     for object_id, class_name in ids_to_revert:
-        schema_type = class_manager.resolve_schema_type(class_name)
-        model_class = class_manager.import_model_class(class_name, schema_type)
+        model_class = class_manager.import_class(class_name)
 
         obj = await (
             model_class.objects.filter(_address__object_id=object_id)
@@ -327,16 +368,17 @@ async def async_rollback_to_timestamp(timestamp: float) -> None:
             .first()
             .aexecute()
         )
-        old_obj = await (
-            model_class.objects.filter(_address__object_id=object_id, _metadata__updated_at__lte=timestamp)
-            .using(LAKEHOUSE_DB_ALIAS)
-            .order_by('-_metadata__updated_at')
-            .first()
-            .aexecute()
-        )
+        old_obj = await obj.aprevious_version()  # type: ignore[union-attr]
+        # old_obj = await (
+        #     model_class.objects.filter(_address__object_id=object_id, _metadata__updated_at__lte=timestamp)
+        #     .using(LAKEHOUSE_DB_ALIAS)
+        #     .order_by('-_metadata__updated_at')
+        #     .first()
+        #     .aexecute()
+        # )
 
         if obj and old_obj:
-            for field, value in old_obj.model_dump().items():
+            for field, value in (await old_obj.amodel_dump()).items():
                 setattr(obj, field, value)
 
             await obj.asave()
@@ -373,9 +415,12 @@ async def async_rollback_transaction(transaction_id: str) -> None:
             table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
             where=glue.Conditions(
                 glue.Condition(
-                    field=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
                     lookup=glue.FieldLookup.EQ,
-                    value=glue.Value(transaction_id),
+                    right=glue.Value(transaction_id, output_type=str),
                 )
             ),
             order_by=[

amsdal/utils/tests/helpers.py

@@ -2,31 +2,36 @@ import os
 import random
 import string
 import tempfile
+from collections.abc import AsyncGenerator
+from collections.abc import AsyncIterator
 from collections.abc import Generator
 from collections.abc import Iterator
+from contextlib import asynccontextmanager
 from contextlib import contextmanager
 from contextlib import suppress
 from pathlib import Path
 from typing import Any
 
 import yaml
+from amsdal_data.application import AsyncDataApplication
+from amsdal_models.migration import MigrationsLoader
+from amsdal_models.migration import migrations
+from amsdal_models.migration.executors.default_executor import DefaultAsyncMigrationExecutor
+from amsdal_models.migration.executors.default_executor import DefaultMigrationExecutor
+from amsdal_models.migration.file_migration_executor import FileMigrationExecutorManager
+from amsdal_models.migration.file_migration_generator import FileMigrationGenerator
+from amsdal_models.migration.file_migration_writer import FileMigrationWriter
+from amsdal_models.migration.migrations import MigrateData
+from amsdal_models.migration.migrations import MigrationSchemas
+from amsdal_models.migration.schemas_loaders import JsonClassSchemaLoader  # type: ignore
+from amsdal_models.migration.utils import contrib_to_module_root_path
 from amsdal_utils.config.manager import AmsdalConfigManager
-from amsdal_utils.models.enums import SchemaTypes
+from amsdal_utils.models.enums import ModuleType
 
 from amsdal.configs.constants import CORE_MIGRATIONS_PATH
 from amsdal.configs.main import settings
 from amsdal.manager import AmsdalManager
-from amsdal.migration import migrations
-from amsdal.migration.data_classes import ModuleTypes
-from amsdal.migration.executors.default_executor import DefaultMigrationExecutor
-from amsdal.migration.file_migration_executor import FileMigrationExecutorManager
-from amsdal.migration.file_migration_generator import FileMigrationGenerator
-from amsdal.migration.file_migration_writer import FileMigrationWriter
-from amsdal.migration.migrations import MigrateData
-from amsdal.migration.migrations import MigrationSchemas
-from amsdal.migration.migrations_loader import MigrationsLoader
-from amsdal.migration.schemas_loaders import JsonClassSchemaLoader
-from amsdal.migration.utils import contrib_to_module_root_path
+from amsdal.manager import AsyncAmsdalManager
 from amsdal.utils.tests.enums import DbExecutionType
 from amsdal.utils.tests.enums import LakehouseOption
 from amsdal.utils.tests.enums import StateOption
@@ -89,19 +94,32 @@ def _get_config_template(
     db_execution_type: DbExecutionType,
     lakehouse_option: LakehouseOption,
     state_option: StateOption,
+    *,
+    is_async: bool = False,
 ) -> str:
     config_object: dict[str, Any] = {
         'application_name': 'test_client_app',
+        'async_mode': is_async,
         'connections': [
-            {'name': 'lock', 'backend': 'amsdal_data.lock.implementations.thread_lock.ThreadLock'},
+            {
+                'name': 'lock',
+                'backend': 'amsdal_data.lock.implementations.thread_lock.ThreadLock',
+            },
         ],
-        'resources_config': {'lakehouse': 'lakehouse', 'lock': 'lock', 'repository': {'default': 'state'}},
+        'resources_config': {
+            'lakehouse': 'lakehouse',
+            'lock': 'lock',
+            'repository': {'default': 'state'},
+        },
     }
-    if lakehouse_option in [LakehouseOption.postgres, LakehouseOption.postgres_immutable]:
+    if lakehouse_option in [
+        LakehouseOption.postgres,
+        LakehouseOption.postgres_immutable,
+    ]:
         config_object['connections'].append(
             {
                 'name': 'lakehouse',
-                'backend': 'postgres-historical',
+                'backend': 'postgres-historical-async' if is_async else 'postgres-historical',
                 'credentials': [
                     {
                         'db_host': '{{ db_host }}',
@@ -117,7 +135,7 @@ def _get_config_template(
         config_object['connections'].append(
             {
                 'name': 'lakehouse',
-                'backend': 'sqlite-historical',
+                'backend': 'sqlite-historical-async' if is_async else 'sqlite-historical',
                 'credentials': [{'db_path': '{{ db_dir }}/sqlite_lakehouse.sqlite3'}],
             }
         )
@@ -131,7 +149,7 @@ def _get_config_template(
         config_object['connections'].append(
             {
                 'name': 'state',
-                'backend': 'postgres',
+                'backend': 'postgres-async' if is_async else 'postgres',
                 'credentials': [
                     {
                         'db_host': '{{ db_host }}',
@@ -148,7 +166,7 @@ def _get_config_template(
         config_object['connections'].append(
             {
                 'name': 'state',
-                'backend': 'sqlite',
+                'backend': 'sqlite-async' if is_async else 'sqlite',
                 'credentials': [{'db_path': '{{ db_dir }}/sqlite_state.sqlite3'}],
             }
         )
@@ -157,11 +175,13 @@ def _get_config_template(
 
 
 @contextmanager
-def init_manager(
+def _init_manager(
     db_execution_type: DbExecutionType,
     lakehouse_option: LakehouseOption,
     state_option: StateOption,
-) -> Generator[AmsdalManager, Any, None]:
+    *,
+    is_async: bool = False,
+) -> Generator[tuple[Path, Path], Any, None]:
     Path('.tmp').mkdir(exist_ok=True)
     with tempfile.TemporaryDirectory(dir='.tmp') as temp_dir:
         db_dir = Path(temp_dir) / 'db_dir'
@@ -169,9 +189,12 @@ def init_manager(
 
         lakehouse_database = ''.join(random.sample(string.ascii_letters, 16))
         state_database = ''.join(random.sample(string.ascii_letters, 16))
-        config_text = _get_config_template(db_execution_type, lakehouse_option, state_option)
+        config_text = _get_config_template(db_execution_type, lakehouse_option, state_option, is_async=is_async)
 
-        if lakehouse_option in [LakehouseOption.postgres, LakehouseOption.postgres_immutable]:
+        if lakehouse_option in [
+            LakehouseOption.postgres,
+            LakehouseOption.postgres_immutable,
+        ]:
             (
                 db_host,
                 db_port,
@@ -186,7 +209,10 @@ def init_manager(
                 .replace('{{ db_password }}', db_password)
                 .replace('{{ lakehouse_postgres_db }}', lakehouse_database)
             )
-        elif lakehouse_option in [LakehouseOption.sqlite, LakehouseOption.sqlite_immutable]:
+        elif lakehouse_option in [
+            LakehouseOption.sqlite,
+            LakehouseOption.sqlite_immutable,
+        ]:
             config_text = config_text.replace('{{ db_dir }}', db_dir.absolute().as_posix())
 
         if state_option == StateOption.postgres:
@@ -204,6 +230,19 @@ def init_manager(
         config_path = Path(temp_dir) / 'config.yml'
         config_path.write_text(config_text)
 
+        yield db_dir, config_path
+
+
+@contextmanager
+def init_manager(
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+) -> Generator[AmsdalManager, Any, None]:
+    with _init_manager(db_execution_type, lakehouse_option, state_option) as (
+        db_dir,
+        config_path,
+    ):
         with override_settings(APP_PATH=db_dir, CONFIG_PATH=config_path):
             config_manager = AmsdalConfigManager()
             config_manager.load_config(config_path)
@@ -219,6 +258,83 @@ def init_manager(
                 AmsdalConfigManager.invalidate()
 
 
+@asynccontextmanager
+async def async_init_manager(
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+) -> AsyncIterator[AsyncAmsdalManager]:
+    with _init_manager(db_execution_type, lakehouse_option, state_option, is_async=True) as (db_dir, config_path):
+        with override_settings(
+            APP_PATH=db_dir,
+            CONFIG_PATH=config_path,
+        ):
+            config_manager = AmsdalConfigManager()
+            config_manager.load_config(config_path)
+            manager = AsyncAmsdalManager()
+            await manager.setup()
+            await manager.post_setup()  # type: ignore[call-arg,misc]
+
+            try:
+                yield manager
+            finally:
+                await manager.teardown()
+                await AsyncDataApplication().teardown()
+                AsyncAmsdalManager.invalidate()
+                AmsdalConfigManager.invalidate()
+                AsyncDataApplication.invalidate()
+
+
+def migrate(app_schemas_path: Path) -> None:
+    schemas = MigrationSchemas()
+    executor = DefaultMigrationExecutor(schemas)
+    _migrate_per_loader(
+        executor,
+        MigrationsLoader(
+            migrations_dir=CORE_MIGRATIONS_PATH,
+            module_type=ModuleType.CORE,
+        ),
+    )
+
+    for contrib in settings.CONTRIBS:
+        contrib_root_path = contrib_to_module_root_path(contrib)
+        _migrate_per_loader(
+            executor,
+            MigrationsLoader(
+                migrations_dir=contrib_root_path / settings.MIGRATIONS_DIRECTORY_NAME,
+                module_type=ModuleType.CONTRIB,
+                module_name=contrib,
+            ),
+        )
+
+    schema_loader = JsonClassSchemaLoader(app_schemas_path)
+
+    for class_schema in schema_loader.iter_app_schemas():
+        for _operation_data in FileMigrationGenerator.build_operations(
+            ModuleType.USER,
+            class_schema.object_schema,
+            None,
+        ):
+            _operation_name = FileMigrationWriter.operation_name_map[_operation_data.type]
+            kwargs = {}
+
+            if _operation_data.old_schema:
+                kwargs['old_schema'] = _operation_data.old_schema.model_dump()
+
+            if _operation_data.new_schema:
+                kwargs['new_schema'] = _operation_data.new_schema.model_dump()
+
+            _operation = getattr(migrations, _operation_name)(
+                schema_type=ModuleType.USER,
+                class_name=_operation_data.class_name,
+                **kwargs,
+            )
+
+            _operation.forward(executor)
+
+    executor.flush_buffer()
+
+
 def _migrate_per_loader(executor: DefaultMigrationExecutor, loader: MigrationsLoader) -> None:
     for _migration in loader:
         migration_class = FileMigrationExecutorManager.get_migration_class(_migration)
@@ -233,24 +349,24 @@ def _migrate_per_loader(executor: DefaultMigrationExecutor, loader: MigrationsLo
         executor.flush_buffer()
 
 
-def migrate(app_schemas_path: Path) -> None:
+async def async_migrate(app_schemas_path: Path) -> None:
     schemas = MigrationSchemas()
-    executor = DefaultMigrationExecutor(schemas)
-    _migrate_per_loader(
+    executor = DefaultAsyncMigrationExecutor(schemas)
+    await _async_migrate_per_loader(
         executor,
         MigrationsLoader(
             migrations_dir=CORE_MIGRATIONS_PATH,
-            module_type=ModuleTypes.CORE,
+            module_type=ModuleType.CORE,
         ),
     )
 
     for contrib in settings.CONTRIBS:
         contrib_root_path = contrib_to_module_root_path(contrib)
-        _migrate_per_loader(
+        await _async_migrate_per_loader(
             executor,
             MigrationsLoader(
                 migrations_dir=contrib_root_path / settings.MIGRATIONS_DIRECTORY_NAME,
-                module_type=ModuleTypes.CONTRIB,
+                module_type=ModuleType.CONTRIB,
                 module_name=contrib,
             ),
         )
@@ -259,7 +375,7 @@ def migrate(app_schemas_path: Path) -> None:
 
     for class_schema in schema_loader.iter_app_schemas():
         for _operation_data in FileMigrationGenerator.build_operations(
-            SchemaTypes.USER,
+            ModuleType.USER,
             class_schema.object_schema,
             None,
         ):
@@ -273,14 +389,28 @@ def migrate(app_schemas_path: Path) -> None:
                 kwargs['new_schema'] = _operation_data.new_schema.model_dump()
 
             _operation = getattr(migrations, _operation_name)(
-                schema_type=SchemaTypes.USER,
+                schema_type=ModuleType.USER,
                 class_name=_operation_data.class_name,
                 **kwargs,
             )
 
             _operation.forward(executor)
 
-    executor.flush_buffer()
+    await executor.flush_buffer()
+
+
+async def _async_migrate_per_loader(executor: DefaultAsyncMigrationExecutor, loader: MigrationsLoader) -> None:
+    for _migration in loader:
+        migration_class = FileMigrationExecutorManager.get_migration_class(_migration)
+        migration_class_instance = migration_class()
+
+        for _operation in migration_class_instance.operations:
+            if isinstance(_operation, MigrateData):
+                await executor.flush_buffer()
+
+            _operation.forward(executor)
+
+        await executor.flush_buffer()
 
 
 @contextmanager
@@ -294,16 +424,35 @@ def init_manager_and_migrate(
         lakehouse_option=lakehouse_option,
         state_option=state_option,
     ) as manager:
-        app_schemas_path = TESTS_DIR / 'src' / 'models'
+        manager.authenticate()
+        migrate()
+        manager.init_classes()
+
+        yield manager
+
+
+@asynccontextmanager
+async def async_init_manager_and_migrate(
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+) -> AsyncGenerator[AsyncAmsdalManager, Any]:
+    async with async_init_manager(
+        db_execution_type=db_execution_type,
+        lakehouse_option=lakehouse_option,
+        state_option=state_option,
+    ) as manager:
+        src_dir_path = TESTS_DIR / 'src'
+        app_schemas_path = src_dir_path / 'models'
         manager.build(
             source_models_path=app_schemas_path,
             source_fixtures_path=app_schemas_path,
-            source_transactions_path=app_schemas_path,
+            source_transactions_path=src_dir_path,
             source_static_files_path=app_schemas_path,
             source_migrations_path=app_schemas_path,
         )
         manager.authenticate()
-        migrate(app_schemas_path)
+        await async_migrate(app_schemas_path)
         manager.init_classes()
 
         yield manager