beanis 0.0.1__py3-none-any.whl

beanis/__init__.py

@@ -0,0 +1,50 @@
+from beanis.migrations.controllers.free_fall import free_fall_migration
+from beanis.migrations.controllers.iterative import iterative_migration
+from beanis.odm.actions import (
+    After,
+    Before,
+    Delete,
+    Insert,
+    Replace,
+    Save,
+    SaveChanges,
+    Update,
+    ValidateOnSave,
+    after_event,
+    before_event,
+)
+from beanis.odm.custom_types import DecimalAnnotation
+from beanis.odm.documents import (
+    Document,
+    MergeStrategy,
+)
+
+from beanis.odm.settings.timeseries import Granularity, TimeSeriesConfig
+from beanis.odm.utils.init import init_beanis
+
+__version__ = "0.0.1"
+__all__ = [
+    # ODM
+    "Document",
+    "init_beanis",
+    "TimeSeriesConfig",
+    "Granularity",
+    "MergeStrategy",
+    # Actions
+    "before_event",
+    "after_event",
+    "Insert",
+    "Replace",
+    "Save",
+    "SaveChanges",
+    "ValidateOnSave",
+    "Delete",
+    "Before",
+    "After",
+    "Update",
+    # Migrations
+    "iterative_migration",
+    "free_fall_migration",
+    # Custom Types
+    "DecimalAnnotation",
+]

beanis/exceptions.py

@@ -0,0 +1,74 @@
+class WrongDocumentUpdateStrategy(Exception):
+    pass
+
+
+class DocumentNotFound(Exception):
+    pass
+
+
+class DocumentAlreadyCreated(Exception):
+    pass
+
+
+class DocumentWasNotSaved(Exception):
+    pass
+
+
+class CollectionWasNotInitialized(Exception):
+    pass
+
+
+class MigrationException(Exception):
+    pass
+
+
+class ReplaceError(Exception):
+    pass
+
+
+class StateManagementIsTurnedOff(Exception):
+    pass
+
+
+class StateNotSaved(Exception):
+    pass
+
+
+class RevisionIdWasChanged(Exception):
+    pass
+
+
+class NotSupported(Exception):
+    pass
+
+
+class MongoDBVersionError(Exception):
+    pass
+
+
+class ViewWasNotInitialized(Exception):
+    pass
+
+
+class ViewHasNoSettings(Exception):
+    pass
+
+
+class UnionHasNoRegisteredDocs(Exception):
+    pass
+
+
+class UnionDocNotInited(Exception):
+    pass
+
+
+class DocWasNotRegisteredInUnionClass(Exception):
+    pass
+
+
+class Deprecation(Exception):
+    pass
+
+
+class ApplyChangesException(Exception):
+    pass

beanis/executors/migrate.py

