iceaxe 0.7.1__cp313-cp313-macosx_11_0_arm64.whl

iceaxe/generics.py

@@ -0,0 +1,140 @@
+import types
+from inspect import isclass
+from typing import Any, Type, TypeGuard, TypeVar, Union, get_args, get_origin
+
+
+def mro_distance(obj_type: Type, target_type: Type) -> float:
+    """
+    Calculate the MRO distance between obj_type and target_type.
+    Returns a large number if no match is found.
+
+    """
+    if not isclass(obj_type):
+        obj_type = type(obj_type)
+    if not isclass(target_type):
+        target_type = type(target_type)
+
+    # Compare class types for exact match
+    if obj_type == target_type:
+        return 0
+
+    # Check if obj_type is a subclass of target_type using the MRO
+    try:
+        return obj_type.mro().index(target_type)  # type: ignore
+    except ValueError:
+        return float("inf")
+
+
+T = TypeVar("T")
+
+
+def is_type_compatible(obj_type: Type, target_type: T) -> TypeGuard[T]:
+    return _is_type_compatible(obj_type, target_type) < float("inf")
+
+
+def _is_type_compatible(obj_type: Type, target_type: Any) -> float:
+    """
+    Relatively comprehensive type compatibility checker. This function is
+    used to check if a type has has a registered object that can
+    handle it.
+
+    Specifically returns the MRO distance where 0 indicates
+    an exact match, 1 indicates a direct ancestor, and so on. Returns a large number
+    if no compatibility is found.
+
+    """
+    # Any type is compatible with any other type
+    if target_type is Any:
+        return 0
+
+    # If obj_type is a nested type, each of these types must be compatible
+    # with the corresponding type in target_type
+    if get_origin(obj_type) is Union or isinstance(obj_type, types.UnionType):
+        return max(_is_type_compatible(t, target_type) for t in get_args(obj_type))
+
+    # Handle OR types
+    if get_origin(target_type) is Union or isinstance(target_type, types.UnionType):
+        return min(_is_type_compatible(obj_type, t) for t in get_args(target_type))
+
+    # Handle Type[Values] like typehints where we want to typehint a class
+    if get_origin(target_type) == type:  # noqa: E721
+        return _is_type_compatible(obj_type, get_args(target_type)[0])
+
+    # Handle dict[str, str] like typehints
+    # We assume that each arg in order must be matched with the target type
+    obj_origin = get_origin(obj_type)
+    target_origin = get_origin(target_type)
+    if obj_origin and target_origin:
+        if obj_origin == target_origin:
+            return max(
+                _is_type_compatible(t1, t2)
+                for t1, t2 in zip(get_args(obj_type), get_args(target_type))
+            )
+        else:
+            return float("inf")
+
+    # For lists, sets, and tuple objects make sure that each object matches
+    # the target type
+    if isinstance(obj_type, (list, set, tuple)):
+        if type(obj_type) != get_origin(target_type):  # noqa: E721
+            return float("inf")
+        return max(
+            _is_type_compatible(obj, get_args(target_type)[0]) for obj in obj_type
+        )
+
+    if isinstance(target_type, type):
+        return mro_distance(obj_type, target_type)
+
+    # Default case
+    return float("inf")
+
+
+def remove_null_type(typehint: Type) -> Type:
+    if get_origin(typehint) is Union or isinstance(typehint, types.UnionType):
+        return Union[  # type: ignore
+            tuple(  # type: ignore
+                [t for t in get_args(typehint) if t != type(None)]  # noqa: E721
+            )
+        ]
+    return typehint
+
+
+def has_null_type(typehint: Type) -> bool:
+    if get_origin(typehint) is Union or isinstance(typehint, types.UnionType):
+        return any(arg == type(None) for arg in get_args(typehint))  # noqa: E721
+    return typehint == type(None)  # noqa: E721
+
+
+def get_typevar_mapping(cls):
+    """
+    Get the raw typevar mappings {typvar: generic values} for each
+    typevar in the class hierarchy of `cls`.
+
+    Shared logic with Mountaineer. TODO: Move to a shared package.
+
+    """
+    mapping: dict[Any, Any] = {}
+
+    # Traverse MRO in reverse order, except `object`
+    for base in reversed(cls.__mro__[:-1]):
+        # Skip non-generic classes
+        if not hasattr(base, "__orig_bases__"):
+            continue
+
+        for origin_base in base.__orig_bases__:
+            origin = get_origin(origin_base)
+            if origin:
+                base_params = getattr(origin, "__parameters__", [])
+                instantiated_params = get_args(origin_base)
+
+                # Update mapping with current base's mappings
+                base_mapping = dict(zip(base_params, instantiated_params))
+                for key, value in base_mapping.items():
+                    # If value is another TypeVar, resolve it if possible
+                    if isinstance(value, TypeVar) and value in mapping:
+                        mapping[key] = mapping[value]
+                    else:
+                        mapping[key] = value
+
+    # Exclude TypeVars from the final mapping
+    return mapping

