# import logging # from typing import Dict, List, Optional, AsyncGenerator # from pydantic import BaseSettings, PostgresDsn # import pg8000 # from pg8000 import Connection, Cursor # from pg8000.exceptions import DatabaseError # import asyncio # from contextlib import asynccontextmanager # from dataclasses import dataclass # from threading import Lock # # 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 # class Config: # env_file = ".env" # @dataclass # class DatabaseConfig: # username: str # password: str # hostname: str # port: int # database: str # class DatabaseError(Exception): # """Custom exception for database errors.""" # pass # class Database: # def __init__(self, db_url: str, pool_size: int): # self.db_url = db_url # self.pool_size = pool_size # self.pool: List[Connection] = [] # self.lock = Lock() # self.config = self._parse_db_url() # def _parse_db_url(self) -> DatabaseConfig: # """Parse the database URL into components.""" # result = urlparse(self.db_url) # return DatabaseConfig( # username=result.username, # password=result.password, # hostname=result.hostname, # port=result.port or 5432, # database=result.path.lstrip("/"), # ) # async def connect(self) -> None: # """Create a connection pool.""" # try: # for _ in range(self.pool_size): # conn = await self._create_connection() # self.pool.append(conn) # logger.info( # f"Database connection pool created with {self.pool_size} connections." # ) # except DatabaseError as e: # logger.error(f"Failed to create database connection pool: {e}") # raise # async def _create_connection(self) -> Connection: # """Create a single database connection.""" # try: # conn = pg8000.connect( # user=self.config.username, # password=self.config.password, # host=self.config.hostname, # port=self.config.port, # database=self.config.database, # ) # return conn # except DatabaseError as e: # logger.error(f"Failed to create database connection: {e}") # raise DatabaseError("Failed to create database connection.") # 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: # raise DatabaseError("Database connection pool is empty.") # conn = self.pool.pop() # try: # yield conn # finally: # with self.lock: # self.pool.append(conn) # async def fetch(self, query: str, *args) -> List[Dict]: # """ # Execute a SELECT query and return the results as a list of dictionaries. # Args: # query (str): The SQL query to execute. # *args: Query parameters. # Returns: # List[Dict]: A list of dictionaries where keys are column names and values are column values. # """ # try: # async with self.get_connection() as conn: # cursor: Cursor = conn.cursor() # cursor.execute(query, args) # rows = cursor.fetchall() # columns = [desc[0] for desc in cursor.description] # return [dict(zip(columns, row)) for row in rows] # except DatabaseError as e: # logger.error(f"Error executing query: {query}. Error: {e}") # raise DatabaseError(f"Failed to execute query: {query}") # 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. # """ # try: # async with self.get_connection() as conn: # cursor: Cursor = conn.cursor() # cursor.execute(query, args) # conn.commit() # except DatabaseError as e: # logger.error(f"Error executing query: {query}. Error: {e}") # raise DatabaseError(f"Failed to execute query: {query}") # # 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: # # Example query # query = """ # SELECT # ppt.type AS product_type, # pc.name AS product_category # FROM # product_producttype ppt # INNER JOIN # product_category pc # ON # ppt.category_id = pc.id # """ # result = await db.fetch(query) # print(result) # finally: # await db.disconnect() # asyncio.run(main()) # import logging # from urllib.parse import urlparse # from typing import Dict, List, Optional, AsyncGenerator # from pydantic_settings import BaseSettings # from pydantic import PostgresDsn # import pg8000 # from pg8000 import Connection, Cursor # from pg8000.exceptions import DatabaseError # import asyncio # from contextlib import asynccontextmanager # from dataclasses import dataclass # from threading import Lock # # 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 # class Config: # env_file = ".env" # @dataclass # class DatabaseConfig: # username: str # password: str # hostname: str # port: int # database: str # class DatabaseError(Exception): # """Custom exception for database errors.""" # pass # class Database: # def __init__(self, db_url: str, pool_size: int): # self.db_url = db_url # self.pool_size = pool_size # self.pool: List[Connection] = [] # self.lock = Lock() # self.config = self._parse_db_url() # def _parse_db_url(self) -> DatabaseConfig: # """Parse the database URL into components.""" # # Convert PostgresDsn to a string # db_url_str = str(self.db_url) # result = urlparse(db_url_str) # return DatabaseConfig( # username=result.username, # password=result.password, # hostname=result.hostname, # port=result.port or 5432, # database=result.path.lstrip("/"), # ) # async def connect(self) -> None: # """Create a connection pool.""" # try: # for _ in range(self.pool_size): # conn = await self._create_connection() # self.pool.append(conn) # logger.info( # f"Database connection pool created with {self.pool_size} connections." # ) # except DatabaseError as e: # logger.error(f"Failed to create database connection pool: {e}") # raise # async def _create_connection(self) -> Connection: # """Create a single database connection.""" # try: # conn = pg8000.connect( # user=self.config.username, # password=self.config.password, # host=self.config.hostname, # port=self.config.port, # database=self.config.database, # ) # return conn # except DatabaseError as e: # logger.error(f"Failed to create database connection: {e}") # raise DatabaseError("Failed to create database connection.") # 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: # raise DatabaseError("Database connection pool is empty.") # conn = self.pool.pop() # try: # yield conn # finally: # with self.lock: # self.pool.append(conn) # async def fetch(self, query: str, *args) -> List[Dict]: # """ # Execute a SELECT query and return the results as a list of dictionaries. # Args: # query (str): The SQL query to execute. # *args: Query parameters. # Returns: # List[Dict]: A list of dictionaries where keys are column names and values are column values. # """ # try: # async with self.get_connection() as conn: # cursor: Cursor = conn.cursor() # cursor.execute(query, args) # rows = cursor.fetchall() # columns = [desc[0] for desc in cursor.description] # return [dict(zip(columns, row)) for row in rows] # except DatabaseError as e: # logger.error(f"Error executing query: {query}. Error: {e}") # raise DatabaseError(f"Failed to execute query: {query}") # 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. # """ # try: # async with self.get_connection() as conn: # cursor: Cursor = conn.cursor() # cursor.execute(query, args) # conn.commit() # except DatabaseError as e: # logger.error(f"Error executing query: {query}. Error: {e}") # raise DatabaseError(f"Failed to execute query: {query}") # # 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: # # Example query # query = "SELECT * FROM your_table LIMIT 10" # query = """ # SELECT # ppt.type AS product_type, # pc.name AS product_category # FROM # product_producttype ppt # INNER JOIN # product_category pc # ON # ppt.category_id = pc.id # """ # result = await db.fetch(query) # print(result) # finally: # await db.disconnect() # asyncio.run(main()) 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) -> List[Dict]: """ Execute a SELECT query and return the results as a list of dictionaries. Args: query (str): The SQL query to execute. *args: Query parameters. Returns: List[Dict]: A list of dictionaries where keys are column names and values are 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] return [dict(zip(columns, row)) for row in rows] 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())