amsdal 0.3.3__cp311-cp311-macosx_10_9_universal2.whl → 0.5.29__cp311-cp311-macosx_10_9_universal2.whl

amsdal/utils/tests/helpers.py

@@ -1,68 +1,55 @@
-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 ExitStack
+from contextlib import asynccontextmanager
 from contextlib import contextmanager
-from contextlib import suppress
+from functools import partial
 from pathlib import Path
 from typing import Any
 
-import yaml
+from amsdal_data.aliases.db import POSTGRES_HISTORICAL_ALIAS
+from amsdal_data.aliases.db import POSTGRES_HISTORICAL_ASYNC_ALIAS
+from amsdal_data.aliases.db import POSTGRES_STATE_ALIAS
+from amsdal_data.aliases.db import POSTGRES_STATE_ASYNC_ALIAS
+from amsdal_data.aliases.db import SQLITE_ALIAS
+from amsdal_data.aliases.db import SQLITE_ASYNC_ALIAS
+from amsdal_data.aliases.db import SQLITE_HISTORICAL_ALIAS
+from amsdal_data.aliases.db import SQLITE_HISTORICAL_ASYNC_ALIAS
+from amsdal_data.application import AsyncDataApplication
+from amsdal_data.connections.db_alias_map import CONNECTION_BACKEND_ALIASES
+from amsdal_data.connections.historical.schema_version_manager import AsyncHistoricalSchemaVersionManager
+from amsdal_data.connections.historical.schema_version_manager import HistoricalSchemaVersionManager
+from amsdal_data.test_utils.common import temp_dir
+from amsdal_data.test_utils.config import build_config
+from amsdal_data.test_utils.config import postgres_async_config
+from amsdal_data.test_utils.config import postgres_async_lakehouse_only_config
+from amsdal_data.test_utils.config import postgres_config
+from amsdal_data.test_utils.config import postgres_lakehouse_only_config
+from amsdal_data.test_utils.config import sqlite_async_config
+from amsdal_data.test_utils.config import sqlite_async_lakehouse_only_config
+from amsdal_data.test_utils.config import sqlite_config
+from amsdal_data.test_utils.config import sqlite_lakehouse_only_config
+from amsdal_data.test_utils.constants import PG_TEST_HOST
+from amsdal_data.test_utils.constants import PG_TEST_PASSWORD
+from amsdal_data.test_utils.constants import PG_TEST_PORT
+from amsdal_data.test_utils.constants import PG_TEST_USER
+from amsdal_data.test_utils.db import create_postgres_database
+from amsdal_data.test_utils.db import create_postgres_extension
+from amsdal_data.test_utils.db import drop_postgres_database
+from amsdal_utils.config.data_models.amsdal_config import AmsdalConfig
 from amsdal_utils.config.manager import AmsdalConfigManager
-from amsdal_utils.models.enums import SchemaTypes
 
-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
-
-TESTS_DIR = Path(os.getcwd())
-
-
-def create_postgres_database(database: str) -> tuple[str, str, str, str]:
-    import psycopg
-
-    db_host = os.getenv('POSTGRES_HOST', 'localhost')
-    db_port = os.getenv('POSTGRES_PORT', '5432')
-    db_user = os.getenv('POSTGRES_USER', 'postgres')
-    db_password = os.getenv('POSTGRES_PASSWORD', 'example')
-
-    conn = psycopg.connect(
-        host=db_host,
-        port=db_port,
-        user=db_user,
-        password=db_password,
-        autocommit=True,
-    )
-    cur = conn.cursor()
-
-    with suppress(psycopg.errors.DuplicateDatabase):
-        cur.execute(f'CREATE DATABASE "{database}"')
-
-    cur.close()
-    conn.close()
-
-    return (
-        db_host,
-        db_port,
-        db_user,
-        db_password,
-    )
+from amsdal.utils.tests.migrations import async_migrate
+from amsdal.utils.tests.migrations import migrate
 
 
 @contextmanager
