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

"""
SQLAlchemy representation of database contents.
"""

from __future__ import annotations

from sqlalchemy import (Column, DateTime, ForeignKey, Integer, String,
                        UniqueConstraint)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship

ORMBaseModel = declarative_base()


class UserCapability(ORMBaseModel):
    __tablename__ = "user_capabilities"

    user_name = Column(
        String,
        ForeignKey("users.name"),
        primary_key=True,
        index=True,
    )
    capability = Column(String, primary_key=True)


class User(ORMBaseModel):
    __tablename__ = "users"

    name = Column(String, primary_key=True, index=True)
    password = Column(String, nullable=False)
    email = Column(String)

    country = Column(String(2))
    state = Column(String)
    city = Column(String)
    organization = Column(String)
    organizational_unit = Column(String)

    capabilities: list[UserCapability] = relationship(
        "UserCapability", lazy="joined", cascade="all, delete-orphan"
    )
    devices: list[Device] = relationship(
        "Device", lazy="select",  back_populates="owner"
    )

    @classmethod
    def from_db(
        cls,
        db: Session,
        name: str,
    ) -> User | None:
        """
        Load user from database by name.
        """

        return (
            db
            .query(cls)
            .filter(cls.name == name)
            .first()
        )


class Device(ORMBaseModel):
    __tablename__ = "devices"

    id = Column(Integer, primary_key=True, autoincrement=True)

    owner_name = Column(String, ForeignKey("users.name"))
    name = Column(String)
    type = Column(String)
    expiry = Column(DateTime)

    owner: User = relationship(
        "User", lazy="joined", back_populates="devices"
    )

    UniqueConstraint(
        owner_name,
        name,
    )