LGTM Devlog 29: ORM for Firestore and __init__subclass__ dunders and metaclasses
Tags: python, devjournal
Published: 2021-02-01 11:25:56

Uh oh, I refactored again. I wasn't quite happy with how messy it was, how each Game, User, and Quest entity was handling their data, there was a lot of duplicate code when it came to communicating with the database, and when that happened.

For example, the game object prior to the latest change looks like this:

```py
class Game:
    @classmethod
    def from_user(cls, user: User) -> Union[Game, NoGameType]:
        """ Create a game from a user """
        key = cls.make_key(user)
        game = cls(key)
        game.user = user

        docs = db.collection("game").where("user_key", "==", user.key).stream()
        for _ in docs:
            return game
        return NoGame

    @classmethod
    def new_from_fork(cls, user: User, fork_url: str) -> Game:
        """ Save game with fork """

        if not fork_url:
            raise ValueError("Fork can't be blank")

        key = cls.make_key(user)
        game = cls(key)
        game.user = user

        game_doc = db.collection("game").document(game.key).get()
        if game_doc.exists:
            game_doc.reference.set(
                {
                    "fork_url": fork_url,
                    "user_uid": user.uid,
                    "user_key": user.key,
                },
                merge=True,
            )

        else:
            game_doc.reference.set(
                {
                    "fork_url": fork_url,
                    "user_uid": user.uid,
                    "user_key": user.key,
                    "joined": firestore.SERVER_TIMESTAMP,
                }
            )

        return game

    @staticmethod
    def make_key(user: User) -> str:
        """ Game's key ARE user key due to 1:1 relationship """
        return user.key

    key: str
    user: Union[User, NoUserType]

    def __init__(self, key: str):
        self.key = key
        self.user = NoUser

    def assign_to_uid(self, uid: str) -> None:
        """ Assign a user to this game """
        doc = db.collection("game").document(self.key).get()
        if doc.exists:
            doc.reference.set({"user_uid": uid}, merge=True)

    def start_first_quest(self) -> None:
        """ Create starting quest if not exist """
        QuestClass = Quest.get_first_quest()
        quest = QuestClass.from_game(self)
        quest.execute_stages(TickType.FULL)
        quest.save()

    def __repr__(self):
        return f"{self.__class__.__name__}(key={self.key})"
```

We have two functions for creating `Game` objects: `new_from_fork()` and `from_user()`. There's also the method `assign_to_uid()` which writes a single property to the database.

This seems ok, but there's a lot of duplicate code in the `User` object, and `Quest` objects.

So I decided to warp a lot of this duplicate function into an ORM base class so that `User`, `Quest` and `Game` can all inherit these functions that will save and restore themselves from the database.

After the change, the new `Game` object looks like this:

```py
class Game(Orm, collection="game", parent_orm=User):
    data: GameData
    storage_model = GameData

    @classmethod
    def from_user(cls, user: User) -> Game:
        key = cls.make_key(user)
        game = cls(key)
        game.parent_key = user.key
        return game

    @staticmethod
    def make_key(user: User) -> str:
        """ Game's key ARE user key due to 1:1 relationship """
        return user.key

    def set_fork_url(self, fork_url: str) -> None:
        self.data.fork_url = fork_url
```

Much cleaner! The first thing you'll notice is the line:

```py
class Game(Orm, collection="game", parent_orm=User):
```

