malepati/custom_template_working

1from collections.abc import AsyncGenerator2import os3from sqlalchemy.ext.asyncio import (4    AsyncEngine,5    create_async_engine,6    async_sessionmaker,7    AsyncSession,8)9from sqlmodel import SQLModel10 11from models.sql.async_presentation_generation_status import (12    AsyncPresentationGenerationTaskModel,13)14from models.sql.image_asset import ImageAsset15from models.sql.key_value import KeyValueSqlModel16from models.sql.ollama_pull_status import OllamaPullStatus17from models.sql.presentation import PresentationModel18from models.sql.slide import SlideModel19from models.sql.presentation_layout_code import PresentationLayoutCodeModel20from models.sql.template import TemplateModel21from models.sql.webhook_subscription import WebhookSubscription22from utils.db_utils import get_database_url_and_connect_args23 24 25database_url, connect_args = get_database_url_and_connect_args()26 27sql_engine: AsyncEngine = create_async_engine(database_url, connect_args=connect_args)28async_session_maker = async_sessionmaker(sql_engine, expire_on_commit=False)29 30 31async def get_async_session() -> AsyncGenerator[AsyncSession, None]:32    async with async_session_maker() as session:33        yield session34 35 36# Container DB (Lives inside the container)37# Container DB (Lives inside the container)38from utils.get_env import get_app_data_directory_env39# Ensure the directory exists40os.makedirs(get_app_data_directory_env(), exist_ok=True)41 42container_db_url = f"sqlite+aiosqlite:///{os.path.join(get_app_data_directory_env(), 'container.db')}"43container_db_engine: AsyncEngine = create_async_engine(44    container_db_url, connect_args={"check_same_thread": False}45)46container_db_async_session_maker = async_sessionmaker(47    container_db_engine, expire_on_commit=False48)49 50 51async def get_container_db_async_session() -> AsyncGenerator[AsyncSession, None]:52    async with container_db_async_session_maker() as session:53        yield session54 55 56# Create Database and Tables57async def create_db_and_tables():58    async with sql_engine.begin() as conn:59        await conn.run_sync(60            lambda sync_conn: SQLModel.metadata.create_all(61                sync_conn,62                tables=[63                    PresentationModel.__table__,64                    SlideModel.__table__,65                    KeyValueSqlModel.__table__,66                    ImageAsset.__table__,67                    PresentationLayoutCodeModel.__table__,68                    TemplateModel.__table__,69                    WebhookSubscription.__table__,70                    AsyncPresentationGenerationTaskModel.__table__,71                ],72            )73        )74 75    async with container_db_engine.begin() as conn:76        await conn.run_sync(77            lambda sync_conn: SQLModel.metadata.create_all(78                sync_conn,79                tables=[OllamaPullStatus.__table__],80            )81        )82