db.py

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

# 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": "5432",
                "user": os.getenv("DB_USER"),
                "password": os.getenv("DB_PASSWORD"),
                "database": "postgres",
                "min_size": 1, # Start with a small pool, e.g., 1-5 connections
                "max_size": 10, # Adjust based on expected load and Supabase limits. Common values: 10-50
                "timeout": 30, # Connection timeout in seconds [16]
                # "ssl": True, # Enable SSL if required by Supabase for production
                # "statement_cache_size": 0, # Optional: Keep for completeness if other issues arise, but primary fix is connection mode
                "max_queries": 50000, # Max queries before a connection is closed and replaced [15]
                "max_inactive_connection_lifetime": 300.0, # Max idle time before a connection is closed [15]
            },

            # Pass the custom connection class and disable the cache
            "connect_args": {
                "statement_cache_size": 0,
                "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()