users.py 1.88 KB
Newer Older
1
2
3
4
5
from pydantic import BaseModel
from typing import List, Union, Optional
from pymongo import ReturnDocument
import time

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

9
10
11
12
13
14
15
16
17
18
19
from config import DB

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


class UserModel(BaseModel):
    id: str
    name: str
    email: str
Timothy J. Baek's avatar
Timothy J. Baek committed
20
21
    role: str = "pending"
    profile_image_url: str = "/user.png"
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    created_at: int  # timestamp in epoch


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


class UsersTable:
    def __init__(self, db):
        self.db = db
        self.table = db.users

    def insert_new_user(
Timothy J. Baek's avatar
Timothy J. Baek committed
36
        self, id: str, name: str, email: str, role: str = "pending"
37
38
39
40
41
42
43
    ) -> Optional[UserModel]:
        user = UserModel(
            **{
                "id": id,
                "name": name,
                "email": email,
                "role": role,
Timothy J. Baek's avatar
Timothy J. Baek committed
44
                "profile_image_url": get_gravatar_url(email),
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
                "created_at": int(time.time()),
            }
        )
        result = self.table.insert_one(user.model_dump())

        if result:
            return user
        else:
            return None

    def get_user_by_email(self, email: str) -> Optional[UserModel]:
        user = self.table.find_one({"email": email}, {"_id": False})

        if user:
            return UserModel(**user)
        else:
            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

    def get_users(self, skip: int = 0, limit: int = 50) -> Optional[UserModel]:
        return [
            UserModel(**user)
            for user in list(self.table.find({}, {"_id": False}))
            .skip(skip)
            .limit(limit)
        ]


Users = UsersTable(DB)