@@ -85,225 +72,260 @@ def override_settings(**kwargs: Any) -> Iterator[None]:
         settings.override(**original_settings)
 
 
-def _get_config_template(
-    db_execution_type: DbExecutionType,
-    lakehouse_option: LakehouseOption,
-    state_option: StateOption,
-) -> str:
-    config_object: dict[str, Any] = {
-        'application_name': 'test_client_app',
-        'connections': [
-            {'name': 'lock', 'backend': 'amsdal_data.lock.implementations.thread_lock.ThreadLock'},
-        ],
-        'resources_config': {'lakehouse': 'lakehouse', 'lock': 'lock', 'repository': {'default': 'state'}},
-    }
-    if lakehouse_option in [LakehouseOption.postgres, LakehouseOption.postgres_immutable]:
-        config_object['connections'].append(
-            {
-                'name': 'lakehouse',
-                'backend': 'postgres-historical',
-                'credentials': [
-                    {
-                        'db_host': '{{ db_host }}',
-                        'db_port': '{{ db_port }}',
-                        'db_user': '{{ db_user }}',
-                        'db_password': '{{ db_password }}',
-                        'db_name': '{{ lakehouse_postgres_db }}',
-                    }
-                ],
-            }
-        )
-    elif lakehouse_option in [LakehouseOption.sqlite, LakehouseOption.sqlite_immutable]:
-        config_object['connections'].append(
-            {
-                'name': 'lakehouse',
-                'backend': 'sqlite-historical',
-                'credentials': [{'db_path': '{{ db_dir }}/sqlite_lakehouse.sqlite3'}],
-            }
-        )
-
-    if db_execution_type == DbExecutionType.lakehouse_only:
-        config_object['resources_config']['repository']['default'] = 'lakehouse'
-
-        return yaml.dump(config_object)
-
-    if state_option == StateOption.postgres:
-        config_object['connections'].append(
-            {
-                'name': 'state',
-                'backend': 'postgres',
-                'credentials': [
-                    {
-                        'db_host': '{{ db_host }}',
-                        'db_port': '{{ db_port }}',
-                        'db_user': '{{ db_user }}',
-                        'db_password': '{{ db_password }}',
-                        'db_name': '{{ state_postgres_db }}',
-                    }
-                ],
-            }
-        )
-
-    elif state_option == StateOption.sqlite:
-        config_object['connections'].append(
-            {
-                'name': 'state',
-                'backend': 'sqlite',
-                'credentials': [{'db_path': '{{ db_dir }}/sqlite_state.sqlite3'}],
-            }
-        )
-
-    return yaml.dump(config_object)
-
-
 @contextmanager
-def init_manager(
+def init_config(
     db_execution_type: DbExecutionType,
     lakehouse_option: LakehouseOption,
     state_option: StateOption,
-) -> Generator[AmsdalManager, Any, None]:
-    Path('.tmp').mkdir(exist_ok=True)
-    with tempfile.TemporaryDirectory(dir='.tmp') as temp_dir:
-        db_dir = Path(temp_dir) / 'db_dir'
-        (db_dir / 'warehouse').mkdir(exist_ok=True, parents=True)
-
+    *,
+    is_async: bool = False,
+    db_name_prefix: str | None = None,
+    db_path: Path | None = None,
+    drop_database: bool = False,
+) -> Iterator[AmsdalConfig]:
+    if db_name_prefix:
+        lakehouse_database = f'{db_name_prefix}_lakehouse'
+        state_database = f'{db_name_prefix}_state'
+    else:
         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)