iceaxe/io.py

@@ -0,0 +1,107 @@
+import asyncio
+import importlib.metadata
+from functools import lru_cache, wraps
+from json import loads as json_loads
+from pathlib import Path
+from re import search as re_search
+from typing import Any, Callable, Coroutine, TypeVar
+
+T = TypeVar("T")
+
+
+def lru_cache_async(
+    maxsize: int | None = 100,
+):
+    def decorator(
+        async_function: Callable[..., Coroutine[Any, Any, T]],
+    ):
+        @lru_cache(maxsize=maxsize)
+        @wraps(async_function)
+        def internal(*args, **kwargs):
+            coroutine = async_function(*args, **kwargs)
+            # Unlike regular coroutine functions, futures can be awaited multiple times
+            # so our caller functions can await the same future on multiple cache hits
+            return asyncio.ensure_future(coroutine)
+
+        return internal
+
+    return decorator
+
+
+def resolve_package_path(package_name: str):
+    """
+    Given a package distribution, returns the local file directory where the code
+    is located. This is resolved to the original reference if installed with `-e`
+    otherwise is a copy of the package.
+
+    NOTE: Copied from Mountaineer, refactor to a shared library.
+
+    """
+    dist = importlib.metadata.distribution(package_name)
+
+    def normalize_package(package: str):
+        return package.replace("-", "_").lower()
+
+    # Recent versions of poetry install development packages (-e .) as direct URLs
+    # https://the-hitchhikers-guide-to-packaging.readthedocs.io/en/latest/introduction.html
+    # "Path configuration files have an extension of .pth, and each line must
+    # contain a single path that will be appended to sys.path."
+    package_name = normalize_package(dist.name)
+    symbolic_links = [
+        path
+        for path in (dist.files or [])
+        if path.name.lower() == f"{package_name}.pth"
+    ]
+    dist_links = [
+        path
+        for path in (dist.files or [])
+        if path.name == "direct_url.json"
+        and re_search(package_name + r"-[0-9-.]+\.dist-info", path.parent.name.lower())
+    ]
+    explicit_links = [
+        path
+        for path in (dist.files or [])
+        if path.parent.name.lower() == package_name
+        and (
+            # Sanity check that the parent is the high level project directory
+            # by looking for common base files
+            path.name == "__init__.py"
+        )
+    ]
+
+    # The user installed code as an absolute package (ie. with pip install .) instead of
+    # as a reference. There's no need to sniff for the additional package path since
+    # we've already found it
+    if explicit_links:
+        # Right now we have a file pointer to __init__.py. Go up one level
+        # to the main package directory to return a directory
+        explicit_link = explicit_links[0]
+        return Path(str(dist.locate_file(explicit_link.parent)))
+
+    # Raw path will capture the path to the pyproject.toml file directory,
+    # not the actual package code directory
+    # Find the root, then resolve the package directory
+    raw_path: Path | None = None
+
+    if symbolic_links:
+        direct_url_path = symbolic_links[0]
+        raw_path = Path(str(dist.locate_file(direct_url_path.read_text().strip())))
+    elif dist_links:
+        dist_link = dist_links[0]
+        direct_metadata = json_loads(dist_link.read_text())
+        package_path = "/" + direct_metadata["url"].lstrip("file://").lstrip("/")
+        raw_path = Path(str(dist.locate_file(package_path)))
+
+    if not raw_path:
+        raise ValueError(
+            f"Could not find a valid path for package {dist.name}, found files: {dist.files}"
+        )
+
+    # Sniff for the presence of the code directory
+    for path in raw_path.iterdir():
+        if path.is_dir() and normalize_package(path.name) == package_name:
+            return path
+
+    raise ValueError(
+        f"No matching package found in root path: {raw_path} {list(raw_path.iterdir())}"
+    )

iceaxe/logging.py

