amsdal 0.4.13__cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

amsdal/services/transaction_execution.pyi

@@ -0,0 +1,93 @@
+import ast
+from _typeshed import Incomplete
+from amsdal.configs.main import settings as settings
+from amsdal.errors import TransactionNotFoundError as TransactionNotFoundError
+from amsdal_utils.utils.singleton import Singleton
+from collections.abc import Callable as Callable, Generator
+from pathlib import Path
+from typing import Any
+
+logger: Incomplete
+
+def is_transaction(statement: ast.AST) -> bool:
+    """
+    Determines if a given AST statement is a transaction function.
+
+    This function checks if the provided AST statement is an asynchronous or synchronous function
+    definition that is decorated with the `transaction` decorator.
+
+    Args:
+        statement (ast.AST): The AST statement to check.
+
+    Returns:
+        bool: True if the statement is a transaction function, False otherwise.
+    """
+def annotation_is_model(annotation: Any) -> bool: ...
+
+class TransactionExecutionService(metaclass=Singleton):
+    """
+    Service for executing transactions.
+
+    This class provides methods to execute transactions, load transaction functions,
+    and handle asynchronous transactions. It ensures that transactions are executed
+    with the correct arguments and handles any necessary preprocessing of arguments.
+    """
+    _transactions: dict[str, Callable[..., Any]]
+    def __init__(self) -> None: ...
+    def execute_transaction(self, transaction_name: str, args: dict[str, Any], *, load_references: bool = True) -> Any:
+        """
+        Executes a transaction with the given name and arguments.
+
+        This method retrieves the transaction function by its name, processes the arguments,
+        and executes the transaction. It handles both synchronous and asynchronous transactions
+        and performs necessary preprocessing of arguments, such as loading references.
+
+        Args:
+            transaction_name (str): The name of the transaction to execute.
+            args (dict[str, Any]): The arguments to pass to the transaction function.
+            load_references (bool, optional): Whether to load references in the arguments. Defaults to True.
+
+        Returns:
+            Any: The result of the transaction execution.
+        """
+    async def async_execute_transaction(self, transaction_name: str, args: dict[str, Any], *, load_references: bool = True) -> Any:
+        """
+        Executes a transaction with the given name and arguments.
+
+        This method retrieves the transaction function by its name, processes the arguments,
+        and executes the transaction. It handles both synchronous and asynchronous transactions
+        and performs necessary preprocessing of arguments, such as loading references.
+
+        Args:
+            transaction_name (str): The name of the transaction to execute.
+            args (dict[str, Any]): The arguments to pass to the transaction function.
+            load_references (bool, optional): Whether to load references in the arguments. Defaults to True.
+
+        Returns:
+            Any: The result of the transaction execution.
+        """
+    def get_transaction_func(self, transaction_name: str) -> Callable[..., Any]:
+        """
+        Retrieves the transaction function by its name.
+
+        This method checks if the transaction function is already loaded in the `_transactions` dictionary.
+        If not, it attempts to load the transaction function from the available transaction definitions.
+
+        Args:
+            transaction_name (str): The name of the transaction function to retrieve.
+
+        Returns:
+            Callable[..., Any]: The transaction function corresponding to the given name.
+
+        Raises:
+            TransactionNotFoundError: If the transaction function with the specified name is not found.
+        """
+    @staticmethod
+    def _run_async_transaction(transaction_func: Callable[..., Any], args: dict[str, Any]) -> Any: ...
+    def _load_transaction(self, transaction_name: str) -> Callable[..., Any]: ...
+    @classmethod
+    def _get_transaction_definitions(cls) -> Generator[tuple[ast.FunctionDef | ast.AsyncFunctionDef, Path], None, None]: ...
+    @classmethod
+    def _iterate_module(cls, module_path: Path) -> Generator[tuple[ast.FunctionDef | ast.AsyncFunctionDef, Path], None, None]: ...
+    @classmethod
+    def _iterate_file(cls, file_path: Path) -> Generator[tuple[ast.FunctionDef | ast.AsyncFunctionDef, Path], None, None]: ...

amsdal/transactions/__init__.py

@@ -0,0 +1,13 @@
+from amsdal_data.transactions import async_transaction
+from amsdal_data.transactions import transaction
+from amsdal_data.transactions.background.schedule import SCHEDULE_TYPE
+from amsdal_data.transactions.background.schedule import Crontab
+from amsdal_data.transactions.background.schedule import ScheduleConfig
+
+__all__ = [
+    'SCHEDULE_TYPE',
+    'Crontab',
+    'ScheduleConfig',
+    'async_transaction',
+    'transaction',
+]

amsdal/transactions/__init__.pyi

