reflex 0.8.14.post1__py3-none-any.whl → 0.8.15a1__py3-none-any.whl

reflex/model.py

@@ -5,69 +5,22 @@ from __future__ import annotations
 import re
 from collections import defaultdict
 from contextlib import suppress
-from typing import Any, ClassVar
-
-import alembic.autogenerate
-import alembic.command
-import alembic.config
-import alembic.operations.ops
-import alembic.runtime.environment
-import alembic.script
-import alembic.util
-import sqlalchemy
-import sqlalchemy.exc
-import sqlalchemy.ext.asyncio
-import sqlalchemy.orm
-from alembic.runtime.migration import MigrationContext
-from alembic.script.base import Script
-
-from reflex.base import Base
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, Any, ClassVar
+
 from reflex.config import get_config
 from reflex.environment import environment
 from reflex.utils import console
-from reflex.utils.compat import sqlmodel, sqlmodel_field_has_primary_key
-
-_ENGINE: dict[str, sqlalchemy.engine.Engine] = {}
-_ASYNC_ENGINE: dict[str, sqlalchemy.ext.asyncio.AsyncEngine] = {}
-_AsyncSessionLocal: dict[str | None, sqlalchemy.ext.asyncio.async_sessionmaker] = {}
-
-# Import AsyncSession _after_ reflex.utils.compat
-from sqlmodel.ext.asyncio.session import AsyncSession  # noqa: E402
-
-
-def format_revision(
-    rev: Script,
-    current_rev: str | None,
-    current_reached_ref: list[bool],
-) -> str:
-    """Format a single revision for display.
-
-    Args:
-        rev: The alembic script object
-        current_rev: The currently applied revision ID
-        current_reached_ref: Mutable reference to track if we've reached current revision
-
-    Returns:
-        Formatted string for display
-    """
-    current = rev.revision
-    message = rev.doc
-
-    # Determine if this migration is applied
-    if current_rev is None:
-        is_applied = False
-    elif current == current_rev:
-        is_applied = True
-        current_reached_ref[0] = True
-    else:
-        is_applied = not current_reached_ref[0]
+from reflex.utils.compat import sqlmodel_field_has_primary_key
+from reflex.utils.serializers import serializer
 
-    # Show checkmark or X with colors
-    status_icon = "[green]✓[/green]" if is_applied else "[red]✗[/red]"
-    head_marker = " (head)" if rev.is_head else ""
+if TYPE_CHECKING:
+    import sqlalchemy
+    import sqlmodel
 
-    # Format output with message
-    return f"  [{status_icon}] {current}{head_marker}, {message}"
+    SQLModelOrSqlAlchemy = (
+        type[sqlmodel.SQLModel] | type[sqlalchemy.orm.DeclarativeBase]
+    )
 
 
 def _safe_db_url_for_logging(url: str) -> str:
@@ -82,536 +35,606 @@ def _safe_db_url_for_logging(url: str) -> str:
     return re.sub(r"://[^@]+@", "://<username>:<password>@", url)
 
 
