Coverage for src/dataknobs_data/pooling/postgres.py: 74%

1"""PostgreSQL-specific connection pooling implementation.""" 

2 

3from dataclasses import dataclass 

4from typing import Optional, Any 

5 

6from .base import BasePoolConfig 

7 

8 

9@dataclass 

10class PostgresPoolConfig(BasePoolConfig): 

11 """Configuration for PostgreSQL connection pools.""" 

12 host: str = "localhost" 

13 port: int = 5432 

14 database: str = "postgres" 

15 user: str = "postgres" 

16 password: str = "" 

17 min_size: int = 10 

18 max_size: int = 10 

19 command_timeout: Optional[float] = None 

20 ssl: Optional[Any] = None 

21 

22 def to_connection_string(self) -> str: 

23 """Convert to PostgreSQL connection string.""" 

24 return f"postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}" 

25 

26 def to_hash_key(self) -> tuple: 

27 """Create a hashable key for this configuration.""" 

28 return (self.host, self.port, self.database, self.user) 

29 

30 @classmethod 

31 def from_dict(cls, config: dict) -> "PostgresPoolConfig": 

32 """Create from configuration dictionary.""" 

33 return cls( 

34 host=config.get("host", "localhost"), 

35 port=config.get("port", 5432), 

36 database=config.get("database", "postgres"), 

37 user=config.get("user", "postgres"), 

38 password=config.get("password", ""), 

39 min_size=config.get("min_pool_size", 10), 

40 max_size=config.get("max_pool_size", 10), 

41 command_timeout=config.get("command_timeout"), 

42 ssl=config.get("ssl") 

43 ) 

44 

45 

46async def create_asyncpg_pool(config: PostgresPoolConfig): 

47 """Create an asyncpg connection pool.""" 

48 import asyncpg 

49 return await asyncpg.create_pool( 

50 config.to_connection_string(), 

51 min_size=config.min_size, 

52 max_size=config.max_size, 

53 command_timeout=config.command_timeout, 

54 ssl=config.ssl 

55 ) 

56 

57 

58async def validate_asyncpg_pool(pool) -> None: 

59 """Validate an asyncpg pool by running a simple query.""" 

60 async with pool.acquire() as conn: 

61 await conn.fetchval("SELECT 1")