Possible way to mixin with DeclarativeBase

pamelafox · pamelafox · commit d89c21526482 · 2023-05-24T13:25:36.000-07:00
diff --git a/.pdm-python b/.pdm-python
@@ -0,0 +1 @@
+/Users/pamelafox/flask-sqlalchemy/.venv/bin/python
diff --git a/src/flask_sqlalchemy/extension.py b/src/flask_sqlalchemy/extension.py
@@ -503,7 +503,11 @@ def _make_declarative_base(
         .. versionchanged:: 2.3
             ``model`` can be an already created declarative model class.
         """
-        if not isinstance(model, sa.orm.DeclarativeMeta):
+        if (
+            not isinstance(model, sa.orm.DeclarativeMeta)
+            and not issubclass(model, sa.orm.DeclarativeBase)
+            and not issubclass(model, sa.orm.DeclarativeBaseNoMeta)
+        ):
             metadata = self._make_metadata(None)
             model = sa.orm.declarative_base(
                 metadata=metadata, cls=model, name="Model", metaclass=DefaultMeta
diff --git a/src/flask_sqlalchemy/model.py b/src/flask_sqlalchemy/model.py
@@ -182,7 +182,8 @@ def should_set_tablename(cls: type) -> bool:
     joined-table inheritance. If no primary key is found, the name will be unset.
     """
     if cls.__dict__.get("__abstract__", False) or not any(
-        isinstance(b, sa.orm.DeclarativeMeta) for b in cls.__mro__[1:]
+        (isinstance(b, sa.orm.DeclarativeMeta) or b is sa.orm.DeclarativeBase)
+        for b in cls.__mro__[1:]
     ):
         return False
 
@@ -196,7 +197,10 @@ def should_set_tablename(cls: type) -> bool:
         return not (
             base is cls
             or base.__dict__.get("__abstract__", False)
-            or not isinstance(base, sa.orm.DeclarativeMeta)
+            or not (
+                isinstance(base, sa.orm.DeclarativeMeta)
+                or base is sa.orm.DeclarativeBase
+            )
         )
 
     return True
@@ -212,3 +216,99 @@ class DefaultMeta(BindMetaMixin, NameMetaMixin, sa.orm.DeclarativeMeta):
     """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and
     ``__tablename__`` support.
     """
+
+
+class DefaultMixin:
+    """A mixin that provides Flask-SQLAlchemy default functionality:
+    * sets a model's ``__tablename__`` by converting the
+    ``CamelCase`` class name to ``snake_case``. A name is set for non-abstract models
+    that do not otherwise define ``__tablename__``. If a model does not define a primary
+    key, it will not generate a name or ``__table__``, for single-table inheritance.
+    * sets a model's ``metadata`` based on its ``__bind_key__``.
+    If the model sets ``metadata`` or ``__table__`` directly, ``__bind_key__`` is
+    ignored. If the ``metadata`` is the same as the parent model, it will not be set
+    directly on the child model.
+    * Provides a default repr based on the model's primary key.
+    """
+
+    __fsa__: SQLAlchemy
+    metadata: sa.MetaData
+    __tablename__: str
+    __table__: sa.Table
+
+    def __init_subclass__(cls, **kwargs):
+        if not ("metadata" in cls.__dict__ or "__table__" in cls.__dict__):
+            bind_key = getattr(cls, "__bind_key__", None)
+            parent_metadata = getattr(cls, "metadata", None)
+            metadata = cls.__fsa__._make_metadata(bind_key)
+
+            if metadata is not parent_metadata:
+                cls.metadata = metadata
+
+        if should_set_tablename(cls):
+            cls.__tablename__ = camel_to_snake_case(cls.__name__)
+
+        super().__init_subclass__(**kwargs)
+
+        # __table_cls__ has run. If no table was created, use the parent table.
+        if (
+            "__tablename__" not in cls.__dict__
+            and "__table__" in cls.__dict__
+            and cls.__dict__["__table__"] is None
+        ):
+            del cls.__table__
+
+    @classmethod
+    def __table_cls__(cls, *args: t.Any, **kwargs: t.Any) -> sa.Table | None:
+        """This is called by SQLAlchemy during mapper setup. It determines the final
+        table object that the model will use.
+
+        If no primary key is found, that indicates single-table inheritance, so no table
+        will be created and ``__tablename__`` will be unset.
+        """
+        schema = kwargs.get("schema")
+
+        if schema is None:
+            key = args[0]
+        else:
+            key = f"{schema}.{args[0]}"
+
+        # Check if a table with this name already exists. Allows reflected tables to be
+        # applied to models by name.
+        if key in cls.metadata.tables:
+            return sa.Table(*args, **kwargs)
+
+        # If a primary key is found, create a table for joined-table inheritance.
+        for arg in args:
+            if (isinstance(arg, sa.Column) and arg.primary_key) or isinstance(
+                arg, sa.PrimaryKeyConstraint
+            ):
+                return sa.Table(*args, **kwargs)
+
+        # If no base classes define a table, return one that's missing a primary key
+        # so SQLAlchemy shows the correct error.
+        for base in cls.__mro__[1:-1]:
+            if "__table__" in base.__dict__:
+                break
+        else:
+            return sa.Table(*args, **kwargs)
+
+        # Single-table inheritance, use the parent table name. __init__ will unset
+        # __table__ based on this.
+        if "__tablename__" in cls.__dict__:
+            del cls.__tablename__
+
+        return None
+
+    def __repr__(self) -> str:
+        state = sa.inspect(self)
+        assert state is not None
+
+        if state.transient:
+            pk = f"(transient {id(self)})"
+        elif state.pending:
+            pk = f"(pending {id(self)})"
+        else:
+            pk = ", ".join(map(str, state.identity))
+
+        return f"<{type(self).__name__} {pk}>"
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -6,9 +6,13 @@
 import sqlalchemy as sa
 import sqlalchemy.orm
 from flask import Flask
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
 
 from flask_sqlalchemy import SQLAlchemy
 from flask_sqlalchemy.model import DefaultMeta
+from flask_sqlalchemy.model import DefaultMixin
 from flask_sqlalchemy.model import Model
 
 
@@ -28,6 +32,45 @@ class CustomModel(Model):
     assert isinstance(db.Model, DefaultMeta)
 
 
+@pytest.mark.usefixtures("app_ctx")
+def test_custom_model_sqlalchemy20_class(app: Flask) -> None:
+    from sqlalchemy.orm.decl_api import DeclarativeAttributeIntercept
+
+    class Base(DeclarativeBase):
+        pass
+
+    db = SQLAlchemy(app, model_class=Base)
+
+    # Check the model class is instantiated with the correct metaclass
+    assert issubclass(db.Model, Base)
+    assert isinstance(db.Model, type)
+    assert isinstance(db.Model, DeclarativeAttributeIntercept)
+    # Check that additional attributes are added to the model class
+    assert db.Model.query_class is db.Query
+
+    # Now create a model that inherits from that declarative base
+    class Quiz(DefaultMixin, db.Model):
+        id: Mapped[int] = mapped_column(
+            db.Integer, primary_key=True, autoincrement=True
+        )
+        title: Mapped[str] = mapped_column(db.String(255), nullable=False)
+
+    assert Quiz.__tablename__ == "quiz"
+    assert isinstance(Quiz, DeclarativeAttributeIntercept)
+
+    db.create_all()
+    quiz = Quiz(title="Python trivia")
+    db.session.add(quiz)
+    db.session.commit()
+
+    # Check column types are correct
+    quiz_id: int = quiz.id
+    quiz_title: str = quiz.title
+    assert quiz_id == 1
+    assert quiz_title == "Python trivia"
+    assert repr(quiz) == f"<Quiz {quiz.id}>"
+
+
 @pytest.mark.usefixtures("app_ctx")
 @pytest.mark.parametrize("base", [Model, object])
 def test_custom_declarative_class(app: Flask, base: t.Any) -> None:

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+/Users/pamelafox/flask-sqlalchemy/.venv/bin/python`