@@ -0,0 +1,221 @@
+import asyncio
+import logging
+import os
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import click
+import toml
+
+from beanis.migrations import template
+from beanis.migrations.database import DBHandler
+from beanis.migrations.models import RunningDirections, RunningMode
+from beanis.migrations.runner import MigrationNode
+
+logging.basicConfig(format="%(message)s", level=logging.INFO)
+
+
+class MigrationSettings:
+    def __init__(self, **kwargs):
+        self.direction = (
+            kwargs.get("direction")
+            or self.get_env_value("direction")
+            or self.get_from_toml("direction")
+            or RunningDirections.FORWARD
+        )
+
+        self.distance = int(
+            kwargs.get("distance")
+            or self.get_env_value("distance")
+            or self.get_from_toml("distance")
+            or 0
+        )
+        self.connection_uri = str(
+            kwargs.get("connection_uri")
+            or self.get_env_value("connection_uri")
+            or self.get_from_toml("connection_uri")
+        )
+        self.database_name = str(
+            kwargs.get("database_name")
+            or self.get_env_value("database_name")
+            or self.get_from_toml("database_name")
+        )
+        self.path = Path(
+            kwargs.get("path")
+            or self.get_env_value("path")
+            or self.get_from_toml("path")
+        )
+        self.allow_index_dropping = bool(
+            kwargs.get("allow_index_dropping")
+            or self.get_env_value("allow_index_dropping")
+            or self.get_from_toml("allow_index_dropping")
+            or False
+        )
+        self.use_transaction = bool(kwargs.get("use_transaction"))
+
+    @staticmethod
+    def get_env_value(field_name) -> Any:
+        if field_name == "connection_uri":
+            value = (
+                os.environ.get("BEANIE_URI")
+                or os.environ.get("BEANIE_CONNECTION_URI")
+                or os.environ.get("BEANIE_CONNECTION_STRING")
+                or os.environ.get("BEANIE_MONGODB_DSN")
+                or os.environ.get("BEANIE_MONGODB_URI")
+                or os.environ.get("beanie_uri")
+                or os.environ.get("beanie_connection_uri")
+                or os.environ.get("beanie_connection_string")
+                or os.environ.get("beanie_mongodb_dsn")
+                or os.environ.get("beanie_mongodb_uri")
+            )
+        elif field_name == "database_name":
+            value = (
+                os.environ.get("BEANIE_DB")
+                or os.environ.get("BEANIE_DB_NAME")
+                or os.environ.get("BEANIE_DATABASE_NAME")
+                or os.environ.get("beanie_db")
+                or os.environ.get("beanie_db_name")
+                or os.environ.get("beanie_database_name")
+            )
+        else:
+            value = os.environ.get(
+                f"BEANIE_{field_name.upper()}"
+            ) or os.environ.get(f"beanie_{field_name.lower()}")
+        return value
+
+    @staticmethod
+    def get_from_toml(field_name) -> Any:
+        path = Path("pyproject.toml")
+        if path.is_file():
+            val = (
+                toml.load(path)
+                .get("tool", {})
+                .get("beanis", {})
+                .get("migrations", {})
+            )
+        else:
+            val = {}
+        return val.get(field_name)
+
+
+@click.group()
+def migrations():
+    pass
+
+
+async def run_migrate(settings: MigrationSettings):
+    DBHandler.set_db(settings.connection_uri, settings.database_name)
+    root = await MigrationNode.build(settings.path)
+    mode = RunningMode(
+        direction=settings.direction, distance=settings.distance
+    )
+    await root.run(
+        mode=mode,
+        allow_index_dropping=settings.allow_index_dropping,
+        use_transaction=settings.use_transaction,
+    )
+
+
+@migrations.command()
+@click.option(
+    "--forward",
+    "direction",
+    required=False,
+    flag_value="FORWARD",
+    help="Roll the migrations forward. This is default",
+)
+@click.option(
+    "--backward",
+    "direction",
+    required=False,
+    flag_value="BACKWARD",
+    help="Roll the migrations backward",
+)
+@click.option(
+    "-d",
+    "--distance",
+    required=False,
+    help="How many migrations should be done since the current? "
+    "0 - all the migrations. Default is 0",
+)
+@click.option(
+    "-uri",
+    "--connection-uri",
+    required=False,
+    type=str,
+    help="MongoDB connection URI",
+)
+@click.option(
+    "-db", "--database_name", required=False, type=str, help="DataBase name"
+)
+@click.option(
+    "-p",
+    "--path",
+    required=False,
+    type=str,
+    help="Path to the migrations directory",
+)
+@click.option(
+    "--allow-index-dropping/--forbid-index-dropping",
+    required=False,
+    default=False,
+    help="if allow-index-dropping is set, Beanie will drop indexes from your collection",
+)
+@click.option(
+    "--use-transaction/--no-use-transaction",
+    required=False,
+    default=True,
+    help="Enable or disable the use of transactions during migration. "
+    "When enabled (--use-transaction), Beanie uses transactions for migration, "
+    "which necessitates a replica set. When disabled (--no-use-transaction), "
+    "migrations occur without transactions.",
+)
+def migrate(
+    direction,
+    distance,
+    connection_uri,
+    database_name,
+    path,
+    allow_index_dropping,
+    use_transaction,
+):
+    settings_kwargs = {}
+    if direction:
+        settings_kwargs["direction"] = direction
+    if distance:
+        settings_kwargs["distance"] = distance
+    if connection_uri:
+        settings_kwargs["connection_uri"] = connection_uri
+    if database_name:
+        settings_kwargs["database_name"] = database_name
+    if path:
+        settings_kwargs["path"] = path
+    if allow_index_dropping:
+        settings_kwargs["allow_index_dropping"] = allow_index_dropping
+    settings_kwargs["use_transaction"] = use_transaction
+    settings = MigrationSettings(**settings_kwargs)
+
+    asyncio.run(run_migrate(settings))
+
+
+@migrations.command()
+@click.option("-n", "--name", required=True, type=str, help="Migration name")
+@click.option(
+    "-p",
+    "--path",
+    required=True,
+    type=str,
+    help="Path to the migrations directory",
+)
+def new_migration(name, path):
+    path = Path(path)
+    ts_string = datetime.now().strftime("%Y%m%d%H%M%S")
+    file_name = f"{ts_string}_{name}.py"
+
+    shutil.copy(template.__file__, path / file_name)
+
+
+if __name__ == "__main__":
+    migrations()

