users.py 3.25 KB
Newer Older
1
from pydantic import BaseModel
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
from peewee import *
from playhouse.shortcuts import model_to_dict
4
5
6
from typing import List, Union, Optional
import time

Timothy J. Baek's avatar
Timothy J. Baek committed
7
8
9
from utils.utils import decode_token
from utils.misc import get_gravatar_url

Timothy J. Baek's avatar
Timothy J. Baek committed
10
from apps.web.internal.db import DB
11
12
from apps.web.models.chats import Chat

13
14
15
16
17
18

####################
# User DB Schema
####################


Timothy J. Baek's avatar
Timothy J. Baek committed
19
20
21
22
23
24
25
26
27
28
29
30
class User(Model):
    id = CharField(unique=True)
    name = CharField()
    email = CharField()
    role = CharField()
    profile_image_url = CharField()
    timestamp = DateField()

    class Meta:
        database = DB


31
32
33
34
class UserModel(BaseModel):
    id: str
    name: str
    email: str
Timothy J. Baek's avatar
Timothy J. Baek committed
35
36
    role: str = "pending"
    profile_image_url: str = "/user.png"
Timothy J. Baek's avatar
Timothy J. Baek committed
37
    timestamp: int  # timestamp in epoch
38
39
40
41
42
43
44


####################
# Forms
####################


Timothy J. Baek's avatar
Timothy J. Baek committed
45
46
47
48
49
class UserRoleUpdateForm(BaseModel):
    id: str
    role: str


50
51
52
class UsersTable:
    def __init__(self, db):
        self.db = db
Timothy J. Baek's avatar
Timothy J. Baek committed
53
        self.db.create_tables([User])
54
55

    def insert_new_user(
Timothy J. Baek's avatar
Timothy J. Baek committed
56
        self, id: str, name: str, email: str, role: str = "pending"
57
58
59
60
61
62
63
    ) -> Optional[UserModel]:
        user = UserModel(
            **{
                "id": id,
                "name": name,
                "email": email,
                "role": role,
Timothy J. Baek's avatar
Timothy J. Baek committed
64
                "profile_image_url": get_gravatar_url(email),
Timothy J. Baek's avatar
Timothy J. Baek committed
65
                "timestamp": int(time.time()),
66
67
            }
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
68
        result = User.create(**user.model_dump())
69
70
71
72
73
        if result:
            return user
        else:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
74
75
76
77
78
79
    def get_user_by_id(self, id: str) -> Optional[UserModel]:
        try:
            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None
80

Timothy J. Baek's avatar
Timothy J. Baek committed
81
82
83
84
85
    def get_user_by_email(self, email: str) -> Optional[UserModel]:
        try:
            user = User.get(User.email == email)
            return UserModel(**model_to_dict(user))
        except:
86
87
88
89
90
91
92
93
94
95
            return None

    def get_user_by_token(self, token: str) -> Optional[UserModel]:
        data = decode_token(token)

        if data != None and "email" in data:
            return self.get_user_by_email(data["email"])
        else:
            return None

96
    def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
97
        return [
Timothy J. Baek's avatar
Timothy J. Baek committed
98
99
            UserModel(**model_to_dict(user))
            for user in User.select().limit(limit).offset(skip)
100
101
        ]

102
    def get_num_users(self) -> Optional[int]:
Timothy J. Baek's avatar
Timothy J. Baek committed
103
        return User.select().count()
Timothy J. Baek's avatar
Timothy J. Baek committed
104
105

    def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
106
107
108
109
110
111
112
113
        try:
            query = User.update(role=role).where(User.id == id)
            query.execute()

            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None
Timothy J. Baek's avatar
Timothy J. Baek committed
114

115
116
117
    def delete_user_by_id(self, id: str) -> bool:
        try:
            # Delete User Chats
118
            result = Chat.delete_chats_by_user_id(id)
119

120
121
122
123
            if result:
                # Delete User
                query = User.delete().where(User.id == id)
                query.execute()  # Remove the rows, return number of rows removed.
124

125
126
127
                return True
            else:
                return False
128
129
130
        except:
            return False

131
132

Users = UsersTable(DB)