@@ -0,0 +1,91 @@
+import logging
+from contextlib import contextmanager
+from json import dumps as json_dumps
+from logging import Formatter, StreamHandler, getLogger
+from os import environ
+from time import monotonic_ns
+
+from click import secho
+from rich.console import Console
+
+VERBOSITY_MAPPING = {
+    "INFO": logging.INFO,
+    "DEBUG": logging.DEBUG,
+    "WARNING": logging.WARNING,
+    "ERROR": logging.ERROR,
+}
+
+
+class JsonFormatter(Formatter):
+    def format(self, record):
+        log_record = {
+            "level": record.levelname,
+            "name": record.name,
+            "timestamp": self.formatTime(record, self.datefmt),
+            "message": record.getMessage(),
+        }
+        if record.exc_info:
+            log_record["exception"] = self.formatException(record.exc_info)
+        return json_dumps(log_record)
+
+
+class ColorHandler(StreamHandler):
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            if record.levelno == logging.WARNING:
+                secho(msg, fg="yellow")
+            elif record.levelno >= logging.ERROR:
+                secho(msg, fg="red")
+            else:
+                secho(msg)
+        except Exception:
+            self.handleError(record)
+
+
+def setup_logger(name, log_level=logging.DEBUG):
+    """
+    Constructor for the main logger used by Mountaineer. Provided
+    convenient defaults for log level and formatting, alongside coloring
+    of stdout/stderr messages and JSON fields for structured parsing.
+
+    """
+    logger = getLogger(name)
+    logger.setLevel(log_level)
+
+    # Create a handler that writes log records to the standard error
+    handler = ColorHandler()
+    handler.setLevel(log_level)
+
+    formatter = JsonFormatter()
+    handler.setFormatter(formatter)
+
+    logger.addHandler(handler)
+
+    return logger
+
+
+@contextmanager
+def log_time_duration(message: str):
+    """
+    Context manager to time a code block at runtime.
+
+    ```python
+    with log_time_duration("Long computation"):
+        # Simulate work
+        sleep(10)
+    ```
+
+    """
+    start = monotonic_ns()
+    yield
+    LOGGER.debug(f"{message} : Took {(monotonic_ns() - start) / 1e9:.2f}s")
+
+
+# Our global logger should only surface warnings and above by default
+LOGGER = setup_logger(
+    __name__,
+    log_level=VERBOSITY_MAPPING[environ.get("ICEAXE_LOG_LEVEL", "WARNING")],
+)
+
+CONSOLE = Console()

iceaxe/migrations/__init__.py

@@ -0,0 +1,5 @@
+from .cli import (
+    handle_apply as handle_apply,
+    handle_generate as handle_generate,
+    handle_rollback as handle_rollback,
+)

iceaxe/migrations/action_sorter.py

@@ -0,0 +1,98 @@
+from collections import defaultdict
+
+from iceaxe.schemas.db_stubs import DBObject
+
+
+class ActionTopologicalSorter:
+    """
+    Extends Python's native TopologicalSorter to group nodes by table_name. This provides
+    better semantic grouping within the migrations since most actions are oriented by their
+    parent table.
+
+    - Places cross-table dependencies and non-table actions (like types) according
+        to their default DAG order.
+    - Tables are processed in alphabetical order of their names.
+
+    """
+
+    def __init__(self, graph: dict[DBObject, list[DBObject]]):
+        self.graph = graph
+        self.in_degree = defaultdict(int)
+        self.nodes = set(graph.keys())
+
+        for node, dependencies in list(graph.items()):
+            for dep in dependencies:
+                self.in_degree[node] += 1
+                if dep not in self.nodes:
+                    self.nodes.add(dep)
+                    self.graph[dep] = []
+
+        # Order based on the original yield / creation order of the nodes
+        self.node_to_ordering = {node: i for i, node in enumerate(self.graph.keys())}
+
+    def sort(self):
+        result = []
+        root_nodes_queued = sorted(
+            [node for node in self.nodes if self.in_degree[node] == 0],
+            key=self.node_key,
+        )
+        if not root_nodes_queued and self.nodes:
+            raise ValueError("Graph contains a cycle")
+        elif not root_nodes_queued:
+            return []
+
+        # Sort by the table name and then by the node representation
+        root_nodes_queued.sort(key=self.node_key)
+
+        # Always put the non-table actions first (things like global types)
+        non_table_nodes = [
+            node for node in root_nodes_queued if not hasattr(node, "table_name")
+        ]
+        root_nodes_queued = [
+            node for node in root_nodes_queued if node not in non_table_nodes
+        ]
+        root_nodes_queued = non_table_nodes + root_nodes_queued
+
+        queue = [root_nodes_queued.pop(0)]
+        processed = set()
+
+        while True:
+            if not queue:
+                # Pop another root node, if available
+                if root_nodes_queued:
+                    queue.append(root_nodes_queued.pop(0))
+                    continue
+                else:
+                    # If no more root nodes are available and no more work to be done
+                    break
+
+            current_node = queue.pop(0)
+
+            result.append(current_node)
+            processed.add(current_node)
+
+            # Newly unblocked nodes, since we've resolved their dependencies
+            # with the current processing
+            new_ready = []
+            for dependent, deps in self.graph.items():
+                if current_node in deps and dependent not in processed:
+                    self.in_degree[dependent] -= 1
+                    if self.in_degree[dependent] == 0:
+                        new_ready.append(dependent)
+
+            # Add newly ready nodes to queue in sorted order
+            queue.extend(
+                sorted(new_ready, key=lambda node: self.node_to_ordering[node])
+            )
+
+        if len(result) != len(self.nodes):
+            raise ValueError("Graph contains a cycle")
+
+        return result
+
+    @staticmethod
+    def node_key(node: DBObject):
+        # Not all objects specify a table_name, but if they do we want to explicitly
+        # sort before the representation
+        table_name = getattr(node, "table_name", "")
+        return (table_name, node.representation())