beanis/migrations/controllers/base.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from typing import List, Type
+
+from beanis.odm.documents import Document
+
+
+class BaseMigrationController(ABC):
+    def __init__(self, function):
+        self.function = function
+
+    @abstractmethod
+    async def run(self, session):
+        pass
+
+    @property
+    @abstractmethod
+    def models(self) -> List[Type[Document]]:
+        pass

beanis/migrations/controllers/free_fall.py

@@ -0,0 +1,28 @@
+from inspect import signature
+from typing import List, Type
+
+from beanis.migrations.controllers.base import BaseMigrationController
+from beanis.odm.documents import Document
+
+
+def free_fall_migration(document_models: List[Type[Document]]):
+    class FreeFallMigrationController(BaseMigrationController):
+        def __init__(self, function):
+            self.function = function
+            self.function_signature = signature(function)
+            self.document_models = document_models
+
+        def __call__(self, *args, **kwargs):
+            pass
+
+        @property
+        def models(self) -> List[Type[Document]]:
+            return self.document_models
+
+        async def run(self, session):
+            function_kwargs = {"session": session}
+            if "self" in self.function_signature.parameters:
+                function_kwargs["self"] = None
+            await self.function(**function_kwargs)
+
+    return FreeFallMigrationController

beanis/migrations/controllers/iterative.py

