mirror of
https://github.com/PaiGramTeam/PamGram.git
synced 2024-11-17 04:18:03 +00:00
233e7ab58d
Co-authored-by: luoshuijs <luoshuijs@outlook.com> Co-authored-by: Karako <karakohear@gmail.com> Co-authored-by: xtaodada <xtao@xtaolink.cn>
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
import contextlib
|
|
from typing import Optional
|
|
|
|
from sqlalchemy.engine import URL
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from typing_extensions import Self
|
|
|
|
from core.base_service import BaseService
|
|
from core.config import ApplicationConfig
|
|
from core.sqlmodel.session import AsyncSession
|
|
|
|
__all__ = ("MySQL",)
|
|
|
|
|
|
class MySQL(BaseService.Dependence):
|
|
@classmethod
|
|
def from_config(cls, config: ApplicationConfig) -> Self:
|
|
return cls(**config.mysql.dict())
|
|
|
|
def __init__(
|
|
self,
|
|
host: Optional[str] = None,
|
|
port: Optional[int] = None,
|
|
username: Optional[str] = None,
|
|
password: Optional[str] = None,
|
|
database: Optional[str] = None,
|
|
):
|
|
self.database = database
|
|
self.password = password
|
|
self.username = username
|
|
self.port = port
|
|
self.host = host
|
|
self.url = URL.create(
|
|
"mysql+asyncmy",
|
|
username=self.username,
|
|
password=self.password,
|
|
host=self.host,
|
|
port=self.port,
|
|
database=self.database,
|
|
)
|
|
self.engine = create_async_engine(self.url)
|
|
self.Session = sessionmaker(bind=self.engine, class_=AsyncSession)
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def session(self) -> AsyncSession:
|
|
yield self.Session()
|
|
|
|
async def shutdown(self):
|
|
self.Session.close_all()
|