Init commit

database/__init__.py

@@ -0,0 +1,2 @@
+__all__ = ["register", "get_all_users"]
+from .worker import register, get_all_users, set_group_settings, get_group

database/models.py

@@ -0,0 +1,26 @@
+from datetime import datetime
+
+from peewee import Model, BigIntegerField, CharField, DateTimeField
+
+from load import db
+
+
+
+class Users(Model):
+    user_id = BigIntegerField(null=False)
+    first_name = CharField(null=True)
+    last_name = CharField(null=True)
+    username = CharField(null=True)
+    date = DateTimeField(default=datetime.now)
+
+    class Meta:
+        database = db
+
+
+class Chat(Model):
+    chat_id = BigIntegerField(null=False)
+    group = CharField(null=False)
+
+    class Meta:
+        database = db
+

database/worker.py

@@ -0,0 +1,52 @@
+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
+