iceaxe/migrations/cli.py

@@ -0,0 +1,228 @@
+from inspect import isclass
+from time import monotonic_ns
+
+from iceaxe.base import DBModelMetaclass, TableBase
+from iceaxe.io import resolve_package_path
+from iceaxe.logging import CONSOLE
+from iceaxe.schemas.db_serializer import DatabaseSerializer
+from iceaxe.session import DBConnection
+
+
+async def handle_generate(
+    package: str, db_connection: DBConnection, message: str | None = None
+):
+    """
+    Creates a new migration definition file, comparing the previous version
+    (if it exists) with the current schema.
+
+    :param package: The current python package name. This should match the name of the
+        project that's specified in pyproject.toml or setup.py.
+
+    :param message: An optional message to include in the migration file. Helps
+        with describing changes and searching for past migration logic over time.
+
+    ```python {{sticky: True}}
+    from iceaxe.migrations.cli import handle_generate
+    from click import command, option
+
+    @command()
+    @option("--message", help="A message to include in the migration file.")
+    def generate_migration(message: str):
+        db_connection = DBConnection(...)
+        handle_generate("my_project", db_connection, message=message)
+    ```
+    """
+    # Any local imports must be done here to avoid circular imports because migrations.__init__
+    # imports this file.
+    from iceaxe.migrations.client_io import fetch_migrations
+    from iceaxe.migrations.generator import MigrationGenerator
+    from iceaxe.migrations.migrator import Migrator
+
+    CONSOLE.print("[bold blue]Generating migration to current schema")
+
+    CONSOLE.print(
+        "[grey58]Note that Iceaxe's migration support is well tested but still in beta."
+    )
+    CONSOLE.print(
+        "[grey58]File an issue @ https://github.com/piercefreeman/iceaxe/issues if you encounter any problems."
+    )
+
+    # Locate the migrations directory that belongs to this project
+    package_path = resolve_package_path(package)
+    migrations_path = package_path / "migrations"
+
+    # Create the path if it doesn't exist
+    migrations_path.mkdir(exist_ok=True)
+    if not (migrations_path / "__init__.py").exists():
+        (migrations_path / "__init__.py").touch()
+
+    # Get all of the instances that have been registered
+    # in memory scope by the user.
+    models = [
+        cls
+        for cls in DBModelMetaclass.get_registry()
+        if isclass(cls) and issubclass(cls, TableBase)
+    ]
+
+    db_serializer = DatabaseSerializer()
+    db_objects = []
+    async for values in db_serializer.get_objects(db_connection):
+        db_objects.append(values)
+
+    migration_generator = MigrationGenerator()
+    up_objects = list(migration_generator.serializer.delegate(models))
+
+    # Get the current revision from the database, this should represent the "down" revision
+    # for the new migration
+    migrator = Migrator(db_connection)
+    await migrator.init_db()
+    current_revision = await migrator.get_active_revision()
+
+    # Make sure there's not a duplicate revision that already have this down revision. If so that means
+    # that we will have two conflicting migration chains
+    migration_revisions = fetch_migrations(migrations_path)
+    conflict_migrations = [
+        migration
+        for migration in migration_revisions
+        if migration.down_revision == current_revision
+    ]
+    if conflict_migrations:
+        up_revisions = {migration.up_revision for migration in conflict_migrations}
+        raise ValueError(
+            f"Found conflicting migrations with down revision {current_revision} (conflicts: {up_revisions}).\n"
+            "If you're trying to generate a new migration, make sure to apply the previous migration first - or delete the old one and recreate."
+        )
+
+    migration_code, revision = await migration_generator.new_migration(
+        db_objects,
+        up_objects,
+        down_revision=current_revision,
+        user_message=message,
+    )
+
+    # Create the migration file. The change of a conflict with this timestamp is very low, but we make sure
+    # not to override any existing files anyway.
+    migration_file_path = migrations_path / f"rev_{revision}.py"
+    if migration_file_path.exists():
+        raise ValueError(
+            f"Migration file {migration_file_path} already exists. Wait a second and try again."
+        )
+
+    migration_file_path.write_text(migration_code)
+
+    CONSOLE.print(f"[bold green]New migration added: {migration_file_path.name}")
+
+
+async def handle_apply(
+    package: str,
+    db_connection: DBConnection,
+):
+    """
+    Applies all migrations that have not been applied to the database.
+
+    :param package: The current python package name. This should match the name of the
+        project that's specified in pyproject.toml or setup.py.
+
+    """
+    from iceaxe.migrations.client_io import fetch_migrations, sort_migrations
+    from iceaxe.migrations.migrator import Migrator
+
+    migrations_path = resolve_package_path(package) / "migrations"
+    if not migrations_path.exists():
+        raise ValueError(f"Migrations path {migrations_path} does not exist.")
+
+    # Load all the migration files into memory and locate the subclasses of MigrationRevisionBase
+    migration_revisions = fetch_migrations(migrations_path)
+    migration_revisions = sort_migrations(migration_revisions)
+
+    # Get the current revision from the database
+    migrator = Migrator(db_connection)
+    await migrator.init_db()
+    current_revision = await migrator.get_active_revision()
+
+    CONSOLE.print(f"Current revision: {current_revision}")
+
+    # Find the item in the sequence that has down_revision equal to the current_revision
+    # This indicates the next migration to apply
+    next_migration_index = None
+    for i, revision in enumerate(migration_revisions):
+        if revision.down_revision == current_revision:
+            next_migration_index = i
+            break
+
+    if next_migration_index is None:
+        raise ValueError(
+            f"Could not find a migration to apply after revision {current_revision}."
+        )
+
+    # Get the chain after this index, this should indicate the next migration to apply
+    migration_chain = migration_revisions[next_migration_index:]
+    CONSOLE.print(f"Applying {len(migration_chain)} migrations...")
+
+    for migration in migration_chain:
+        with CONSOLE.status(
+            f"[bold blue]Applying {migration.up_revision}...", spinner="dots"
+        ):
+            start = monotonic_ns()
+            await migration._handle_up(db_connection)
+
+        CONSOLE.print(
+            f"[bold green]🚀 Applied {migration.up_revision} in {(monotonic_ns() - start) / 1e9:.2f}s"
+        )
+
+
+async def handle_rollback(
+    package: str,
+    db_connection: DBConnection,
+):
+    """
+    Rolls back the last migration that was applied to the database.
+
+    :param package: The current python package name. This should match the name of the
+        project that's specified in pyproject.toml or setup.py.
+
+    """
+    from iceaxe.migrations.client_io import fetch_migrations, sort_migrations
+    from iceaxe.migrations.migrator import Migrator
+
+    migrations_path = resolve_package_path(package) / "migrations"
+    if not migrations_path.exists():
+        raise ValueError(f"Migrations path {migrations_path} does not exist.")
+
+    # Load all the migration files into memory and locate the subclasses of MigrationRevisionBase
+    migration_revisions = fetch_migrations(migrations_path)
+    migration_revisions = sort_migrations(migration_revisions)
+
+    # Get the current revision from the database
+    migrator = Migrator(db_connection)
+    await migrator.init_db()
+    current_revision = await migrator.get_active_revision()
+
+    CONSOLE.print(f"Current revision: {current_revision}")
+
+    # Find the item in the sequence that has down_revision equal to the current_revision
+    # This indicates the next migration to apply
+    this_migration_index = None
+    for i, revision in enumerate(migration_revisions):
+        if revision.up_revision == current_revision:
+            this_migration_index = i
+            break
+
+    if this_migration_index is None:
+        raise ValueError(
+            f"Could not find a migration matching {current_revision} for rollback."
+        )
+
+    # Get the chain after this index, this should indicate the next migration to apply
+    this_migration = migration_revisions[this_migration_index]
+
+    with CONSOLE.status(
+        f"[bold blue]Rolling back revision {this_migration.up_revision} to {this_migration.down_revision}...",
+        spinner="dots",
+    ):
+        start = monotonic_ns()
+        await this_migration._handle_down(db_connection)
+
+    CONSOLE.print(
+        f"[bold green]🪃 Rolled back migration to {this_migration.down_revision} in {(monotonic_ns() - start) / 1e9:.2f}s"
+    )