-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdiff_app_db.py
More file actions
95 lines (90 loc) · 3.37 KB
/
diff_app_db.py
File metadata and controls
95 lines (90 loc) · 3.37 KB
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
diff --git a/app/db.py b/app/db.py
index b7b00e9d..1d299bce 100644
--- a/app/db.py
+++ b/app/db.py
@@ -3,10 +3,12 @@ Minimal database configuration for Supabase PostgreSQL connectivity
"""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session, DeclarativeBase
-from contextlib import contextmanager
+from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
+from contextlib import contextmanager, asynccontextmanager
import os
import logging
import time
+from typing import AsyncGenerator
# Configure logging
logging.basicConfig(level=logging.INFO)
@@ -30,7 +32,19 @@ if DATABASE_URL and "sslmode=" not in DATABASE_URL and not DATABASE_URL.startswi
else:
DATABASE_URL += "?sslmode=require"
-# Create engine with production-ready settings and timeout protection
+# Helper to build async database URL
+def _build_async_database_url(url: str) -> str:
+ if url.startswith("sqlite+aiosqlite://") or url.startswith("postgresql+asyncpg://"):
+ return url
+ if url.startswith("sqlite://"):
+ return url.replace("sqlite://", "sqlite+aiosqlite://", 1)
+ if url.startswith("postgresql://"):
+ return url.replace("postgresql://", "postgresql+asyncpg://", 1)
+ return url
+
+ASYNC_DATABASE_URL = _build_async_database_url(DATABASE_URL)
+
+# Create sync engine with production-ready settings and timeout protection
engine_kwargs = {
"pool_pre_ping": True, # Verify connections before use
"pool_recycle": 300, # Recycle connections every 5 minutes
@@ -63,7 +77,20 @@ except Exception as e:
logger.error(f"Failed to create database engine: {e}")
raise
+# Async engine (mirrors sync configuration for async usage)
+async_engine_kwargs = {
+ "pool_pre_ping": True,
+ "echo": engine_kwargs.get("echo", False)
+}
+
+try:
+ async_engine = create_async_engine(ASYNC_DATABASE_URL, **async_engine_kwargs)
+except Exception as e:
+ logger.error(f"Failed to create async database engine: {e}")
+ raise
+
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+AsyncSessionLocal = async_sessionmaker(async_engine, expire_on_commit=False)
# Modern SQLAlchemy 2.0 style Base class
class Base(DeclarativeBase):
@@ -121,6 +148,27 @@ def get_db_session(timeout: int = 30):
finally:
db.close()
+@asynccontextmanager
+async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+ """Async context manager for database sessions"""
+ async with AsyncSessionLocal() as session:
+ try:
+ yield session
+ except Exception as e:
+ await session.rollback()
+ logger.error(f"Async database session error: {e}")
+ raise
+
+async def get_async_db() -> AsyncGenerator[AsyncSession, None]:
+ """FastAPI dependency for async database access"""
+ async with AsyncSessionLocal() as session:
+ try:
+ yield session
+ except Exception as e:
+ await session.rollback()
+ logger.error(f"Async database dependency error: {e}")
+ raise
+
def get_database_url() -> str:
"""
Get the database URL for external scripts
@@ -141,4 +189,4 @@ def test_connection():
return True
except Exception as e:
logger.error(f"Database connection test failed: {e}")
- return False
\ No newline at end of file
+ return False