diff --git a/api/kiwi_vpn_api/db_new/device.py b/api/kiwi_vpn_api/db_new/device.py new file mode 100644 index 0000000..c493ebb --- /dev/null +++ b/api/kiwi_vpn_api/db_new/device.py @@ -0,0 +1,75 @@ +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() diff --git a/api/kiwi_vpn_api/db_new/user.py b/api/kiwi_vpn_api/db_new/user.py index 3d3eea7..7d0ec57 100644 --- a/api/kiwi_vpn_api/db_new/user.py +++ b/api/kiwi_vpn_api/db_new/user.py @@ -9,6 +9,7 @@ from sqlmodel import Field, Relationship, SQLModel from ..config import Config from .capabilities import Capability, UserCapability from .connection import Connection +from .device import Device class UserBase(SQLModel): @@ -33,6 +34,10 @@ class User(UserBase, table=True): }, ) + devices: list[Device] = Relationship( + back_populates="owner", + ) + @classmethod def create(cls, **kwargs) -> User | None: """