Coverage for src/dataknobs_data/pooling/elasticsearch.py: 38%

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

2 

3from dataclasses import dataclass 

4from typing import Optional, Any, List 

5 

6from .base import BasePoolConfig 

7 

8 

9@dataclass 

10class ElasticsearchPoolConfig(BasePoolConfig): 

11 """Configuration for Elasticsearch connection pools.""" 

12 hosts: List[str] = None 

13 index: str = "records" 

14 api_key: Optional[str] = None 

15 basic_auth: Optional[tuple] = None 

16 verify_certs: bool = True 

17 ca_certs: Optional[str] = None 

18 client_cert: Optional[str] = None 

19 client_key: Optional[str] = None 

20 ssl_show_warn: bool = True 

21 

22 def __post_init__(self): 

23 """Set default hosts if not provided.""" 

24 if self.hosts is None: 

25 self.hosts = ["http://localhost:9200"] 

26 

27 def to_connection_string(self) -> str: 

28 """Convert to connection string (not used for ES, but required by base).""" 

29 return ";".join(self.hosts) 

30 

31 def to_hash_key(self) -> tuple: 

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

33 return (tuple(self.hosts), self.index) 

34 

35 @classmethod 

36 def from_dict(cls, config: dict) -> "ElasticsearchPoolConfig": 

37 """Create from configuration dictionary.""" 

38 # Handle both old-style (host, port) and new-style (hosts) configuration 

39 if "hosts" in config: 

40 hosts = config["hosts"] 

41 elif "host" in config: 

42 host = config["host"] 

43 port = config.get("port", 9200) 

44 # Check if it already has a scheme 

45 if host.startswith("http://") or host.startswith("https://"): 

46 hosts = [f"{host}:{port}" if ":" not in host.split("://")[1] else host] 

47 else: 

48 hosts = [f"http://{host}:{port}"] 

49 else: 

50 hosts = ["http://localhost:9200"] 

51 

52 return cls( 

53 hosts=hosts, 

54 index=config.get("index", "records"), 

55 api_key=config.get("api_key"), 

56 basic_auth=config.get("basic_auth"), 

57 verify_certs=config.get("verify_certs", True), 

58 ca_certs=config.get("ca_certs"), 

59 client_cert=config.get("client_cert"), 

60 client_key=config.get("client_key"), 

61 ssl_show_warn=config.get("ssl_show_warn", True) 

62 ) 

63 

64 

65async def create_async_elasticsearch_client(config: ElasticsearchPoolConfig): 

66 """Create an async Elasticsearch client.""" 

67 from elasticsearch import AsyncElasticsearch 

68 

69 # Build client configuration 

70 client_config = { 

71 "hosts": config.hosts, 

72 } 

73 

74 # Add authentication if provided 

75 if config.api_key: 

76 client_config["api_key"] = config.api_key 

77 elif config.basic_auth: 

78 client_config["basic_auth"] = config.basic_auth 

79 

80 # Add SSL configuration 

81 if config.ca_certs: 

82 client_config["ca_certs"] = config.ca_certs 

83 if config.client_cert: 

84 client_config["client_cert"] = config.client_cert 

85 if config.client_key: 

86 client_config["client_key"] = config.client_key 

87 

88 client_config["verify_certs"] = config.verify_certs 

89 client_config["ssl_show_warn"] = config.ssl_show_warn 

90 

91 # Create and return the client 

92 return AsyncElasticsearch(**client_config) 

93 

94 

95async def validate_elasticsearch_client(client) -> None: 

96 """Validate an Elasticsearch client by pinging it.""" 

97 if not await client.ping(): 

98 raise ConnectionError("Failed to ping Elasticsearch") 

99 

100 

101async def close_elasticsearch_client(client) -> None: 

102 """Properly close an Elasticsearch client and its underlying connections.""" 

103 if client: 

104 try: 

105 await client.close() 

106 except Exception: 

107 pass # Ignore errors during cleanup