53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
|
from typing import Union, List
|
||
|
|
|
||
|
|
from .models import Users, Chat, db
|
||
|
|
|
||
|
|
|
||
|
|
db.create_tables([Users, Chat])
|
||
|
|
|
||
|
|
|
||
|
|
def register(
|
||
|
|
user_id: int,
|
||
|
|
username: Union[str, None],
|
||
|
|
first_name: str,
|
||
|
|
last_name: Union[str, None]
|
||
|
|
) -> None:
|
||
|
|
if not Users.select().where(Users.user_id == user_id).exists():
|
||
|
|
Users.create(
|
||
|
|
user_id=user_id,
|
||
|
|
username=username,
|
||
|
|
first_name=first_name,
|
||
|
|
last_name=last_name
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
(Users.update(
|
||
|
|
username=username,
|
||
|
|
first_name=first_name,
|
||
|
|
last_name=last_name
|
||
|
|
)
|
||
|
|
.where(Users.user_id == user_id).execute())
|
||
|
|
|
||
|
|
|
||
|
|
def get_all_users() -> List[int]:
|
||
|
|
usr = []
|
||
|
|
for user in Users.select():
|
||
|
|
usr.append(user.user_id)
|
||
|
|
return usr
|
||
|
|
|
||
|
|
|
||
|
|
def set_group_settings(chat_id: int, group: Union[int, str]) -> None:
|
||
|
|
if Chat.select().where(Chat.chat_id == chat_id).exists():
|
||
|
|
Chat.update(group=group).where(Chat.chat_id == chat_id).execute()
|
||
|
|
else:
|
||
|
|
Chat.create(
|
||
|
|
chat_id=chat_id,
|
||
|
|
group=group
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_group(chat_id: int) -> Union[str, None]:
|
||
|
|
if Chat.select().where(Chat.chat_id == chat_id).exists():
|
||
|
|
return Chat.get(Chat.chat_id == chat_id).group
|
||
|
|
return None
|
||
|
|
|