-
-        if lakehouse_option in [LakehouseOption.postgres, LakehouseOption.postgres_immutable]:
-            (
-                db_host,
-                db_port,
-                db_user,
-                db_password,
-            ) = create_postgres_database(lakehouse_database)
-
-            config_text = (
-                config_text.replace('{{ db_host }}', db_host)
-                .replace('{{ db_port }}', db_port)
-                .replace('{{ db_user }}', db_user)
-                .replace('{{ db_password }}', db_password)
-                .replace('{{ lakehouse_postgres_db }}', lakehouse_database)
-            )
-        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:
-            create_postgres_database(state_database)
-            config_text = (
-                config_text.replace('{{ db_host }}', db_host)
-                .replace('{{ db_port }}', db_port)
-                .replace('{{ db_user }}', db_user)
-                .replace('{{ db_password }}', db_password)
-                .replace('{{ state_postgres_db }}', state_database)
-            )
-        elif state_option == StateOption.sqlite:
-            config_text = config_text.replace('{{ db_dir }}', db_dir.absolute().as_posix())
 
-        config_path = Path(temp_dir) / 'config.yml'
-        config_path.write_text(config_text)
+    if db_execution_type == DbExecutionType.lakehouse_only:
+        ctx_managers = {
+            LakehouseOption.sqlite: (
+                partial(sqlite_async_lakehouse_only_config, db_path)
+                if is_async
+                else partial(sqlite_lakehouse_only_config, db_path)
+            ),
+            LakehouseOption.postgres: (
+                partial(postgres_async_lakehouse_only_config, lakehouse_database, drop_database=drop_database)
+                if is_async
+                else partial(postgres_lakehouse_only_config, lakehouse_database, drop_database=drop_database)
+            ),
+        }
+        with ctx_managers[lakehouse_option]() as _amsdal_config:
+            yield _amsdal_config
+        return
+
+    if lakehouse_option.value == state_option.value:
+        ctx_managers = {
+            LakehouseOption.sqlite: (
+                partial(
+                    sqlite_async_config,
+                    db_path=db_path,
+                )
+                if is_async
+                else partial(sqlite_config, db_path)
+            ),
+            LakehouseOption.postgres: (
+                partial(
+                    postgres_async_config,
+                    lakehouse_database=lakehouse_database,
+                    state_database=state_database,
+                    drop_database=drop_database,
+                )
+                if is_async
+                else partial(
+                    postgres_config,
+                    lakehouse_database=lakehouse_database,
+                    state_database=state_database,
+                    drop_database=drop_database,
+                )
+            ),
+        }
 
-        with override_settings(APP_PATH=db_dir, CONFIG_PATH=config_path):
-            config_manager = AmsdalConfigManager()
-            config_manager.load_config(config_path)
-            manager = AmsdalManager()
-            manager.setup()
-            manager.post_setup()  # type: ignore[call-arg]
+        with ctx_managers[lakehouse_option]() as _amsdal_config:
+            yield _amsdal_config
+        return
 
+    with ExitStack() as stack:
+        db_path = db_path if db_path else stack.enter_context(temp_dir())
+
+        if lakehouse_option == LakehouseOption.postgres:
+            config = build_config(
+                lakehouse_backend=CONNECTION_BACKEND_ALIASES[
+                    POSTGRES_HISTORICAL_ASYNC_ALIAS if is_async else POSTGRES_HISTORICAL_ALIAS
+                ],
+                lakehouse_credentials={
+                    'dsn': f'postgresql://{PG_TEST_USER}:{PG_TEST_PASSWORD}@{PG_TEST_HOST}:{PG_TEST_PORT}/{lakehouse_database}',
+                },
+                state_backend=CONNECTION_BACKEND_ALIASES[SQLITE_ASYNC_ALIAS if is_async else SQLITE_ALIAS],
+                state_credentials={
+                    'db_path': f'{db_path}/amsdal_state.sqlite3',
+                },
+                is_async_mode=is_async,
+            )
+            AmsdalConfigManager().set_config(config)
+            create_postgres_database(lakehouse_database)
+            create_postgres_extension(lakehouse_database, 'vector')
             try:
-                yield manager
+                yield config
             finally:
-                manager.teardown()
-                AmsdalManager.invalidate()
                 AmsdalConfigManager.invalidate()
 
+                if drop_database:
+                    drop_postgres_database(lakehouse_database)
+            return
+
+        config = build_config(
+            lakehouse_backend=CONNECTION_BACKEND_ALIASES[
+                SQLITE_HISTORICAL_ASYNC_ALIAS if is_async else SQLITE_HISTORICAL_ALIAS
+            ],
+            lakehouse_credentials={
+                'db_path': f'{db_path}/amsdal_historical.sqlite3',
+            },
+            state_backend=CONNECTION_BACKEND_ALIASES[POSTGRES_STATE_ASYNC_ALIAS if is_async else POSTGRES_STATE_ALIAS],
+            state_credentials={
+                'dsn': f'postgresql://{PG_TEST_USER}:{PG_TEST_PASSWORD}@{PG_TEST_HOST}:{PG_TEST_PORT}/{state_database}',
+            },
+            is_async_mode=True,
+        )
+        AmsdalConfigManager().set_config(config)
+        create_postgres_database(state_database)
+        create_postgres_extension(state_database, 'vector')
 
-def _migrate_per_loader(executor: DefaultMigrationExecutor, 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):
-                executor.flush_buffer()
-
-            _operation.forward(executor)
-
-        executor.flush_buffer()
+        try:
+            yield config
+        finally:
+            AmsdalConfigManager.invalidate()
 
+            if drop_database:
+                drop_postgres_database(state_database)
 