@@ -0,0 +1,4 @@
+from amsdal_data.transactions import async_transaction as async_transaction, transaction as transaction
+from amsdal_data.transactions.background.schedule import Crontab as Crontab, SCHEDULE_TYPE as SCHEDULE_TYPE, ScheduleConfig as ScheduleConfig
+
+__all__ = ['SCHEDULE_TYPE', 'Crontab', 'ScheduleConfig', 'async_transaction', 'transaction']

amsdal/utils/contrib_paths.py

@@ -0,0 +1,23 @@
+from collections.abc import Generator
+from importlib import import_module
+from pathlib import Path
+
+
+def get_contrib_schemas_paths() -> Generator[Path, None, None]:
+    """
+    Retrieves paths to contribution schemas.
+
+    This function iterates over the configured contributions in the settings and yields
+    the paths to the 'models' directories for each contribution.
+
+    Returns:
+        Generator[Path, None, None]: A generator that yields paths to the 'models' directories
+        of the configured contributions.
+    """
+    from amsdal.configs.main import settings
+
+    for contrib in settings.CONTRIB:
+        module_name, *_ = contrib.rsplit('.', 2)
+        module = import_module(module_name)
+
+        yield Path(module.__file__).parent / 'models'  # type: ignore[arg-type]

amsdal/utils/contrib_paths.pyi

@@ -0,0 +1,14 @@
+from collections.abc import Generator
+from pathlib import Path
+
+def get_contrib_schemas_paths() -> Generator[Path, None, None]:
+    """
+    Retrieves paths to contribution schemas.
+
+    This function iterates over the configured contributions in the settings and yields
+    the paths to the 'models' directories for each contribution.
+
+    Returns:
+        Generator[Path, None, None]: A generator that yields paths to the 'models' directories
+        of the configured contributions.
+    """

amsdal/utils/rollback/__init__.py

