users.py 5.27 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
from typing import List, Union, Optional
import time
Timothy J. Baek's avatar
Timothy J. Baek committed
6
7
from utils.misc import get_gravatar_url

Timothy J. Baek's avatar
Timothy J. Baek committed
8
from apps.web.internal.db import DB
Timothy J. Baek's avatar
Timothy J. Baek committed
9
from apps.web.models.chats import Chats
10

11
12
13
14
15
####################
# User DB Schema
####################


Timothy J. Baek's avatar
Timothy J. Baek committed
16
17
18
19
20
class User(Model):
    id = CharField(unique=True)
    name = CharField()
    email = CharField()
    role = CharField()
21
    profile_image_url = TextField()
Timothy J. Baek's avatar
Timothy J. Baek committed
22
23
24
25
26

    last_active_at = BigIntegerField()
    updated_at = BigIntegerField()
    created_at = BigIntegerField()

27
    api_key = CharField(null=True, unique=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
30
31
32

    class Meta:
        database = DB


33
34
35
36
class UserModel(BaseModel):
    id: str
    name: str
    email: str
Timothy J. Baek's avatar
Timothy J. Baek committed
37
    role: str = "pending"
38
    profile_image_url: str
Timothy J. Baek's avatar
Timothy J. Baek committed
39
40
41
42
43

    last_active_at: int  # timestamp in epoch
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch

44
    api_key: Optional[str] = None
45
46
47
48
49
50
51


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


Timothy J. Baek's avatar
Timothy J. Baek committed
52
53
54
55
56
class UserRoleUpdateForm(BaseModel):
    id: str
    role: str


Timothy J. Baek's avatar
Timothy J. Baek committed
57
58
59
60
61
62
class UserUpdateForm(BaseModel):
    name: str
    email: str
    profile_image_url: str
    password: Optional[str] = None

63

Timothy J. Baek's avatar
Timothy J. Baek committed
64
class UsersTable:
65
66
    def __init__(self, db):
        self.db = db
Timothy J. Baek's avatar
Timothy J. Baek committed
67
        self.db.create_tables([User])
68

Timothy J. Baek's avatar
Timothy J. Baek committed
69
    def insert_new_user(
Danny Liu's avatar
Danny Liu committed
70
71
72
73
        self,
        id: str,
        name: str,
        email: str,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
74
        profile_image_url: str = "/user.png",
Danny Liu's avatar
Danny Liu committed
75
        role: str = "pending",
Timothy J. Baek's avatar
Timothy J. Baek committed
76
    ) -> Optional[UserModel]:
77
78
79
80
81
82
        user = UserModel(
            **{
                "id": id,
                "name": name,
                "email": email,
                "role": role,
83
                "profile_image_url": profile_image_url,
Timothy J. Baek's avatar
Timothy J. Baek committed
84
85
86
                "last_active_at": int(time.time()),
                "created_at": int(time.time()),
                "updated_at": int(time.time()),
Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
            }
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
89
        result = User.create(**user.model_dump())
90
91
92
93
94
        if result:
            return user
        else:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
95
96
97
98
99
100
    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
101

102
103
104
105
106
107
108
    def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
        try:
            user = User.get(User.api_key == api_key)
            return UserModel(**model_to_dict(user))
        except:
            return None

Timothy J. Baek's avatar
Timothy J. Baek committed
109
110
111
112
113
    def get_user_by_email(self, email: str) -> Optional[UserModel]:
        try:
            user = User.get(User.email == email)
            return UserModel(**model_to_dict(user))
        except:
114
115
            return None

116
    def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
117
        return [
Timothy J. Baek's avatar
Timothy J. Baek committed
118
            UserModel(**model_to_dict(user))
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
            for user in User.select()
            # .limit(limit).offset(skip)
121
122
        ]

123
    def get_num_users(self) -> Optional[int]:
Timothy J. Baek's avatar
Timothy J. Baek committed
124
        return User.select().count()
Timothy J. Baek's avatar
Timothy J. Baek committed
125

Timothy J. Baek's avatar
Timothy J. Baek committed
126
    def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
127
128
129
130
131
132
133
134
        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
135

136
137
138
139
140
141
142
143
144
145
146
147
148
149
    def update_user_profile_image_url_by_id(
        self, id: str, profile_image_url: str
    ) -> Optional[UserModel]:
        try:
            query = User.update(profile_image_url=profile_image_url).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
150
151
152
153
154
155
156
157
158
159
    def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
        try:
            query = User.update(last_active_at=int(time.time())).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
160
161
162
163
164
165
166
167
168
169
    def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
        try:
            query = User.update(**updated).where(User.id == id)
            query.execute()

            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None

170
171
172
    def delete_user_by_id(self, id: str) -> bool:
        try:
            # Delete User Chats
Timothy J. Baek's avatar
Timothy J. Baek committed
173
            result = Chats.delete_chats_by_user_id(id)
174

175
176
177
            if result:
                # Delete User
                query = User.delete().where(User.id == id)
Timothy J. Baek's avatar
Timothy J. Baek committed
178
                query.execute()  # Remove the rows, return number of rows removed.
179

180
181
182
                return True
            else:
                return False
183
184
185
        except:
            return False

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
    def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
        try:
            query = User.update(api_key=api_key).where(User.id == id)
            result = query.execute()

            return True if result == 1 else False
        except:
            return False

    def get_user_api_key_by_id(self, id: str) -> Optional[str]:
        try:
            user = User.get(User.id == id)
            return user.api_key
        except:
            return None

202
203

Users = UsersTable(DB)