File size: 2,356 Bytes
9d4bd7c
183d33e
0d2d920
0d12e8a
 
 
 
 
 
 
 
 
9d4bd7c
60217c0
 
 
0d12e8a
9d4bd7c
183d33e
 
 
 
5d7715a
 
183d33e
 
 
 
0d12e8a
60217c0
 
0d12e8a
60217c0
 
183d33e
 
9d4bd7c
 
 
839f19d
 
 
 
 
 
 
9d4bd7c
 
 
839f19d
9d4bd7c
 
 
 
839f19d
 
 
 
 
 
 
 
 
9d4bd7c
 
 
839f19d
9d4bd7c
 
 
 
 
839f19d
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
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()