Spaces:
Running
Running
File size: 1,672 Bytes
f50a610 a8a02ba ca23981 f50a610 ca23981 f50a610 | 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 | import os
import asyncpg
import redis.asyncio as redis_async
from typing import Optional
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://ciq:ciq_password@localhost:5432/conflictiq")
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
if DATABASE_URL.startswith("postgresql+asyncpg://"):
DATABASE_URL = DATABASE_URL.replace("postgresql+asyncpg://", "postgresql://", 1)
class Database:
pool: Optional[asyncpg.Pool] = None
redis: Optional[redis_async.Redis] = None
db = Database()
async def connect_db():
db_url = DATABASE_URL
if '?' in db_url:
db_url = db_url.split('?')[0]
ssl_context = None
if "neon.tech" in db_url:
import ssl
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
db.pool = await asyncpg.create_pool(db_url, ssl=ssl_context)
# Use Upstash Serverless REST if provided, otherwise Local Redis
upstash_url = os.getenv("UPSTASH_REDIS_REST_URL")
upstash_token = os.getenv("UPSTASH_REDIS_REST_TOKEN")
if upstash_url and upstash_token:
from upstash_redis.asyncio import Redis as UpstashRedis
db.redis = UpstashRedis(url=upstash_url, token=upstash_token)
else:
db.redis = redis_async.from_url(REDIS_URL, decode_responses=True)
async def disconnect_db():
if db.pool:
await db.pool.close()
if db.redis and hasattr(db.redis, 'close') and callable(db.redis.close):
await db.redis.close()
async def get_db_connection():
async with db.pool.acquire() as conn:
yield conn
def get_redis():
return db.redis
|