Spaces:
Running
Running
File size: 7,205 Bytes
2cb9dec ad86a2b 2cb9dec ad86a2b 2cb9dec ad86a2b 2cb9dec ad86a2b 2cb9dec |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 |
import logging
from typing import AsyncGenerator, List, Optional, Dict
from pydantic_settings import BaseSettings
from pydantic import PostgresDsn
import pg8000
from pg8000 import Connection
from pg8000.exceptions import DatabaseError as Pg8000DatabaseError
import asyncio
from contextlib import asynccontextmanager
from threading import Lock
from urllib.parse import urlparse
# Set up structured logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
class DatabaseSettings(BaseSettings):
db_url: PostgresDsn
pool_size: int = 5 # Default pool size is 5
class Config:
env_file = ".env"
# Custom database errors
class DatabaseError(Exception):
"""Base exception for database errors."""
pass
class ConnectionError(DatabaseError):
"""Exception raised when a database connection fails."""
pass
class PoolExhaustedError(DatabaseError):
"""Exception raised when the connection pool is exhausted."""
pass
class QueryExecutionError(DatabaseError):
"""Exception raised when a query execution fails."""
pass
class HealthCheckError(DatabaseError):
"""Exception raised when a health check fails."""
pass
class Database:
def __init__(self, db_url: PostgresDsn, pool_size: int):
self.db_url = db_url
self.pool_size = pool_size
self.pool: List[Connection] = []
self.lock = Lock()
async def connect(self) -> None:
"""Create a connection pool."""
try:
# Convert PostgresDsn to a string
db_url_str = str(self.db_url)
result = urlparse(db_url_str)
for _ in range(self.pool_size):
conn = pg8000.connect(
user=result.username,
password=result.password,
host=result.hostname,
port=result.port or 5432,
database=result.path.lstrip("/"),
)
self.pool.append(conn)
logger.info(
f"Database connection pool created with {self.pool_size} connections."
)
except Pg8000DatabaseError as e:
logger.error(f"Failed to create database connection pool: {e}")
raise ConnectionError("Failed to create database connection pool.") from e
async def disconnect(self) -> None:
"""Close all connections in the pool."""
with self.lock:
for conn in self.pool:
conn.close()
self.pool.clear()
logger.info("Database connection pool closed.")
@asynccontextmanager
async def get_connection(self) -> AsyncGenerator[Connection, None]:
"""Acquire a connection from the pool."""
with self.lock:
if not self.pool:
logger.error("Connection pool is exhausted.")
raise PoolExhaustedError("No available connections in the pool.")
conn = self.pool.pop()
try:
yield conn
except Pg8000DatabaseError as e:
logger.error(f"Connection error: {e}")
raise ConnectionError("Failed to use database connection.") from e
finally:
with self.lock:
self.pool.append(conn)
async def fetch(self, query: str, *args) -> Dict[str, List]:
"""
Execute a SELECT query and return the results as a dictionary of lists.
Args:
query (str): The SQL query to execute.
*args: Query parameters.
Returns:
Dict[str, List]: A dictionary where keys are column names and values are lists of column values.
Raises:
QueryExecutionError: If the query execution fails.
"""
try:
async with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(query, args)
rows = cursor.fetchall()
columns = [desc[0] for desc in cursor.description]
# Convert the list of dictionaries into a dictionary of lists
data_dict = {column: [] for column in columns}
for row in rows:
for i, value in enumerate(row):
data_dict[columns[i]].append(value)
return data_dict
except Pg8000DatabaseError as e:
logger.error(f"Query execution failed: {e}")
raise QueryExecutionError(f"Failed to execute query: {query}") from e
async def execute(self, query: str, *args) -> None:
"""
Execute an INSERT, UPDATE, or DELETE query.
Args:
query (str): The SQL query to execute.
*args: Query parameters.
Raises:
QueryExecutionError: If the query execution fails.
"""
try:
async with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute(query, args)
conn.commit()
except Pg8000DatabaseError as e:
logger.error(f"Query execution failed: {e}")
raise QueryExecutionError(f"Failed to execute query: {query}") from e
async def health_check(self) -> bool:
"""
Perform a health check by executing a simple query (e.g., SELECT 1).
Returns:
bool: True if the database is healthy, False otherwise.
Raises:
HealthCheckError: If the health check fails.
"""
try:
async with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT 1")
result = cursor.fetchone()
cursor.close()
# Check if the result is as expected
if result and result[0] == 1:
logger.info("Database health check succeeded.")
return True
else:
logger.error("Database health check failed: Unexpected result.")
raise HealthCheckError("Unexpected result from health check query.")
except Pg8000DatabaseError as e:
logger.error(f"Health check failed: {e}")
raise HealthCheckError("Failed to perform health check.") from e
# Dependency to get the database instance
async def get_db() -> AsyncGenerator[Database, None]:
settings = DatabaseSettings()
db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
await db.connect()
try:
yield db
finally:
await db.disconnect()
# Example usage
if __name__ == "__main__":
async def main():
settings = DatabaseSettings()
db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
await db.connect()
try:
# Perform a health check
is_healthy = await db.health_check()
print(f"Database health check: {'Success' if is_healthy else 'Failure'}")
except HealthCheckError as e:
print(f"Health check failed: {e}")
finally:
await db.disconnect()
asyncio.run(main())
|