Skip to content

Commit

Permalink
feat: support mysql
Browse files Browse the repository at this point in the history
  • Loading branch information
krau committed Oct 13, 2023
1 parent 776f487 commit e055b46
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 23 deletions.
56 changes: 34 additions & 22 deletions kmua/models.py
Original file line number Diff line number Diff line change
@@ -1,28 +1,36 @@
from sqlalchemy import (
BLOB,
LargeBinary,
Boolean,
CheckConstraint,
Column,
DateTime,
Float,
ForeignKey,
Integer,
BigInteger,
String,
func,
)
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base
from kmua.dao._db import engine
from kmua.config import data_dir
from kmua.logger import logger


Base = declarative_base()


class UserChatAssociation(Base):
__tablename__ = "user_chat_association"
user_id = Column(Integer, ForeignKey("user_data.id"), primary_key=True)
chat_id = Column(Integer, ForeignKey("chat_data.id"), primary_key=True)
waifu_id = Column(Integer, ForeignKey("user_data.id"), default=None)
user_id = Column(
BigInteger, ForeignKey("user_data.id"), primary_key=True, autoincrement=False
)
chat_id = Column(
BigInteger, ForeignKey("chat_data.id"), primary_key=True, autoincrement=False
)
waifu_id = Column(
BigInteger, ForeignKey("user_data.id"), default=None, autoincrement=False
)
is_bot_admin = Column(Boolean, default=False)

created_at = Column(DateTime, default=func.now())
Expand All @@ -31,12 +39,12 @@ class UserChatAssociation(Base):

class UserData(Base):
__tablename__ = "user_data"
id = Column(Integer, primary_key=True, index=True)
username = Column(String)
full_name = Column(String, nullable=False)
avatar_small_blob = Column(BLOB, default=None)
avatar_big_blob = Column(BLOB, default=None)
avatar_big_id = Column(String, default=None)
id = Column(BigInteger, primary_key=True, index=True, autoincrement=False)
username = Column(String(32))
full_name = Column(String(128), nullable=False)
avatar_small_blob = Column(LargeBinary(65536), default=None)
avatar_big_blob = Column(LargeBinary(65536), default=None)
avatar_big_id = Column(String(128), default=None)

chats = relationship(
"ChatData",
Expand All @@ -49,7 +57,7 @@ class UserData(Base):
quotes = relationship("Quote", back_populates="user")

is_married = Column(Boolean, default=False)
married_waifu_id = Column(Integer, default=None)
married_waifu_id = Column(BigInteger, default=None)
waifu_mention = Column(Boolean, default=True)

is_bot = Column(Boolean, default=False)
Expand All @@ -68,17 +76,17 @@ class UserData(Base):

class ChatData(Base):
__tablename__ = "chat_data"
id = Column(Integer, primary_key=True, index=True)
id = Column(BigInteger, primary_key=True, index=True, autoincrement=False)
quote_probability = Column(Float, default=0.001)
title = Column(String, nullable=False)
title = Column(String(128), nullable=False)
members = relationship(
"UserData",
secondary=UserChatAssociation.__tablename__,
back_populates="chats",
primaryjoin="ChatData.id==UserChatAssociation.chat_id",
secondaryjoin="UserData.id==UserChatAssociation.user_id",
)
greet = Column(String, default=None)
greet = Column(String(4096), default=None)

quotes = relationship("Quote", back_populates="chat")
created_at = Column(DateTime, default=func.now())
Expand All @@ -87,13 +95,13 @@ class ChatData(Base):

class Quote(Base):
__tablename__ = "quotes"
chat_id = Column(Integer, ForeignKey("chat_data.id"))
message_id = Column(Integer, nullable=False)
link = Column(String, nullable=False, primary_key=True)
user_id = Column(Integer, ForeignKey("user_data.id"))
qer_id = Column(Integer) # 使用 q 的人
text = Column(String, nullable=True, default=None)
img = Column(String, nullable=True, default=None, comment="图片的 file id")
chat_id = Column(BigInteger, ForeignKey("chat_data.id"))
message_id = Column(BigInteger, nullable=False)
link = Column(String(128), nullable=False, primary_key=True)
user_id = Column(BigInteger, ForeignKey("user_data.id"))
qer_id = Column(BigInteger) # 使用 q 的人
text = Column(String(4096), nullable=True, default=None)
img = Column(String(128), nullable=True, default=None, comment="图片的 file id")
created_at = Column(DateTime, default=func.now())
updated_at = Column(DateTime, default=func.now(), onupdate=func.now())

Expand All @@ -104,4 +112,8 @@ class Quote(Base):
if not data_dir.exists():
data_dir.mkdir()

logger.debug("Connecting to database...")

Base.metadata.create_all(bind=engine)

logger.info("Success")
64 changes: 63 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ sniffio = "^1.3.0"
idna = "^3.4"
loguru = "^0.7.2"
greenlet = "^3.0.0"
mysql-connector-python = "^8.1.0"


[build-system]
Expand Down
2 changes: 2 additions & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@ httpcore==0.18.0 ; python_version >= "3.11" and python_version < "4.0"
httpx==0.25.0 ; python_version >= "3.11" and python_version < "4.0"
idna==3.4 ; python_version >= "3.11" and python_version < "4.0"
loguru==0.7.2 ; python_version >= "3.11" and python_version < "4.0"
mysql-connector-python==8.1.0 ; python_version >= "3.11" and python_version < "4.0"
pillow==9.5.0 ; python_version >= "3.11" and python_version < "4.0"
pilmoji==2.0.4 ; python_version >= "3.11" and python_version < "4.0"
protobuf==4.21.12 ; python_version >= "3.11" and python_version < "4.0"
psutil==5.9.5 ; python_version >= "3.11" and python_version < "4.0"
pypinyin==0.49.0 ; python_version >= "3.11" and python_version < "4"
python-telegram-bot[rate-limiter]==20.6 ; python_version >= "3.11" and python_version < "4.0"
Expand Down

0 comments on commit e055b46

Please sign in to comment.