-def get_engine_args(url: str | None = None) -> dict[str, Any]:
-    """Get the database engine arguments.
-
-    Args:
-        url: The database url.
-
-    Returns:
-        The database engine arguments as a dict.
-    """
-    kwargs: dict[str, Any] = {
-        # Print the SQL queries if the log level is INFO or lower.
-        "echo": environment.SQLALCHEMY_ECHO.get(),
-        # Check connections before returning them.
-        "pool_pre_ping": environment.SQLALCHEMY_POOL_PRE_PING.get(),
-        "pool_size": environment.SQLALCHEMY_POOL_SIZE.get(),
-        "max_overflow": environment.SQLALCHEMY_MAX_OVERFLOW.get(),
-        "pool_recycle": environment.SQLALCHEMY_POOL_RECYCLE.get(),
-        "pool_timeout": environment.SQLALCHEMY_POOL_TIMEOUT.get(),
-    }
-    conf = get_config()
-    url = url or conf.db_url
-    if url is not None and url.startswith("sqlite"):
-        # Needed for the admin dash on sqlite.
-        kwargs["connect_args"] = {"check_same_thread": False}
-    return kwargs
-
-
-def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
-    """Get the database engine.
-
-    Args:
-        url: the DB url to use.
-
-    Returns:
-        The database engine.
-
-    Raises:
-        ValueError: If the database url is None.
-    """
-    conf = get_config()
-    url = url or conf.db_url
-    if url is None:
-        msg = "No database url configured"
-        raise ValueError(msg)
-
-    global _ENGINE
-    if url in _ENGINE:
-        return _ENGINE[url]
-
-    if not environment.ALEMBIC_CONFIG.get().exists():
-        console.warn(
-            "Database is not initialized, run [bold]reflex db init[/bold] first."
-        )
-    _ENGINE[url] = sqlmodel.create_engine(
-        url,
-        **get_engine_args(url),
+def _print_db_not_available(*args, **kwargs):
+    msg = (
+        "Database is not available. Please install the required packages: "
+        "`pip install reflex[db]`."
     )
-    return _ENGINE[url]
-
+    raise ImportError(msg)
 
-def get_async_engine(url: str | None) -> sqlalchemy.ext.asyncio.AsyncEngine:
-    """Get the async database engine.
 
-    Args:
-        url: The database url.
-
-    Returns:
-        The async database engine.
-
-    Raises:
-        ValueError: If the async database url is None.
-    """
-    if url is None:
-        conf = get_config()
-        url = conf.async_db_url
-        if url is not None and conf.db_url is not None:
-            async_db_url_tail = url.partition("://")[2]
-            db_url_tail = conf.db_url.partition("://")[2]
-            if async_db_url_tail != db_url_tail:
-                console.warn(
-                    f"async_db_url `{_safe_db_url_for_logging(url)}` "
-                    "should reference the same database as "
-                    f"db_url `{_safe_db_url_for_logging(conf.db_url)}`."
-                )
-    if url is None:
-        msg = "No async database url configured"
-        raise ValueError(msg)
-
-    global _ASYNC_ENGINE
-    if url in _ASYNC_ENGINE:
-        return _ASYNC_ENGINE[url]
+class _ClassThatErrorsOnInit:
+    def __init__(self, *args, **kwargs):
+        _print_db_not_available(*args, **kwargs)
 
-    if not environment.ALEMBIC_CONFIG.get().exists():
-        console.warn(
-            "Database is not initialized, run [bold]reflex db init[/bold] first."
-        )
-    _ASYNC_ENGINE[url] = sqlalchemy.ext.asyncio.create_async_engine(
-        url,
-        **get_engine_args(url),
-    )
-    return _ASYNC_ENGINE[url]
 
+if find_spec("sqlalchemy"):
+    import sqlalchemy
+    import sqlalchemy.exc
+    import sqlalchemy.ext.asyncio
+    import sqlalchemy.orm
 
-async def get_db_status() -> dict[str, bool]:
-    """Checks the status of the database connection.
+    _ENGINE: dict[str, sqlalchemy.engine.Engine] = {}
+    _ASYNC_ENGINE: dict[str, sqlalchemy.ext.asyncio.AsyncEngine] = {}
 
-    Attempts to connect to the database and execute a simple query to verify connectivity.
+    def get_engine_args(url: str | None = None) -> dict[str, Any]:
+        """Get the database engine arguments.
 
-    Returns:
-        The status of the database connection.
-    """
-    status = True
-    try:
-        engine = get_engine()
-        with engine.connect() as connection:
-            connection.execute(sqlalchemy.text("SELECT 1"))
-    except sqlalchemy.exc.OperationalError:
-        status = False
+        Args:
+            url: The database url.
 
-    return {"db": status}
+        Returns:
+            The database engine arguments as a dict.
+        """
+        kwargs: dict[str, Any] = {
+            # Print the SQL queries if the log level is INFO or lower.
+            "echo": environment.SQLALCHEMY_ECHO.get(),
+            # Check connections before returning them.
+            "pool_pre_ping": environment.SQLALCHEMY_POOL_PRE_PING.get(),
+            "pool_size": environment.SQLALCHEMY_POOL_SIZE.get(),
+            "max_overflow": environment.SQLALCHEMY_MAX_OVERFLOW.get(),
+            "pool_recycle": environment.SQLALCHEMY_POOL_RECYCLE.get(),
+            "pool_timeout": environment.SQLALCHEMY_POOL_TIMEOUT.get(),
+        }
+        conf = get_config()
+        url = url or conf.db_url
+        if url is not None and url.startswith("sqlite"):
+            # Needed for the admin dash on sqlite.
+            kwargs["connect_args"] = {"check_same_thread": False}
+        return kwargs
 
+    def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
+        """Get the database engine.
 
-SQLModelOrSqlAlchemy = type[sqlmodel.SQLModel] | type[sqlalchemy.orm.DeclarativeBase]
+        Args:
+            url: the DB url to use.
 
+        Returns:
+            The database engine.
 
-class ModelRegistry:
-    """Registry for all models."""
+        Raises:
+            ValueError: If the database url is None.
+        """
+        conf = get_config()
+        url = url or conf.db_url
+        if url is None:
+            msg = "No database url configured"
+            raise ValueError(msg)
 
-    models: ClassVar[set[SQLModelOrSqlAlchemy]] = set()
+        global _ENGINE
+        if url in _ENGINE:
+            return _ENGINE[url]
 
-    # Cache the metadata to avoid re-creating it.
-    _metadata: ClassVar[sqlalchemy.MetaData | None] = None
+        if not environment.ALEMBIC_CONFIG.get().exists():
+            console.warn(
+                "Database is not initialized, run [bold]reflex db init[/bold] first."
+            )
+        _ENGINE[url] = sqlalchemy.engine.create_engine(
+            url,
+            **get_engine_args(url),
+        )
+        return _ENGINE[url]
 
-    @classmethod
-    def register(cls, model: SQLModelOrSqlAlchemy):
-        """Register a model. Can be used directly or as a decorator.
+    def get_async_engine(url: str | None) -> sqlalchemy.ext.asyncio.AsyncEngine:
+        """Get the async database engine.
 
         Args:
-            model: The model to register.
+            url: The database url.
 
         Returns:
-            The model passed in as an argument (Allows decorator usage)
-        """
-        cls.models.add(model)
-        return model
-
-    @classmethod
-    def get_models(cls, include_empty: bool = False) -> set[SQLModelOrSqlAlchemy]:
-        """Get registered models.
-
-        Args:
-            include_empty: If True, include models with empty metadata.
+            The async database engine.
 
-        Returns:
-            The registered models.
+        Raises:
+            ValueError: If the async database url is None.
         """
-        if include_empty:
-            return cls.models
-        return {
-            model for model in cls.models if not cls._model_metadata_is_empty(model)
-        }
+        if url is None:
+            conf = get_config()
+            url = conf.async_db_url
+            if url is not None and conf.db_url is not None:
+                async_db_url_tail = url.partition("://")[2]
+                db_url_tail = conf.db_url.partition("://")[2]
+                if async_db_url_tail != db_url_tail:
+                    console.warn(
+                        f"async_db_url `{_safe_db_url_for_logging(url)}` "
+                        "should reference the same database as "
+                        f"db_url `{_safe_db_url_for_logging(conf.db_url)}`."
+                    )
+        if url is None:
+            msg = "No async database url configured"
+            raise ValueError(msg)
+
+        global _ASYNC_ENGINE
+        if url in _ASYNC_ENGINE:
+            return _ASYNC_ENGINE[url]
+
+        if not environment.ALEMBIC_CONFIG.get().exists():
+            console.warn(
+                "Database is not initialized, run [bold]reflex db init[/bold] first."
+            )
+        _ASYNC_ENGINE[url] = sqlalchemy.ext.asyncio.create_async_engine(
+            url,
+            **get_engine_args(url),
+        )
+        return _ASYNC_ENGINE[url]
 
-    @staticmethod
-    def _model_metadata_is_empty(model: SQLModelOrSqlAlchemy) -> bool:
-        """Check if the model metadata is empty.
+    def sqla_session(url: str | None = None) -> sqlalchemy.orm.Session:
+        """Get a bare sqlalchemy session to interact with the database.
 
         Args:
-            model: The model to check.
+            url: The database url.
 
         Returns:
-            True if the model metadata is empty, False otherwise.
+            A database session.
         """
-        return len(model.metadata.tables) == 0
+        return sqlalchemy.orm.Session(get_engine(url))
+
+    class ModelRegistry:
+        """Registry for all models."""
+
+        models: ClassVar[set[SQLModelOrSqlAlchemy]] = set()
+
+        # Cache the metadata to avoid re-creating it.
+        _metadata: ClassVar[sqlalchemy.MetaData | None] = None
+
+        @classmethod
+        def register(cls, model: SQLModelOrSqlAlchemy):
+            """Register a model. Can be used directly or as a decorator.
+
+            Args:
+                model: The model to register.
+
+            Returns:
+                The model passed in as an argument (Allows decorator usage)
+            """
+            cls.models.add(model)
+            return model
+
+        @classmethod
+        def get_models(cls, include_empty: bool = False) -> set[SQLModelOrSqlAlchemy]:
+            """Get registered models.
+
+            Args:
+                include_empty: If True, include models with empty metadata.
+
+            Returns:
+                The registered models.
+            """
+            if include_empty:
+                return cls.models
+            return {
+                model for model in cls.models if not cls._model_metadata_is_empty(model)
+            }
+
+        @staticmethod
+        def _model_metadata_is_empty(model: SQLModelOrSqlAlchemy) -> bool:
+            """Check if the model metadata is empty.
+
+            Args:
+                model: The model to check.
+
+            Returns:
+                True if the model metadata is empty, False otherwise.
+            """
+            return len(model.metadata.tables) == 0
+
+        @classmethod
+        def get_metadata(cls) -> sqlalchemy.MetaData:
+            """Get the database metadata.
+
+            Returns:
+                The database metadata.
+            """
+            if cls._metadata is not None:
+                return cls._metadata
+
+            models = cls.get_models(include_empty=False)
+
+            if len(models) == 1:
+                metadata = next(iter(models)).metadata
+            else:
+                # Merge the metadata from all the models.
+                # This allows mixing bare sqlalchemy models with sqlmodel models in one database.
+                metadata = sqlalchemy.MetaData()
+                for model in cls.get_models():
+                    for table in model.metadata.tables.values():
+                        table.to_metadata(metadata)
+
+            # Cache the metadata
+            cls._metadata = metadata
+
+            return metadata
+
+else:
+    get_engine_args = _print_db_not_available
+    get_engine = _print_db_not_available
+    get_async_engine = _print_db_not_available
+    sqla_session = _print_db_not_available
+    ModelRegistry = _ClassThatErrorsOnInit  # pyright: ignore [reportAssignmentType]
+
+if find_spec("sqlmodel") and find_spec("sqlalchemy") and find_spec("pydantic"):
+    import alembic.autogenerate
+    import alembic.command
+    import alembic.config
+    import alembic.operations.ops
+    import alembic.runtime.environment
+    import alembic.script
+    import sqlmodel
+    from alembic.runtime.migration import MigrationContext
+    from alembic.script.base import Script
+    from sqlmodel.ext.asyncio.session import AsyncSession
+
+    _AsyncSessionLocal: dict[str | None, sqlalchemy.ext.asyncio.async_sessionmaker] = {}
+
+    def format_revision(
+        rev: Script,
+        current_rev: str | None,
+        current_reached_ref: list[bool],
+    ) -> str:
+        """Format a single revision for display.
 
-    @classmethod
-    def get_metadata(cls) -> sqlalchemy.MetaData:
-        """Get the database metadata.
+        Args:
+            rev: The alembic script object
+            current_rev: The currently applied revision ID
+            current_reached_ref: Mutable reference to track if we've reached current revision
 
         Returns:
-            The database metadata.
+            Formatted string for display
         """
-        if cls._metadata is not None:
-            return cls._metadata
-
-        models = cls.get_models(include_empty=False)
-
-        if len(models) == 1:
-            metadata = next(iter(models)).metadata
+        current = rev.revision
+        message = rev.doc
+
+        # Determine if this migration is applied
+        if current_rev is None:
+            is_applied = False
+        elif current == current_rev:
+            is_applied = True
+            current_reached_ref[0] = True
         else:
-            # Merge the metadata from all the models.
-            # This allows mixing bare sqlalchemy models with sqlmodel models in one database.
-            metadata = sqlalchemy.MetaData()
-            for model in cls.get_models():
-                for table in model.metadata.tables.values():
-                    table.to_metadata(metadata)
-
-        # Cache the metadata
-        cls._metadata = metadata
-
-        return metadata
+            is_applied = not current_reached_ref[0]
 
+        # Show checkmark or X with colors
+        status_icon = "[green]✓[/green]" if is_applied else "[red]✗[/red]"
+        head_marker = " (head)" if rev.is_head else ""
 
-class Model(Base, sqlmodel.SQLModel):  # pyright: ignore [reportGeneralTypeIssues,reportIncompatibleVariableOverride]
-    """Base class to define a table in the database."""
+        # Format output with message
+        return f"  [{status_icon}] {current}{head_marker}, {message}"
 
-    # The primary key for the table.
-    id: int | None = sqlmodel.Field(default=None, primary_key=True)
+    async def get_db_status() -> dict[str, bool]:
+        """Checks the status of the database connection.
 
-    def __init_subclass__(cls):
-        """Drop the default primary key field if any primary key field is defined."""
-        non_default_primary_key_fields = [
-            field_name
-            for field_name, field in cls.__fields__.items()
-            if field_name != "id" and sqlmodel_field_has_primary_key(field)
-        ]
-        if non_default_primary_key_fields:
-            cls.__fields__.pop("id", None)
-
-        super().__init_subclass__()
-
-    @classmethod
-    def _dict_recursive(cls, value: Any):
-        """Recursively serialize the relationship object(s).
-
-        Args:
-            value: The value to serialize.
+        Attempts to connect to the database and execute a simple query to verify connectivity.
 
         Returns:
-            The serialized value.
+            The status of the database connection.
         """
-        if hasattr(value, "dict"):
-            return value.dict()
-        if isinstance(value, list):
-            return [cls._dict_recursive(item) for item in value]
-        return value
+        status = True
+        try:
+            engine = get_engine()
+            with engine.connect() as connection:
+                connection.execute(sqlalchemy.text("SELECT 1"))
+        except sqlalchemy.exc.OperationalError:
+            status = False
+
+        return {"db": status}
 
-    def dict(self, **kwargs):
-        """Convert the object to a dictionary.
+    @serializer
+    def serialize_sqlmodel(m: sqlmodel.SQLModel) -> dict[str, Any]:
+        """Serialize a SQLModel object to a dictionary.
 
         Args:
-            kwargs: Ignored but needed for compatibility.
+            m: The SQLModel object to serialize.
 
         Returns:
-            The object as a dictionary.
+            The serialized object as a dictionary.
         """
-        base_fields = {name: getattr(self, name) for name in self.__fields__}
+        base_fields = m.model_dump()
         relationships = {}
         # SQLModel relationships do not appear in __fields__, but should be included if present.
-        for name in self.__sqlmodel_relationships__:
+        for name in m.__sqlmodel_relationships__:
             with suppress(
                 sqlalchemy.orm.exc.DetachedInstanceError  # This happens when the relationship was never loaded and the session is closed.
             ):
-                relationships[name] = self._dict_recursive(getattr(self, name))
+                relationships[name] = getattr(m, name)
         return {
             **base_fields,
             **relationships,
         }
 
-    @staticmethod
-    def create_all():
-        """Create all the tables."""
-        engine = get_engine()
-        ModelRegistry.get_metadata().create_all(engine)
+    class Model(sqlmodel.SQLModel):
+        """Base class to define a table in the database."""
 
-    @staticmethod
-    def get_db_engine():
-        """Get the database engine.
-
-        Returns:
-            The database engine.
-        """
-        return get_engine()
+        # The primary key for the table.
+        id: int | None = sqlmodel.Field(default=None, primary_key=True)
 
-    @staticmethod
-    def _alembic_config():
-        """Get the alembic configuration and script_directory.
+        model_config = {  # pyright: ignore [reportAssignmentType]
+            "arbitrary_types_allowed": True,
+            "use_enum_values": True,
+            "extra": "allow",
+        }
 
-        Returns:
-            tuple of (config, script_directory)
-        """
-        config = alembic.config.Config(environment.ALEMBIC_CONFIG.get())
-        if not config.get_main_option("script_location"):
-            config.set_main_option("script_location", "version")
-        return config, alembic.script.ScriptDirectory.from_config(config)
+        @classmethod
+        def __pydantic_init_subclass__(cls):
+            """Drop the default primary key field if any primary key field is defined."""
+            non_default_primary_key_fields = [
+                field_name
+                for field_name, field_info in cls.model_fields.items()
+                if field_name != "id" and sqlmodel_field_has_primary_key(field_info)
+            ]
+            if non_default_primary_key_fields:
+                cls.model_fields.pop("id", None)
+                console.deprecate(
+                    feature_name="Overriding default primary key",
+                    reason=(
+                        "Register sqlmodel.SQLModel classes with `@rx.ModelRegistry.register`"
+                    ),
+                    deprecation_version="0.8.15",
+                    removal_version="0.9.0",
+                )
+            super().__pydantic_init_subclass__()
+
+        @staticmethod
+        def create_all():
+            """Create all the tables."""
+            engine = get_engine()
+            ModelRegistry.get_metadata().create_all(engine)
+
+        @staticmethod
+        def get_db_engine():
+            """Get the database engine.
+
+            Returns:
+                The database engine.
+            """
+            return get_engine()
+
+        @staticmethod
+        def _alembic_config():
+            """Get the alembic configuration and script_directory.
+
+            Returns:
+                tuple of (config, script_directory)
+            """
+            config = alembic.config.Config(environment.ALEMBIC_CONFIG.get())
+            if not config.get_main_option("script_location"):
+                config.set_main_option("script_location", "version")
+            return config, alembic.script.ScriptDirectory.from_config(config)
+
+        @staticmethod
+        def _alembic_render_item(
+            type_: str,
+            obj: Any,
+            autogen_context: alembic.autogenerate.api.AutogenContext,
+        ):
+            """Alembic render_item hook call.
 
-    @staticmethod
-    def _alembic_render_item(
-        type_: str,
-        obj: Any,
-        autogen_context: alembic.autogenerate.api.AutogenContext,
-    ):
-        """Alembic render_item hook call.
+            This method is called to provide python code for the given obj,
+            but currently it is only used to add `sqlmodel` to the import list
+            when generating migration scripts.
 
-        This method is called to provide python code for the given obj,
-        but currently it is only used to add `sqlmodel` to the import list
-        when generating migration scripts.
+            See https://alembic.sqlalchemy.org/en/latest/api/runtime.html
 
-        See https://alembic.sqlalchemy.org/en/latest/api/runtime.html
+            Args:
+                type_: One of "schema", "table", "column", "index",
+                    "unique_constraint", or "foreign_key_constraint".
+                obj: The object being rendered.
+                autogen_context: Shared AutogenContext passed to each render_item call.
 
-        Args:
-            type_: One of "schema", "table", "column", "index",
-                "unique_constraint", or "foreign_key_constraint".
-            obj: The object being rendered.
-            autogen_context: Shared AutogenContext passed to each render_item call.
+            Returns:
+                False - Indicating that the default rendering should be used.
+            """
+            autogen_context.imports.add("import sqlmodel")
+            return False
 
-        Returns:
-            False - Indicating that the default rendering should be used.
-        """
-        autogen_context.imports.add("import sqlmodel")
-        return False
-
-    @classmethod
-    def alembic_init(cls):
-        """Initialize alembic for the project."""
-        alembic.command.init(
-            config=alembic.config.Config(environment.ALEMBIC_CONFIG.get()),
-            directory=str(environment.ALEMBIC_CONFIG.get().parent / "alembic"),
-        )
+        @classmethod
+        def alembic_init(cls):
+            """Initialize alembic for the project."""
+            alembic.command.init(
+                config=alembic.config.Config(environment.ALEMBIC_CONFIG.get()),
+                directory=str(environment.ALEMBIC_CONFIG.get().parent / "alembic"),
+            )
 
-    @classmethod
-    def get_migration_history(cls):
-        """Get migration history with current database state.
+        @classmethod
+        def get_migration_history(cls):
+            """Get migration history with current database state.
+
+            Returns:
+                tuple: (current_revision, revisions_list) where revisions_list is in chronological order
+            """
+            # Get current revision from database
+            with cls.get_db_engine().connect() as connection:
+                context = MigrationContext.configure(connection)
+                current_rev = context.get_current_revision()
+
+            # Get all revisions from base to head
+            _, script_dir = cls._alembic_config()
+            revisions = list(script_dir.walk_revisions())
+            revisions.reverse()  # Reverse to get chronological order (base first)
+
+            return current_rev, revisions
+
+        @classmethod
+        def alembic_autogenerate(
+            cls,
+            connection: sqlalchemy.engine.Connection,
+            message: str | None = None,
+            write_migration_scripts: bool = True,
+        ) -> bool:
+            """Generate migration scripts for alembic-detectable changes.
+
+            Args:
+                connection: SQLAlchemy connection to use when detecting changes.
+                message: Human readable identifier describing the generated revision.
+                write_migration_scripts: If True, write autogenerated revisions to script directory.
+
+            Returns:
+                True when changes have been detected.
+            """
+            if not environment.ALEMBIC_CONFIG.get().exists():
+                return False
+
+            config, script_directory = cls._alembic_config()
+            revision_context = alembic.autogenerate.api.RevisionContext(
+                config=config,
+                script_directory=script_directory,
+                command_args=defaultdict(
+                    lambda: None,
+                    autogenerate=True,
+                    head="head",
+                    message=message,
+                ),
+            )
+            writer = alembic.autogenerate.rewriter.Rewriter()
 
-        Returns:
-            tuple: (current_revision, revisions_list) where revisions_list is in chronological order
-        """
-        # Get current revision from database
-        with cls.get_db_engine().connect() as connection:
-            context = MigrationContext.configure(connection)
-            current_rev = context.get_current_revision()
-
-        # Get all revisions from base to head
-        _, script_dir = cls._alembic_config()
-        revisions = list(script_dir.walk_revisions())
-        revisions.reverse()  # Reverse to get chronological order (base first)
-
-        return current_rev, revisions
-
-    @classmethod
-    def alembic_autogenerate(
-        cls,
-        connection: sqlalchemy.engine.Connection,
-        message: str | None = None,
-        write_migration_scripts: bool = True,
-    ) -> bool:
-        """Generate migration scripts for alembic-detectable changes.
+            @writer.rewrites(alembic.operations.ops.AddColumnOp)
+            def render_add_column_with_server_default(
+                context: MigrationContext,
+                revision: str | None,
+                op: Any,
+            ):
+                # Carry the sqlmodel default as server_default so that newly added
+                # columns get the desired default value in existing rows.
+                if op.column.default is not None and op.column.server_default is None:
+                    op.column.server_default = sqlalchemy.DefaultClause(
+                        sqlalchemy.sql.expression.literal(op.column.default.arg),
+                    )
+                return op
+
+            def run_autogenerate(rev: str, context: MigrationContext):
+                revision_context.run_autogenerate(rev, context)
+                return []
+
+            with alembic.runtime.environment.EnvironmentContext(
+                config=config,
+                script=script_directory,
+                fn=run_autogenerate,
+            ) as env:
+                env.configure(
+                    connection=connection,
+                    target_metadata=ModelRegistry.get_metadata(),
+                    render_item=cls._alembic_render_item,
+                    process_revision_directives=writer,
+                    compare_type=False,
+                    render_as_batch=True,  # for sqlite compatibility
+                )
+                env.run_migrations()
+            changes_detected = False
+            if revision_context.generated_revisions:
+                upgrade_ops = revision_context.generated_revisions[-1].upgrade_ops
+                if upgrade_ops is not None:
+                    changes_detected = bool(upgrade_ops.ops)
+            if changes_detected and write_migration_scripts:
+                # Must iterate the generator to actually write the scripts.
+                _ = tuple(revision_context.generate_scripts())
+            return changes_detected
+
+        @classmethod
+        def _alembic_upgrade(
+            cls,
+            connection: sqlalchemy.engine.Connection,
+            to_rev: str = "head",
+        ) -> None:
+            """Apply alembic migrations up to the given revision.
+
+            Args:
+                connection: SQLAlchemy connection to use when performing upgrade.
+                to_rev: Revision to migrate towards.
+            """
+            config, script_directory = cls._alembic_config()
+
+            def run_upgrade(rev: str, context: MigrationContext):
+                return script_directory._upgrade_revs(to_rev, rev)
+
+            with alembic.runtime.environment.EnvironmentContext(
+                config=config,
+                script=script_directory,
+                fn=run_upgrade,
+            ) as env:
+                env.configure(connection=connection)
+                env.run_migrations()
+
+        @classmethod
+        def migrate(cls, autogenerate: bool = False) -> bool | None:
+            """Execute alembic migrations for all sqlmodel Model classes.
+
+            If alembic is not installed or has not been initialized for the project,
+            then no action is performed.
+
+            If there are no revisions currently tracked by alembic, then
+            an initial revision will be created based on sqlmodel metadata.
+
+            If models in the app have changed in incompatible ways that alembic
+            cannot automatically generate revisions for, the app may not be able to
+            start up until migration scripts have been corrected by hand.
+
+            Args:
+                autogenerate: If True, generate migration script and use it to upgrade schema
+                    (otherwise, just bring the schema to current "head" revision).
+
+            Returns:
+                True - indicating the process was successful.
+                None - indicating the process was skipped.
+            """
+            if not environment.ALEMBIC_CONFIG.get().exists():
+                return None
+
+            with cls.get_db_engine().connect() as connection:
+                cls._alembic_upgrade(connection=connection)
+                if autogenerate:
+                    changes_detected = cls.alembic_autogenerate(connection=connection)
+                    if changes_detected:
+                        cls._alembic_upgrade(connection=connection)
+                connection.commit()
+            return True
+
+        @classmethod
+        def select(cls):
+            """Select rows from the table.
+
+            Returns:
+                The select statement.
+            """
+            return sqlmodel.select(cls)
+
+    ModelRegistry.register(Model)
+
+    def session(url: str | None = None) -> sqlmodel.Session:
+        """Get a sqlmodel session to interact with the database.
 
         Args:
-            connection: SQLAlchemy connection to use when detecting changes.
-            message: Human readable identifier describing the generated revision.
-            write_migration_scripts: If True, write autogenerated revisions to script directory.
+            url: The database url.
 
         Returns:
-            True when changes have been detected.
-        """
-        if not environment.ALEMBIC_CONFIG.get().exists():
-            return False
-
-        config, script_directory = cls._alembic_config()
-        revision_context = alembic.autogenerate.api.RevisionContext(
-            config=config,
-            script_directory=script_directory,
-            command_args=defaultdict(
-                lambda: None,
-                autogenerate=True,
-                head="head",
-                message=message,
-            ),
-        )
-        writer = alembic.autogenerate.rewriter.Rewriter()
-
-        @writer.rewrites(alembic.operations.ops.AddColumnOp)
-        def render_add_column_with_server_default(
-            context: MigrationContext,
-            revision: str | None,
-            op: Any,
-        ):
-            # Carry the sqlmodel default as server_default so that newly added
-            # columns get the desired default value in existing rows.
-            if op.column.default is not None and op.column.server_default is None:
-                op.column.server_default = sqlalchemy.DefaultClause(
-                    sqlalchemy.sql.expression.literal(op.column.default.arg),
-                )
-            return op
-
-        def run_autogenerate(rev: str, context: MigrationContext):
-            revision_context.run_autogenerate(rev, context)
-            return []
-
-        with alembic.runtime.environment.EnvironmentContext(
-            config=config,
-            script=script_directory,
-            fn=run_autogenerate,
-        ) as env:
-            env.configure(
-                connection=connection,
-                target_metadata=ModelRegistry.get_metadata(),
-                render_item=cls._alembic_render_item,
-                process_revision_directives=writer,
-                compare_type=False,
-                render_as_batch=True,  # for sqlite compatibility
-            )
-            env.run_migrations()
-        changes_detected = False
-        if revision_context.generated_revisions:
-            upgrade_ops = revision_context.generated_revisions[-1].upgrade_ops
-            if upgrade_ops is not None:
-                changes_detected = bool(upgrade_ops.ops)
-        if changes_detected and write_migration_scripts:
-            # Must iterate the generator to actually write the scripts.
-            _ = tuple(revision_context.generate_scripts())
-        return changes_detected
-
-    @classmethod
-    def _alembic_upgrade(
-        cls,
-        connection: sqlalchemy.engine.Connection,
-        to_rev: str = "head",
-    ) -> None:
-        """Apply alembic migrations up to the given revision.
-
-        Args:
-            connection: SQLAlchemy connection to use when performing upgrade.
-            to_rev: Revision to migrate towards.
+            A database session.
         """
-        config, script_directory = cls._alembic_config()
-
-        def run_upgrade(rev: str, context: MigrationContext):
-            return script_directory._upgrade_revs(to_rev, rev)
-
-        with alembic.runtime.environment.EnvironmentContext(
-            config=config,
-            script=script_directory,
-            fn=run_upgrade,
-        ) as env:
-            env.configure(connection=connection)
-            env.run_migrations()
+        return sqlmodel.Session(get_engine(url))
 
-    @classmethod
-    def migrate(cls, autogenerate: bool = False) -> bool | None:
-        """Execute alembic migrations for all sqlmodel Model classes.
+    def asession(url: str | None = None) -> AsyncSession:
+        """Get an async sqlmodel session to interact with the database.
 
-        If alembic is not installed or has not been initialized for the project,
-        then no action is performed.
+        async with rx.asession() as asession:
+            ...
 
-        If there are no revisions currently tracked by alembic, then
-        an initial revision will be created based on sqlmodel metadata.
-
-        If models in the app have changed in incompatible ways that alembic
-        cannot automatically generate revisions for, the app may not be able to
-        start up until migration scripts have been corrected by hand.
+        Most operations against the `asession` must be awaited.
 
         Args:
-            autogenerate: If True, generate migration script and use it to upgrade schema
-                (otherwise, just bring the schema to current "head" revision).
-
-        Returns:
-            True - indicating the process was successful.
-            None - indicating the process was skipped.
-        """
-        if not environment.ALEMBIC_CONFIG.get().exists():
-            return None
-
-        with cls.get_db_engine().connect() as connection:
-            cls._alembic_upgrade(connection=connection)
-            if autogenerate:
-                changes_detected = cls.alembic_autogenerate(connection=connection)
-                if changes_detected:
-                    cls._alembic_upgrade(connection=connection)
-            connection.commit()
-        return True
-
-    @classmethod
-    def select(cls):
-        """Select rows from the table.
+            url: The database url.
 
         Returns:
-            The select statement.
+            An async database session.
         """
-        return sqlmodel.select(cls)
-
-
-ModelRegistry.register(Model)
-
-
-def session(url: str | None = None) -> sqlmodel.Session:
-    """Get a sqlmodel session to interact with the database.
-
-    Args:
-        url: The database url.
-
-    Returns:
-        A database session.
-    """
-    return sqlmodel.Session(get_engine(url))
-
-
-def asession(url: str | None = None) -> AsyncSession:
-    """Get an async sqlmodel session to interact with the database.
-
-    async with rx.asession() as asession:
-        ...
-
-    Most operations against the `asession` must be awaited.
-
-    Args:
-        url: The database url.
-
-    Returns:
-        An async database session.
-    """
-    global _AsyncSessionLocal
-    if url not in _AsyncSessionLocal:
-        _AsyncSessionLocal[url] = sqlalchemy.ext.asyncio.async_sessionmaker(
-            bind=get_async_engine(url),
-            class_=AsyncSession,
-            expire_on_commit=False,
-            autocommit=False,
-            autoflush=False,
-        )
-    return _AsyncSessionLocal[url]()
-
-
-def sqla_session(url: str | None = None) -> sqlalchemy.orm.Session:
-    """Get a bare sqlalchemy session to interact with the database.
-
-    Args:
-        url: The database url.
+        global _AsyncSessionLocal
+        if url not in _AsyncSessionLocal:
+            _AsyncSessionLocal[url] = sqlalchemy.ext.asyncio.async_sessionmaker(
+                bind=get_async_engine(url),
+                class_=AsyncSession,
+                expire_on_commit=False,
+                autocommit=False,
+                autoflush=False,
+            )
+        return _AsyncSessionLocal[url]()
 
-    Returns:
-        A database session.
-    """
-    return sqlalchemy.orm.Session(get_engine(url))
+else:
+    get_db_status = _print_db_not_available
+    session = _print_db_not_available
+    asession = _print_db_not_available
+    Model = _ClassThatErrorsOnInit  # pyright: ignore [reportAssignmentType]