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

from __future__ import annotations

from datetime import datetime
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 DeviceBase(SQLModel):
    name: str
    type: str
    expiry: datetime | None


class DeviceCreate(DeviceBase):
    owner_name: str | None


class Device(DeviceBase, table=True):
    __table_args__ = (UniqueConstraint(
        "owner_name",
        "name",
    ),)

    id: int | None = Field(primary_key=True)
    owner_name: str | None = Field(foreign_key="user.name")

    owner: User = Relationship(
        back_populates="devices",
    )

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

        try:
            with Connection.session as db:
                device = cls.from_orm(DeviceCreate(**kwargs))

                db.add(device)
                db.commit()
                db.refresh(device)

                return device

        except IntegrityError:
            # device already existed
            return None

    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) -> bool:
        """
        Delete this device from the database.
        """

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