28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
import os
|
|
|
|
import ray
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlmodel import SQLModel
|
|
|
|
from pretor.core.database.module.user import AuthDatabase
|
|
from pretor.core.database.module.provider import ProviderDatabase
|
|
|
|
@ray.remote
|
|
class PostgresDatabase:
|
|
def __init__(self):
|
|
user = os.environ.get('POSTGRES_USER')
|
|
password = os.environ.get('POSTGRES_PASSWORD')
|
|
host = os.environ.get('POSTGRES_HOST')
|
|
port = os.environ.get('POSTGRES_PORT')
|
|
database = os.environ.get('POSTGRES_DB')
|
|
database_url = f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{database}"
|
|
self.async_engine = create_async_engine(database_url, echo=True)
|
|
self.async_session_maker = sessionmaker(self.async_engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
self.auth_database = AuthDatabase(self.async_session_maker)
|
|
self.provider_database = ProviderDatabase(self.async_session_maker)
|
|
|
|
async def init_db(self) -> None:
|
|
async with self.async_engine.begin() as conn:
|
|
await conn.run_sync(SQLModel.metadata.create_all) |