kiwi-vpn/api/kiwi_vpn_api/db/device.py

"""
Python representation of `device` table.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING

from sqlalchemy.exc import IntegrityError
from sqlmodel import Field, Relationship, SQLModel, UniqueConstraint

from .connection import Connection

if TYPE_CHECKING:
    from .user import User


class DeviceStatus(Enum):
    uncertified = "uncertified"
    pending = "pending"
    certified = "certified"

    def __repr__(self) -> str:
        return self.value


class DeviceBase(SQLModel):
    """
    Common to all representations of devices
    """

    name: str
    type: str


class DeviceCreate(DeviceBase):
    """
    Representation of a newly created device
    """


class DeviceRead(DeviceBase):
    """
    Representation of a device read via the API
    """

    id: int | None = Field(primary_key=True)
    status_str: str = Field(default=repr(DeviceStatus.uncertified))
    expiry: datetime | None = Field(default=None)
    owner_name: str = Field(foreign_key="user.name")

    @property
    def status(self) -> DeviceStatus:
        return DeviceStatus(self.status_str)

    # property setters don't work with sqlmodel
    def set_status(self, status: DeviceStatus) -> None:
        self.status_str = repr(status)


class Device(DeviceRead, table=True):
    """
    Representation of `device` table
    """

    __table_args__ = (UniqueConstraint(
        "owner_name",
        "name",
    ),)

    # no idea, but "User" (in quotes) doesn't work here
    # might be a future problem?
    owner: User = Relationship(
        back_populates="devices",
        sa_relationship_kwargs={
            "lazy": "joined",
        },
    )

    @classmethod
    def create(
        cls,
        *,
        owner: User,
        device: DeviceCreate,
    ) -> Device | None:
        """
        Create a new device in the database.
        """

        try:
            with Connection.session as db:
                new_device = cls.from_orm(device, {"owner_name": owner.name})

                db.add(new_device)
                db.commit()
                db.refresh(new_device)

                return new_device

        except IntegrityError:
            # device already existed
            return None

    @classmethod
    def get(cls, id: int) -> Device | None:
        """
        Load device from database by id.
        """

        with Connection.session as db:
            return db.get(cls, id)

    def update(self) -> None:
        """
        Update this device in the database.
        """

        with Connection.session as db:
            db.add(self)
            db.commit()
            db.refresh(self)

    def delete(self) -> None:
        """
        Delete this device from the database.
        """

        with Connection.session as db:
            db.delete(self)
            db.commit()