@@ -0,0 +1,440 @@
+import amsdal_glue as glue
+from amsdal_data.application import AsyncDataApplication
+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.class_manager import ClassManager
+from amsdal_models.querysets.executor import LAKEHOUSE_DB_ALIAS
+
+
+@transaction
+def rollback_to_timestamp(timestamp: float) -> None:
+    """
+    Rollback the data to the given timestamp
+    Args:
+        timestamp (float): The timestamp to rollback the data to.
+    Returns:
+        None
+    """
+    class_manager = ClassManager()
+
+    lakehouse_connection = (
+        DataApplication()._application.lakehouse_connection_manager.get_connection_pool('Company').get_connection()
+    )
+
+    metadatas_to_delete = lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
+                    lookup=glue.FieldLookup.GT,
+                    right=glue.Value(timestamp),
+                ),
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
+                    lookup=glue.FieldLookup.ISNULL,
+                    right=glue.Value(True),
+                ),
+            ),
+        )
+    )
+
+    ids_to_ignore = [m.data['object_id'] for m in metadatas_to_delete]
+
+    metadatas_to_revert = lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
+                    lookup=glue.FieldLookup.GT,
+                    right=glue.Value(timestamp),
+                ),
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
+                    lookup=glue.FieldLookup.ISNULL,
+                    right=glue.Value(False),
+                ),
+            ),
+        )
+    )
+
+    transaction_ids = {m.data['transaction']['ref']['object_id'] for m in metadatas_to_revert}
+    transaction_ids.update({m.data['transaction']['ref']['object_id'] for m in metadatas_to_delete})
+    ids_to_revert = [
+        (m.data['object_id'], m.data['class_schema_reference']['ref']['object_id'])
+        for m in metadatas_to_revert
+        if m.data['object_id'] not in ids_to_ignore
+    ]
+
+    if transaction_ids:
+        _conditions = []
+        for transaction_id in transaction_ids:
+            _parent_field = glue.Field(
+                name='transaction',
+                child=glue.Field(
+                    name='ref',
+                    child=glue.Field(name='object_id'),
+                ),
+            )
+            _parent_field.child.parent = _parent_field  # type: ignore[union-attr]
+            _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
+            _conditions.append(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
+                    lookup=glue.FieldLookup.EQ,
+                    right=glue.Value(transaction_id, output_type=str),
+                )
+            )
+
+        conflict_metadata = lakehouse_connection.query(
+            query=glue.QueryStatement(
+                table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+                where=glue.Conditions(
+                    glue.Condition(
+                        left=glue.FieldReferenceExpression(
+                            field_reference=glue.FieldReference(
+                                field=glue.Field(name='updated_at'),
+                                table_name='Metadata',
+                            ),
+                        ),
+                        lookup=glue.FieldLookup.LTE,
+                        right=glue.Value(timestamp),
+                    ),
+                    glue.Conditions(*_conditions, connector=glue.FilterConnector.OR),
+                ),
+            )
+        )
+        if conflict_metadata:
+            msg = 'Cannot rollback to this timestamp because it will conflict with other transactions'
+            raise AmsdalTransactionError(msg)
+
+    for m in metadatas_to_delete:
+        class_name = m.data['class_schema_reference']['ref']['object_id']
+        model_class = class_manager.import_class(class_name)
+        obj = (
+            model_class.objects.filter(_address__object_id=m.data['object_id'])
+            .using(LAKEHOUSE_DB_ALIAS)
+            .latest()
+            .first()
+            .execute()
+        )
+
+        if obj and not obj.get_metadata().is_deleted:
+            obj.delete()
+
+    for object_id, class_name in ids_to_revert:
+        model_class = class_manager.import_class(class_name)
+
+        obj = (
+            model_class.objects.filter(_address__object_id=object_id)
+            .using(LAKEHOUSE_DB_ALIAS)
+            .latest()
+            .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():
+                setattr(obj, field, value)
+
+            obj.save()
+
+            if old_obj.get_metadata().is_deleted:
+                obj.delete()
+
+
+@transaction
+def rollback_transaction(transaction_id: str) -> None:
+    """
+    Rollback the data to the point in time before the given transaction
+    Args:
+        transaction_id (str): The transaction ID to rollback the data to.
+    Returns:
+        None
+    """
+
+    lakehouse_connection = (
+        DataApplication()._application.lakehouse_connection_manager.get_connection_pool('Company').get_connection()
+    )
+
+    _parent_field = glue.Field(
+        name='transaction',
+        child=glue.Field(
+            name='ref',
+            child=glue.Field(name='object_id'),
+        ),
+    )
+    _parent_field.child.parent = _parent_field  # type: ignore[union-attr]
+    _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
+
+    metadatas_to_revert = lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
+                    lookup=glue.FieldLookup.EQ,
+                    right=glue.Value(transaction_id, output_type=str),
+                )
+            ),
+            order_by=[
+                glue.OrderByQuery(
+                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    direction=glue.OrderDirection.DESC,
+                )
+            ],
+        )
+    )
+
+    if not metadatas_to_revert:
+        msg = 'Transaction not found'
+        raise AmsdalTransactionError(msg)
+
+    updated_at = metadatas_to_revert[0].data['updated_at']
+    rollback_to_timestamp(updated_at)
+
+
+@async_transaction
+async def async_rollback_to_timestamp(timestamp: float) -> None:
+    """
+    Rollback the data to the given timestamp
+    Args:
+        timestamp (float): The timestamp to rollback the data to.
+    Returns:
+        None
+    """
+    class_manager = ClassManager()
+
+    lakehouse_connection = await (
+        AsyncDataApplication()._application.lakehouse_connection_manager.get_connection_pool('Company').get_connection()
+    )
+
+    metadatas_to_delete = await lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
+                    lookup=glue.FieldLookup.GT,
+                    right=glue.Value(timestamp),
+                ),
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'),
+                            table_name='Metadata',
+                        ),
+                    ),
+                    lookup=glue.FieldLookup.ISNULL,
+                    right=glue.Value(True),
+                ),
+            ),
+        )
+    )
+
+    ids_to_ignore = [m.data['object_id'] for m in metadatas_to_delete]
+
+    metadatas_to_revert = await lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    ),
+                    lookup=glue.FieldLookup.GT,
+                    right=glue.Value(timestamp),
+                ),
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(
+                            field=glue.Field(name='prior_version'), table_name='Metadata'
+                        ),
+                    ),
+                    lookup=glue.FieldLookup.ISNULL,
+                    right=glue.Value(False),
+                ),
+            ),
+        )
+    )
+
+    transaction_ids = {m.data['transaction']['ref']['object_id'] for m in metadatas_to_revert}
+    transaction_ids.update({m.data['transaction']['ref']['object_id'] for m in metadatas_to_delete})
+    ids_to_revert = [
+        (m.data['object_id'], m.data['class_schema_reference']['ref']['object_id'])
+        for m in metadatas_to_revert
+        if m.data['object_id'] not in ids_to_ignore
+    ]
+
+    if transaction_ids:
+        _conditions = []
+        for transaction_id in transaction_ids:
+            _parent_field = glue.Field(
+                name='transaction',
+                child=glue.Field(
+                    name='ref',
+                    child=glue.Field(name='object_id'),
+                ),
+            )
+            _parent_field.child.parent = _parent_field  # type: ignore[union-attr]
+            _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
+            _conditions.append(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
+                    lookup=glue.FieldLookup.EQ,
+                    right=glue.Value(transaction_id, output_type=str),
+                )
+            )
+
+        conflict_metadata = await lakehouse_connection.query(
+            query=glue.QueryStatement(
+                table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+                where=glue.Conditions(
+                    glue.Condition(
+                        left=glue.FieldReferenceExpression(
+                            field_reference=glue.FieldReference(
+                                field=glue.Field(name='updated_at'),
+                                table_name='Metadata',
+                            ),
+                        ),
+                        lookup=glue.FieldLookup.LTE,
+                        right=glue.Value(timestamp),
+                    ),
+                    glue.Conditions(*_conditions, connector=glue.FilterConnector.OR),
+                ),
+            )
+        )
+        if conflict_metadata:
+            msg = 'Cannot rollback to this timestamp because it will conflict with other transactions'
+            raise AmsdalTransactionError(msg)
+
+    for m in metadatas_to_delete:
+        class_name = m.data['class_schema_reference']['ref']['object_id']
+        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)
+            .latest()
+            .first()
+            .aexecute()
+        )
+
+        if obj and not (await obj.aget_metadata()).is_deleted:
+            await obj.adelete()
+
+    for object_id, class_name in ids_to_revert:
+        model_class = class_manager.import_class(class_name)
+
+        obj = await (
+            model_class.objects.filter(_address__object_id=object_id)
+            .using(LAKEHOUSE_DB_ALIAS)
+            .latest()
+            .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 (await old_obj.amodel_dump()).items():
+                setattr(obj, field, value)
+
+            await obj.asave()
+
+            if (await old_obj.aget_metadata()).is_deleted:
+                await obj.adelete()
+
+
+@async_transaction
+async def async_rollback_transaction(transaction_id: str) -> None:
+    """
+    Rollback the data to the point in time before the given transaction
+    Args:
+        transaction_id (str): The transaction ID to rollback the data to.
+    Returns:
+        None
+    """
+    lakehouse_connection = await (
+        AsyncDataApplication()._application.lakehouse_connection_manager.get_connection_pool('Company').get_connection()
+    )
+
+    _parent_field = glue.Field(
+        name='transaction',
+        child=glue.Field(
+            name='ref',
+            child=glue.Field(name='object_id'),
+        ),
+    )
+    _parent_field.child.parent = _parent_field  # type: ignore[union-attr]
+    _parent_field.child.child.parent = _parent_field.child  # type: ignore[union-attr]
+
+    metadatas_to_revert = await lakehouse_connection.query(
+        query=glue.QueryStatement(
+            table=glue.SchemaReference(name='Metadata', version=glue.Version.LATEST),
+            where=glue.Conditions(
+                glue.Condition(
+                    left=glue.FieldReferenceExpression(
+                        field_reference=glue.FieldReference(field=_parent_field, table_name='Metadata'),
+                        output_type=str,
+                    ),
+                    lookup=glue.FieldLookup.EQ,
+                    right=glue.Value(transaction_id, output_type=str),
+                )
+            ),
+            order_by=[
+                glue.OrderByQuery(
+                    field=glue.FieldReference(field=glue.Field(name='updated_at'), table_name='Metadata'),
+                    direction=glue.OrderDirection.DESC,
+                )
+            ],
+        )
+    )
+
+    if not metadatas_to_revert:
+        msg = 'Transaction not found'
+        raise AmsdalTransactionError(msg)
+
+    updated_at = metadatas_to_revert[0].data['updated_at']
+    await async_rollback_to_timestamp(updated_at)  # type: ignore[misc]

amsdal/utils/rollback/__init__.pyi

@@ -0,0 +1,38 @@
+from amsdal_data.transactions.decorators import async_transaction, transaction
+
+@transaction
+def rollback_to_timestamp(timestamp: float) -> None:
+    """
+    Rollback the data to the given timestamp
+    Args:
+        timestamp (float): The timestamp to rollback the data to.
+    Returns:
+        None
+    """
+@transaction
+def rollback_transaction(transaction_id: str) -> None:
+    """
+    Rollback the data to the point in time before the given transaction
+    Args:
+        transaction_id (str): The transaction ID to rollback the data to.
+    Returns:
+        None
+    """
+@async_transaction
+async def async_rollback_to_timestamp(timestamp: float) -> None:
+    """
+    Rollback the data to the given timestamp
+    Args:
+        timestamp (float): The timestamp to rollback the data to.
+    Returns:
+        None
+    """
+@async_transaction
+async def async_rollback_transaction(transaction_id: str) -> None:
+    """
+    Rollback the data to the point in time before the given transaction
+    Args:
+        transaction_id (str): The transaction ID to rollback the data to.
+    Returns:
+        None
+    """

amsdal/utils/tests/enums.py

@@ -0,0 +1,18 @@
+from enum import Enum
+
+
+class DbExecutionType(str, Enum):
+    lakehouse_only = 'lakehouse_only'
+    include_state_db = 'include_state_db'
+
+
+class StateOption(str, Enum):
+    sqlite = 'sqlite'
+    postgres = 'postgres'
+
+
+class LakehouseOption(str, Enum):
+    postgres = 'postgres'
+    postgres_immutable = 'postgres-immutable'
+    sqlite = 'sqlite'
+    sqlite_immutable = 'sqlite-immutable'