-def migrate(app_schemas_path: Path) -> None:
-    schemas = MigrationSchemas()
-    executor = DefaultMigrationExecutor(schemas)
-    _migrate_per_loader(
-        executor,
-        MigrationsLoader(
-            migrations_dir=CORE_MIGRATIONS_PATH,
-            module_type=ModuleTypes.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=ModuleTypes.CONTRIB,
-                module_name=contrib,
+@contextmanager
+def init_manager(
+    src_dir_path: Path,
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+    app_models_path: Path | None = None,
+    app_transactions_path: Path | None = None,
+    app_fixtures_path: Path | None = None,
+    **settings_options: Any,
+) -> Iterator[AmsdalManager]:
+    app_models_path = app_models_path or src_dir_path / 'models'
+    app_transactions_path = app_transactions_path or src_dir_path / 'transactions'
+    app_fixtures_path = app_fixtures_path or src_dir_path / 'fixtures'
+
+    with ExitStack() as stack:
+        tmp_dir = stack.enter_context(temp_dir())
+        stack.enter_context(init_config(db_execution_type, lakehouse_option, state_option))
+        stack.enter_context(
+            override_settings(
+                APP_PATH=tmp_dir,
+                USER_MODELS_MODULE_PATH=app_models_path,
+                TRANSACTIONS_MODULE_PATH=app_transactions_path,
+                FIXTURES_MODULE_PATH=app_fixtures_path,
+                **settings_options,
             ),
         )
 
-    schema_loader = JsonClassSchemaLoader(app_schemas_path)
+        manager = AmsdalManager()
+        manager.setup()
+        manager.post_setup()  # type: ignore[call-arg]
 
-    for class_schema in schema_loader.iter_app_schemas():
-        for _operation_data in FileMigrationGenerator.build_operations(
-            SchemaTypes.USER,
-            class_schema.object_schema,
-            None,
-        ):
-            _operation_name = FileMigrationWriter.operation_name_map[_operation_data.type]
-            kwargs = {}
+        try:
+            yield manager
+        finally:
+            manager.teardown()
+            AmsdalManager.invalidate()
+            AsyncDataApplication.invalidate()
 
-            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=SchemaTypes.USER,
-                class_name=_operation_data.class_name,
-                **kwargs,
+@asynccontextmanager
+async def async_init_manager(
+    src_dir_path: Path,
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+    app_models_path: Path | None = None,
+    app_transactions_path: Path | None = None,
+    app_fixtures_path: Path | None = None,
+    **settings_options: Any,
+) -> AsyncIterator[AsyncAmsdalManager]:
+    app_models_path = app_models_path or src_dir_path / 'models'
+    app_transactions_path = app_transactions_path or src_dir_path / 'transactions'
+    app_fixtures_path = app_fixtures_path or src_dir_path / 'fixtures'
+
+    with ExitStack() as stack:
+        tmp_dir = stack.enter_context(temp_dir())
+        stack.enter_context(init_config(db_execution_type, lakehouse_option, state_option, is_async=True))
+        stack.enter_context(
+            override_settings(
+                APP_PATH=tmp_dir,
+                USER_MODELS_MODULE_PATH=app_models_path,
+                TRANSACTIONS_MODULE_PATH=app_transactions_path,
+                FIXTURES_MODULE_PATH=app_fixtures_path,
+                **settings_options,
             )
+        )
+        manager = AsyncAmsdalManager()
+        await manager.setup()
+        await manager.post_setup()  # type: ignore[call-arg,misc]
 
-            _operation.forward(executor)
-
-    executor.flush_buffer()
+        try:
+            yield manager
+        finally:
+            await manager.teardown()
+            await AsyncDataApplication().teardown()
+            AsyncAmsdalManager.invalidate()
+            AsyncDataApplication.invalidate()
 
 
 @contextmanager
 def init_manager_and_migrate(
+    src_dir_path: Path,
     db_execution_type: DbExecutionType,
     lakehouse_option: LakehouseOption,
     state_option: StateOption,
+    app_models_path: Path | None = None,
+    app_transactions_path: Path | None = None,
+    app_fixtures_path: Path | None = None,
+    **settings_options: Any,
 ) -> Generator[AmsdalManager, Any, None]:
     with init_manager(
+        src_dir_path=src_dir_path,
         db_execution_type=db_execution_type,
         lakehouse_option=lakehouse_option,
         state_option=state_option,
+        app_models_path=app_models_path,
+        app_transactions_path=app_transactions_path,
+        app_fixtures_path=app_fixtures_path,
+        **settings_options,
     ) as manager:
-        app_schemas_path = TESTS_DIR / 'src' / 'models'
-        manager.build(
-            source_models_path=app_schemas_path,
-            source_fixtures_path=app_schemas_path,
-            source_transactions_path=app_schemas_path,
-            source_static_files_path=app_schemas_path,
-            source_migrations_path=app_schemas_path,
-        )
+        migrate()
+        manager.authenticate()
+        manager.init_classes()
+        HistoricalSchemaVersionManager()._cache_object_classes.clear()
+
+        yield manager
+
+
+@asynccontextmanager
+async def async_init_manager_and_migrate(
+    src_dir_path: Path,
+    db_execution_type: DbExecutionType,
+    lakehouse_option: LakehouseOption,
+    state_option: StateOption,
+    app_models_path: Path | None = None,
+    app_transactions_path: Path | None = None,
+    app_fixtures_path: Path | None = None,
+    **settings_options: Any,
+) -> AsyncGenerator[AsyncAmsdalManager, Any]:
+    async with async_init_manager(
+        src_dir_path=src_dir_path,
+        db_execution_type=db_execution_type,
+        lakehouse_option=lakehouse_option,
+        state_option=state_option,
+        app_models_path=app_models_path,
+        app_transactions_path=app_transactions_path,
+        app_fixtures_path=app_fixtures_path,
+        **settings_options,
+    ) as manager:
+        await async_migrate()
         manager.authenticate()
-        migrate(app_schemas_path)
         manager.init_classes()
+        AsyncHistoricalSchemaVersionManager()._cache_object_classes.clear()
 
         yield manager