@@ -0,0 +1,130 @@
+import asyncio
+from inspect import isclass, signature
+from typing import List, Optional, Type, Union
+
+from beanis.migrations.controllers.base import BaseMigrationController
+from beanis.migrations.utils import update_dict
+from beanis.odm.documents import Document
+from beanis.odm.utils.pydantic import parse_model
+
+
+class DummyOutput:
+    def __init__(self):
+        super(DummyOutput, self).__setattr__("_internal_structure_dict", {})
+
+    def __setattr__(self, key, value):
+        self._internal_structure_dict[key] = value
+
+    def __getattr__(self, item):
+        try:
+            return self._internal_structure_dict[item]
+        except KeyError:
+            self._internal_structure_dict[item] = DummyOutput()
+            return self._internal_structure_dict[item]
+
+    def dict(self, to_parse: Optional[Union[dict, "DummyOutput"]] = None):
+        if to_parse is None:
+            to_parse = self
+        input_dict = (
+            to_parse._internal_structure_dict
+            if isinstance(to_parse, DummyOutput)
+            else to_parse
+        )
+        result_dict = {}
+        for key, value in input_dict.items():
+            if isinstance(value, (DummyOutput, dict)):
+                result_dict[key] = self.dict(to_parse=value)
+            else:
+                result_dict[key] = value
+        return result_dict
+
+
+def iterative_migration(
+    document_models: Optional[List[Type[Document]]] = None,
+    batch_size: int = 10000,
+):
+    class IterativeMigration(BaseMigrationController):
+        def __init__(self, function):
+            self.function = function
+            self.function_signature = signature(function)
+            input_signature = self.function_signature.parameters.get(
+                "input_document"
+            )
+            if input_signature is None:
+                raise RuntimeError("input_signature must not be None")
+            self.input_document_model: Type[
+                Document
+            ] = input_signature.annotation
+            output_signature = self.function_signature.parameters.get(
+                "output_document"
+            )
+            if output_signature is None:
+                raise RuntimeError("output_signature must not be None")
+            self.output_document_model: Type[
+                Document
+            ] = output_signature.annotation
+
+            if (
+                not isclass(self.input_document_model)
+                or not issubclass(self.input_document_model, Document)
+                or not isclass(self.output_document_model)
+                or not issubclass(self.output_document_model, Document)
+            ):
+                raise TypeError(
+                    "input_document and output_document "
+                    "must have annotation of Document subclass"
+                )
+
+            self.batch_size = batch_size
+
+        def __call__(self, *args, **kwargs):
+            pass
+
+        @property
+        def models(self) -> List[Type[Document]]:
+            preset_models = document_models
+            if preset_models is None:
+                preset_models = []
+            return preset_models + [
+                self.input_document_model,
+                self.output_document_model,
+            ]
+
+        async def run(self, session):
+            output_documents = []
+            all_migration_ops = []
+            async for input_document in self.input_document_model.find_all(
+                session=session
+            ):
+                output = DummyOutput()
+                function_kwargs = {
+                    "input_document": input_document,
+                    "output_document": output,
+                }
+                if "self" in self.function_signature.parameters:
+                    function_kwargs["self"] = None
+                await self.function(**function_kwargs)
+                output_dict = input_document.dict()
+                update_dict(output_dict, output.dict())
+                output_document = parse_model(
+                    self.output_document_model, output_dict
+                )
+                output_documents.append(output_document)
+
+                if len(output_documents) == self.batch_size:
+                    all_migration_ops.append(
+                        self.output_document_model.replace_many(
+                            documents=output_documents, session=session
+                        )
+                    )
+                    output_documents = []
+
+            if output_documents:
+                all_migration_ops.append(
+                    self.output_document_model.replace_many(
+                        documents=output_documents, session=session
+                    )
+                )
+            await asyncio.gather(*all_migration_ops)
+
+    return IterativeMigration

beanis/migrations/database.py

@@ -0,0 +1,16 @@
+import motor.motor_asyncio
+
+
+class DBHandler:
+    @classmethod
+    def set_db(cls, uri, db_name):
+        cls.client = motor.motor_asyncio.AsyncIOMotorClient(uri)
+        cls.database = cls.client[db_name]
+
+    @classmethod
+    def get_cli(cls):
+        return cls.client if hasattr(cls, "client") else None
+
+    @classmethod
+    def get_db(cls):
+        return cls.database if hasattr(cls, "database") else None

beanis/migrations/models.py

@@ -0,0 +1,33 @@
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import Field
+from pydantic.main import BaseModel
+
+from beanis.odm.documents import Document
+
+
+class MigrationLog(Document):
+    ts: datetime = Field(default_factory=datetime.now)
+    name: str
+    is_current: bool
+
+    class Settings:
+        name = "migrations_log"
+
+
+class RunningDirections(str, Enum):
+    FORWARD = "FORWARD"
+    BACKWARD = "BACKWARD"
+
+
+class RunningMode(BaseModel):
+    direction: RunningDirections
+    distance: int = 0
+
+
+class ParsedMigrations(BaseModel):
+    path: str
+    names: List[str]
+    current: Optional[MigrationLog] = None