db.py

from tortoise import Tortoise
import os
import ssl
import uuid
from asyncpg import Connection

# 1. Create a custom connection class to generate unique IDs for every statement
class UniquePreparedStatementConnection(Connection):
    def _get_unique_id(self, prefix: str) -> str:
        # This forces every prepared statement to have a new, unique name,
        # preventing any possible collisions when using a connection pooler.
        return f"__asyncpg_{prefix}_{uuid.uuid4()}__"

# It's a good practice to create the SSL context outside the dict
ssl_context = ssl.create_default_context()

# 2. Update your TORTOISE_ORM configuration
TORTOISE_ORM = {
    "connections": {
        "default": {
            "engine": "tortoise.backends.asyncpg",
            "credentials": {
                "host": "aws-0-us-east-2.pooler.supabase.com",
                "port": "6543",
                "user": os.getenv("DB_USER"),
                "password": os.getenv("DB_PASSWORD"),
                "database": "postgres",
            },
            # Pass the custom connection class and disable the cache
            "connect_args": {
                "statement_cache_size": 0,
                "connection_class": UniquePreparedStatementConnection,
                "ssl": ssl_context
            }
        }
    },
    "apps": {
        "models": {
            "models": [
                "App.routers.stocks.models",
                "App.routers.tasks.models",
                "App.routers.utt.models",
                "App.routers.users.models",
                "App.routers.portfolio.models",
                "App.routers.bonds.models",
                "aerich.models",
            ],
            "default_connection": "default",
        }
    },
}

async def init_db():
    await Tortoise.init(
        TORTOISE_ORM  # db_url=DATABASE_URL,
        # modules={'models': [
        #     'App.routers.stocks.models',
        #     'App.routers.tasks.models',
        #     'App.routers.utt.models',
        #     'App.routers.users.models',
        #     'App.routers.portfolio.models',
        #     'App.routers.bonds.models'
        # ]}
    )
    await Tortoise.generate_schemas()


async def close_db():
    await Tortoise.close_connections()


async def clear_db():
    for model in Tortoise.apps.get("models").values():
        await model.all().delete()