This is a neat feature in Python allowing customisation of metaclasses and is described in [PEP 487](https://www.python.org/dev/peps/pep-0487/). It's powered by the `__init_subclass__` dunder, which runs when the class is initialised (compared to `__init__` when the _instance_ is initialised), which in this case looks like this (from the `Orm` base class):

```py
    def __init_subclass__(
        cls, collection: str, parent_orm: Union[Type[Orm], NoParentType] = NoParent
    ):
        """ Set collection and parent """
        cls.collection = collection
        cls.parent_orm = parent_orm
        cls.col_ref = db.collection(collection)
```

So, what is happening here is when `Game` class is initialized, it is setting the `collection` class variable, and the `parent_orm` class variable. While we could have used `__init__` to set the collection and parent_orm, this method provides a nice separation of concerns between the `Game` class, whose class properties relate to the collection of games; and a specific instance of the `Game` class which relate to a specific entry in this collection.  Setting `col_ref` in `__init_subclass__` allows us to do `Game.col_ref` without having to instantiate a Game object (which we shouldn't need to do when referencing all Games).

The other changes is removing `new_from_fork()` and replacing with a short `set_fork_url()` setter. The reason for this is I realised both `new_from_fork()` and `from_user()` require a `user` argument, so there's a duplication here.  `assign_to_uid()` can now use the `Orm` base class's `parent_key` property; and the quest stuff has been moved to a new `QuestPage` object instead, where it's more relevant.


# The ORM
The Underlying ORM base class's job is to map data between Firestore objects and Python objects. It has all the shared code that `Game`, `User` and new `QuestPage` object needs. The top half of it looks like this:

```py

class Orm(ABC):
    """ ORM base class links stuff together """

    @property
    @abstractmethod
    def storage_model(cls) -> Type[BaseModel]:
        """ Storage model """
        return NotImplemented

    collection: ClassVar[str]
    parent_orm: ClassVar[Union[Type[Orm], NoParentType]]
    col_ref: CollectionReference

    def __init_subclass__(
        cls, collection: str, parent_orm: Union[Type[Orm], NoParentType] = NoParent
    ):
        """ Set collection and parent """
        cls.collection = collection
        cls.parent_orm = parent_orm
        cls.col_ref = db.collection(collection)

    key: Union[str, NoKeyType]
    parent_key: Union[str, NoKeyType]
    data: BaseModel

    def __init__(self, key: Union[str, NoKeyType] = NoKey):
        self.key = key
        self.data = self.storage_model()
        self.parent_key = NoKey

    @property
    def parent(self) -> Union[Orm, OrmNotFoundType]:
        if self.parent_orm is not NoParent and self.parent_key is not NoKey:
            return self.parent_orm(self.parent_key)

        return OrmNotFound

    ...
```

With the bottom half related to the mechanics of saving and loading data from the database.

I've described the `__init_subclass__` already, the rest of this ABC enforces that concrete implementations should supply their own `storage_model`, a Pydantic model which the base class's methods will use for creating and restoring data firestore. It also has a property called `parent` which will return an instance of the parent object linked by a `parent_key` property. For example if we had an instantiated `Game` object, we could fetch it's user by doing:

```py
game = Game(key)
game.load()

user = game.parent
```

Because `Game` objects know that their parent ORM object is `User`, and `Game` objects also store a `parent_key` property. This is enough detail to return a User object.

# QuestPages
I've done some big refactoring to the former `Quest` object. It's now been split into two: `Quest` which contains everything related to running the Quest, and serves as a base class to concrete implementations of individual Quests; and `QuestPage` which is subclassed from `Orm` and relates only to saving and loading quests from the database. I imagine `QuestPage` as a page in a quest log, which contains the "save data" relating to a `Quest` which is the actual implementation of the quest.

There's a bit of composition going on, as `QuestPage` has one instance of `quest` as a property.

```py

class QuestPage(Orm, collection="quest", parent_orm=Game):
    data: QuestData
    storage_model = QuestData
    quest: Quest

    @staticmethod
    def make_key(game: Game, quest_name: str) -> str:
        if not quest_name:
            raise ValueError("quest_name must be valid")
        return f"{game.key}:{quest_name}"

    @classmethod
    def from_game_get_first_quest(cls, game: Game) -> QuestPage:
        return cls.from_game_get_quest(game, FIRST_QUEST_NAME)

    @classmethod
    def from_game_get_quest(cls, game: Game, quest_name: str) -> QuestPage:
        key = cls.make_key(game, quest_name)
        quest = cls(key, quest_name)
        return quest

    @classmethod
    def iterate_all(cls) -> Generator[QuestPage, None, None]:
        """ Iterate over all quests, the generator yields loaded quest_pages """
        docs = cls.col_ref.where("complete", "!=", True).stream()
        for doc in docs:
            data = cls.storage_model.parse_obj(doc.to_dict())
            quest_page = cls(doc.id, data.quest_name)
            quest_page.data = data
            quest_page.quest.load_raw(data.version, data.serialized_data)
            yield quest_page

    def __init__(self, key: str, quest_name):
        super().__init__(key)
        self.quest = Quest.from_name(quest_name, self)
        self.data.quest_name = quest_name

    def load(self) -> None:
        """ Additionally parse the quest storage """
        super().load()
        if isinstance(self.quest, Quest):
            self.quest.load_raw(self.data.version, self.data.serialized_data)

    def save(self) -> None:
        """ Additionally parse out the quest storage """
        if isinstance(self.quest, Quest):
            self.data.serialized_data = self.quest.save_raw()
            self.data.version = str(self.quest.version)
        super().save()
    ...
```

This extends the `__init__()`, `save()`, and `load()` of the ORM base class in order to also save and load the quest data model distinct from the quest_page data model. The reason for this split is so that the quests themselves have their own variable space for quest-related data, while all the metadata - the version number, the completed quest list, etc. are store separately.

___
Together with the other changes, this refactor, which still hasn't added any new functionality, has cleaned up the code once again. Hopefully i don't have to do too many more of these