How to define multiple, useful mixin classes, including messing with __table_args__?

[jfb-workhelix]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from datetime import UTC, datetime
from typing import Any
from uuid import UUID

from sqlalchemy import Column, DateTime, ForeignKey, func
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import RelationshipProperty, declarative_mixin, relationship
from sqlmodel import Field, SQLModel

DEFAULT_SCHEMA_NAME: str = "foobar"
SCHEMA_DUNDER_NAME: str = "__schema__"
DATASET_SCHEMA_NAME: str = DEFAULT_SCHEMA_NAME


class SQLModelWeirdnessError(Exception):
    """base."""


class StrangeTableArgsError(SQLModelWeirdnessError):
    """Is __table_args__ defined, and something that's not a tuple or a dict?."""


@declarative_mixin
class MyBaseModel(SQLModel):
    """We use schemata! Let's support this in a first class way."""

    @classmethod
    @declared_attr
    def __table_args__(cls: Any) -> tuple | dict:  # STFU mypy
        """Ensure that the __table_args__ structure has the appropriate 'schema' entry.

        Schema defaults to DEFAULT_SCHEMA_NAME.
        """
        schema_name = getattr(cls, SCHEMA_DUNDER_NAME, DEFAULT_SCHEMA_NAME)
        ta: tuple | dict = getattr(super(cls), "__table_args__", {})

        # Because this POS can be either a tuple or dict, we need to test types.
        # We live in a fallen world.

        # there are five possibilities:
        # 0. it's not defined on `super`, so use a default dict
        # 1. it's a dict: add the `schema` -> schema_name pair
        # 2. it's a zero length tuple?: return the default dict
        # 3. it's a tuple and the last element is a dict: update the last element
        # 4. otherwise, it's a tuple: add the default dict at the end

        if isinstance(ta, dict):
            return ta | {"schema": schema_name}

        if len(ta) == 0:
            return {"schema": schema_name}

        if isinstance(ta[-1], tuple):
            return *ta, {"schema": schema_name}

        if isinstance(ta[-1], dict):
            ta[-1]["schema"] = schema_name
            return ta

        raise StrangeTableArgsError(type(ta))


@declarative_mixin
class Timestamped:
    """Add our timestamp columns."""

    @classmethod
    @declared_attr
    def updated_at(cls: Any) -> Column:
        """Define `updated_at`."""
        return Column(
                name="updated_at",
                type=DateTime(timezone=True),
                default=datetime.now(UTC),
                nullable=False,
                server_default=func.now(),
                onupdate=func.now(),
        )

    @classmethod
    @declared_attr
    def created_at(cls: Any) -> Column:
        """Define `created_at`."""
        return Column(
                name="created_at",
                type=DateTime(timezone=True),
                default=datetime.now(UTC),
                nullable=False,
                server_default=func.now(),
        )


@declarative_mixin
class Datasetable:
    """Add the `dataset_id` column and foreign key relationship."""
    __config__ = None

    @classmethod
    @declared_attr
    def dataset_id(cls: Any) -> Column:
        """Define the DB column."""
        return Column(
                UUID,
                ForeignKey(
                        f"{DATASET_SCHEMA_NAME}.datasets.id",
                        onupdate="CASCADE",
                        deferrable=True,
                ),
                primary_key=True,
        )

    @classmethod
    @declared_attr
    def dataset(cls: Any) -> RelationshipProperty:
        """Define the relationship."""
        return relationship("Dataset")

class Foobar(SQLModel, Datasetable, table=True):
    __tablename__ = "foos"

    id: UUID = Field(..., primary_key=True)

Description

I'm trying to perform some cleanup on our models, and one of the things I'd like to do is to remove boilerplate, of which there is a lot. I've pared this down as much as possible for the example, but briefly; I have a fair number (~100) tables, and I'd really really really like to lift some boilerplate out of the model definitions. I thought, ok, I'll define a number of mixins, and I'll be good to go.

First, we heavily use PostgreSQL schemata, so I want to add first-class support for specifying the schema; something like,

class FooSchema:
  __schema__ = "foo"

class ThisShouldLiveInFoo (FooSchema, SQLModel, table=True):
  ...

The only way I could see to get this to work was by playing horrible games with the __table_args__ variable (see above). OK, horrible games are fine, if they're documented, good, that works. And it does! But then:

Second, there are a subset of tables that need to have certain columns defined on them: the standard updated_at/created_at pair, as well as ones that are used in our dataset abstraction (it's a way of grouping read-only data together for purposes of data lineage/auditing/&c., by joining subsidiary tables into a dataset entity). IOW, with the definitions above, you'd get:

# this class will inherit the `created_at/updated_at` columns
class Timestamped (Timestampable, SQLModel, table=True):
  ...

# this class will have a `dataset_id` and the requisite foreign key constraints, as well as a `dataset` relationship
class LivesInADataset (Datasetable, SQLModel, table=True):
  ...

However, I'm stymied at the first, because if I do use one of these Mixins, I get:

> poetry run python -c 'from sqlmodel_inheritance import Foobar'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/jfb/Projects/sqlmodel-inheritance/sqlmodel_inheritance/__init__.py", line 117, in <module>
    class Foobar(SQLModel, Datasetable, table=True):
  File "/Users/jfb/Projects/sqlmodel-inheritance/.venv/lib/python3.11/site-packages/sqlmodel/main.py", line 322, in __init__
    config = getattr(base, "__config__")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: type object 'Datasetable' has no attribute '__config__'

If I add the __config__ = None into Datasetable, ok, it works? Kind of? But when I run alembic, I can see that I do not get the dataset_id field in the DDL definition. At this point, magic stuff is happening that exceeds my Python understanding, so I thought I'd throw myself on the tender mercies of the internet.

Operating System

macOS

Operating System Details

No response

SQLModel Version

0.0.8

Python Version

3.11.2

Additional Context

No response