agentic-threat-hunting-framework 0.5.0__tar.gz → 0.5.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (84) hide show
  1. {agentic_threat_hunting_framework-0.5.0/agentic_threat_hunting_framework.egg-info → agentic_threat_hunting_framework-0.5.1}/PKG-INFO +1 -1
  2. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1/agentic_threat_hunting_framework.egg-info}/PKG-INFO +1 -1
  3. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/agentic_threat_hunting_framework.egg-info/SOURCES.txt +0 -7
  4. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/agents/base.py +5 -14
  5. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/pyproject.toml +1 -1
  6. agentic_threat_hunting_framework-0.5.0/athf/core/clickhouse_connection.py +0 -396
  7. agentic_threat_hunting_framework-0.5.0/athf/core/metrics_tracker.py +0 -518
  8. agentic_threat_hunting_framework-0.5.0/athf/core/query_executor.py +0 -169
  9. agentic_threat_hunting_framework-0.5.0/athf/core/query_parser.py +0 -203
  10. agentic_threat_hunting_framework-0.5.0/athf/core/query_suggester.py +0 -235
  11. agentic_threat_hunting_framework-0.5.0/athf/core/query_validator.py +0 -240
  12. agentic_threat_hunting_framework-0.5.0/athf/core/session_manager.py +0 -764
  13. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/LICENSE +0 -0
  14. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/MANIFEST.in +0 -0
  15. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/README.md +0 -0
  16. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/USING_ATHF.md +0 -0
  17. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/agentic_threat_hunting_framework.egg-info/dependency_links.txt +0 -0
  18. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/agentic_threat_hunting_framework.egg-info/entry_points.txt +0 -0
  19. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/agentic_threat_hunting_framework.egg-info/requires.txt +0 -0
  20. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/agentic_threat_hunting_framework.egg-info/top_level.txt +0 -0
  21. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/ATHF_level_3.png +0 -0
  22. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf-cli-workflow.gif +0 -0
  23. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf-level0.gif +0 -0
  24. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf-level1.gif +0 -0
  25. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf-level2.gif +0 -0
  26. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf-level3.gif +0 -0
  27. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf_fivelevels.png +0 -0
  28. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf_lock.png +0 -0
  29. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf_logo.png +0 -0
  30. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/assets/athf_manual_v_ai.png +0 -0
  31. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/__init__.py +0 -0
  32. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/__version__.py +0 -0
  33. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/agents/__init__.py +0 -0
  34. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/agents/llm/__init__.py +0 -0
  35. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/agents/llm/hunt_researcher.py +0 -0
  36. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/agents/llm/hypothesis_generator.py +0 -0
  37. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/cli.py +0 -0
  38. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/__init__.py +0 -0
  39. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/agent.py +0 -0
  40. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/context.py +0 -0
  41. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/env.py +0 -0
  42. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/hunt.py +0 -0
  43. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/init.py +0 -0
  44. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/investigate.py +0 -0
  45. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/research.py +0 -0
  46. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/similar.py +0 -0
  47. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/commands/splunk.py +0 -0
  48. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/__init__.py +0 -0
  49. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/attack_matrix.py +0 -0
  50. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/hunt_manager.py +0 -0
  51. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/hunt_parser.py +0 -0
  52. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/investigation_parser.py +0 -0
  53. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/research_manager.py +0 -0
  54. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/splunk_client.py +0 -0
  55. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/template_engine.py +0 -0
  56. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/core/web_search.py +0 -0
  57. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/__init__.py +0 -0
  58. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/CHANGELOG.md +0 -0
  59. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/CLI_REFERENCE.md +0 -0
  60. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/INSTALL.md +0 -0
  61. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/README.md +0 -0
  62. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/environment.md +0 -0
  63. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/getting-started.md +0 -0
  64. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/level4-agentic-workflows.md +0 -0
  65. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/lock-pattern.md +0 -0
  66. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/maturity-model.md +0 -0
  67. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/docs/why-athf.md +0 -0
  68. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/hunts/FORMAT_GUIDELINES.md +0 -0
  69. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/hunts/H-0001.md +0 -0
  70. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/hunts/H-0002.md +0 -0
  71. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/hunts/H-0003.md +0 -0
  72. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/hunts/README.md +0 -0
  73. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/integrations/MCP_CATALOG.md +0 -0
  74. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/integrations/README.md +0 -0
  75. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/integrations/quickstart/splunk.md +0 -0
  76. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/knowledge/hunting-knowledge.md +0 -0
  77. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/prompts/README.md +0 -0
  78. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/prompts/ai-workflow.md +0 -0
  79. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/prompts/basic-prompts.md +0 -0
  80. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/data/templates/HUNT_LOCK.md +0 -0
  81. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/plugin_system.py +0 -0
  82. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/athf/utils/__init__.py +0 -0
  83. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/setup.cfg +0 -0
  84. {agentic_threat_hunting_framework-0.5.0 → agentic_threat_hunting_framework-0.5.1}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentic-threat-hunting-framework
3
- Version: 0.5.0
3
+ Version: 0.5.1
4
4
  Summary: Agentic Threat Hunting Framework - Memory and AI for threat hunters
5
5
  Author-email: Sydney Marrone <athf@nebulock.io>
6
6
  Maintainer-email: Sydney Marrone <athf@nebulock.io>
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentic-threat-hunting-framework
3
- Version: 0.5.0
3
+ Version: 0.5.1
4
4
  Summary: Agentic Threat Hunting Framework - Memory and AI for threat hunters
5
5
  Author-email: Sydney Marrone <athf@nebulock.io>
6
6
  Maintainer-email: Sydney Marrone <athf@nebulock.io>
@@ -41,17 +41,10 @@ athf/commands/similar.py
41
41
  athf/commands/splunk.py
42
42
  athf/core/__init__.py
43
43
  athf/core/attack_matrix.py
44
- athf/core/clickhouse_connection.py
45
44
  athf/core/hunt_manager.py
46
45
  athf/core/hunt_parser.py
47
46
  athf/core/investigation_parser.py
48
- athf/core/metrics_tracker.py
49
- athf/core/query_executor.py
50
- athf/core/query_parser.py
51
- athf/core/query_suggester.py
52
- athf/core/query_validator.py
53
47
  athf/core/research_manager.py
54
- athf/core/session_manager.py
55
48
  athf/core/splunk_client.py
56
49
  athf/core/template_engine.py
57
50
  athf/core/web_search.py
@@ -1,4 +1,4 @@
1
- """Base classes for hunt-vault agents."""
1
+ """Base classes for ATHF agents."""
2
2
 
3
3
  import os
4
4
  from abc import ABC, abstractmethod
@@ -89,6 +89,8 @@ class LLMAgent(Agent[InputT, OutputT]):
89
89
  ) -> None:
90
90
  """Log LLM call metrics to centralized tracker.
91
91
 
92
+ Override this method in subclasses or plugins to implement custom metrics tracking.
93
+
92
94
  Args:
93
95
  agent_name: Name of the agent (e.g., "hypothesis-generator")
94
96
  model_id: Bedrock model ID
@@ -97,19 +99,8 @@ class LLMAgent(Agent[InputT, OutputT]):
97
99
  cost_usd: Estimated cost in USD
98
100
  duration_ms: Call duration in milliseconds
99
101
  """
100
- try:
101
- from athf.core.metrics_tracker import MetricsTracker
102
-
103
- MetricsTracker.get_instance().log_bedrock_call(
104
- agent=agent_name,
105
- model_id=model_id,
106
- input_tokens=input_tokens,
107
- output_tokens=output_tokens,
108
- cost_usd=cost_usd,
109
- duration_ms=duration_ms,
110
- )
111
- except Exception:
112
- pass # Never fail agent execution due to metrics logging
102
+ # No-op by default. Override in plugins for custom metrics tracking.
103
+ pass
113
104
 
114
105
  def _get_llm_client(self) -> Any:
115
106
  """Get AWS Bedrock runtime client for Claude models.
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "agentic-threat-hunting-framework"
7
- version = "0.5.0"
7
+ version = "0.5.1"
8
8
  description = "Agentic Threat Hunting Framework - Memory and AI for threat hunters"
9
9
  readme = {file = "README.md", content-type = "text/markdown"}
10
10
  requires-python = ">=3.8"
@@ -1,396 +0,0 @@
1
- """ClickHouse connection management and configuration."""
2
-
3
- import os
4
- import time
5
- from dataclasses import dataclass
6
- from pathlib import Path
7
- from typing import Any, Dict, Optional
8
-
9
- import yaml
10
-
11
-
12
- # Custom exceptions
13
- class ClickHouseConfigError(Exception):
14
- """Configuration validation or loading error."""
15
-
16
- pass
17
-
18
-
19
- class ClickHouseConnectionError(Exception):
20
- """Connection establishment or network error."""
21
-
22
- pass
23
-
24
-
25
- class ClickHouseQueryError(Exception):
26
- """Query execution error."""
27
-
28
- pass
29
-
30
-
31
- @dataclass
32
- class ClickHouseConfig:
33
- """Configuration for ClickHouse connection.
34
-
35
- Attributes:
36
- host: ClickHouse server hostname
37
- port: ClickHouse server port (default: 8443 for HTTPS)
38
- username: Database username (required)
39
- password: Database password (required)
40
- database: Default database name (default: "default")
41
- secure: Use SSL/TLS encryption (default: True)
42
- """
43
-
44
- host: str
45
- port: int
46
- username: str
47
- password: str
48
- database: str = "default"
49
- secure: bool = True
50
-
51
- @classmethod
52
- def load(cls) -> "ClickHouseConfig":
53
- """Load configuration from environment variables and config file.
54
-
55
- Precedence order:
56
- 1. Environment variables (highest priority)
57
- 2. Config file (~/.athf/clickhouse.yaml)
58
- 3. Hardcoded defaults (host, port, database only)
59
-
60
- Returns:
61
- ClickHouseConfig instance
62
-
63
- Raises:
64
- ClickHouseConfigError: If required credentials are missing
65
- """
66
- config_data: Dict[str, Any] = {}
67
-
68
- # Load from config file first
69
- config_file = Path.home() / ".athf" / "clickhouse.yaml"
70
- if config_file.exists():
71
- try:
72
- with open(config_file, "r") as f:
73
- yaml_data = yaml.safe_load(f) or {}
74
- clickhouse_section = yaml_data.get("clickhouse", {})
75
- config_data.update(clickhouse_section)
76
- except Exception as e:
77
- raise ClickHouseConfigError(f"Failed to load config file {config_file}: {e}")
78
-
79
- # Override with environment variables
80
- if host := os.getenv("CLICKHOUSE_HOST"):
81
- config_data["host"] = host
82
- if port := os.getenv("CLICKHOUSE_PORT"):
83
- config_data["port"] = int(port)
84
- if user := os.getenv("CLICKHOUSE_USER"):
85
- config_data["username"] = user
86
- if password := os.getenv("CLICKHOUSE_PASSWORD"):
87
- config_data["password"] = password
88
- if database := os.getenv("CLICKHOUSE_DATABASE"):
89
- config_data["database"] = database
90
- if secure := os.getenv("CLICKHOUSE_SECURE"):
91
- config_data["secure"] = secure.lower() in ("true", "1", "yes")
92
-
93
- # Apply defaults
94
- config_data.setdefault("host", "ohma99qewu.us-east-1.aws.clickhouse.cloud")
95
- config_data.setdefault("port", 8443)
96
- config_data.setdefault("database", "default")
97
- config_data.setdefault("secure", True)
98
-
99
- # Validate required fields
100
- if "username" not in config_data or not config_data["username"]:
101
- raise ClickHouseConfigError(
102
- "Missing required credential: CLICKHOUSE_USER environment variable not set. "
103
- "Set credentials with: export CLICKHOUSE_USER='your_username'"
104
- )
105
- if "password" not in config_data or not config_data["password"]:
106
- raise ClickHouseConfigError(
107
- "Missing required credential: CLICKHOUSE_PASSWORD environment variable not set. "
108
- "Set credentials with: export CLICKHOUSE_PASSWORD='your_password'"
109
- )
110
-
111
- return cls(
112
- host=config_data["host"],
113
- port=config_data["port"],
114
- username=config_data["username"],
115
- password=config_data["password"],
116
- database=config_data["database"],
117
- secure=config_data["secure"],
118
- )
119
-
120
- def to_dict(self, mask_password: bool = True) -> Dict[str, Any]:
121
- """Convert config to dictionary.
122
-
123
- Args:
124
- mask_password: If True, replace password with asterisks
125
-
126
- Returns:
127
- Dictionary representation of config
128
- """
129
- return {
130
- "host": self.host,
131
- "port": self.port,
132
- "username": self.username,
133
- "password": "***" if mask_password else self.password,
134
- "database": self.database,
135
- "secure": self.secure,
136
- }
137
-
138
-
139
- class ClickHouseConnectionManager:
140
- """Singleton connection manager for ClickHouse queries.
141
-
142
- Manages a single ClickHouse client instance with lazy initialization.
143
- Connection is reused across multiple queries within the same process.
144
- """
145
-
146
- _instance: Optional["ClickHouseConnectionManager"] = None
147
- _client: Optional[Any] = None # Type: clickhouse_connect.driver.Client
148
- _config: Optional[ClickHouseConfig] = None
149
-
150
- def __new__(cls) -> "ClickHouseConnectionManager":
151
- """Ensure only one instance exists (singleton pattern)."""
152
- if cls._instance is None:
153
- cls._instance = super().__new__(cls)
154
- return cls._instance
155
-
156
- @classmethod
157
- def get_instance(cls) -> "ClickHouseConnectionManager":
158
- """Get the singleton instance.
159
-
160
- Returns:
161
- ClickHouseConnectionManager instance
162
- """
163
- if cls._instance is None:
164
- cls._instance = cls()
165
- return cls._instance
166
-
167
- def get_client(self) -> Any:
168
- """Get or create ClickHouse client (lazy initialization).
169
-
170
- Returns:
171
- ClickHouse client instance
172
-
173
- Raises:
174
- ClickHouseConnectionError: If connection fails
175
- """
176
- if self._client is None:
177
- self._client = self._create_client()
178
- return self._client
179
-
180
- def _create_client(self) -> Any:
181
- """Create ClickHouse client from configuration.
182
-
183
- Returns:
184
- ClickHouse client instance
185
-
186
- Raises:
187
- ClickHouseConnectionError: If client creation fails
188
- ClickHouseConfigError: If configuration is invalid
189
- """
190
- try:
191
- import clickhouse_connect
192
- except ImportError:
193
- raise ClickHouseConnectionError(
194
- "clickhouse-connect not installed. Install with: pip install 'hunt-vault[clickhouse]'"
195
- )
196
-
197
- # Load configuration
198
- if self._config is None:
199
- self._config = ClickHouseConfig.load()
200
-
201
- # Create client with retry logic
202
- max_retries = 2
203
- retry_delay = 5 # seconds
204
-
205
- for attempt in range(max_retries):
206
- try:
207
- # Check if running in AWS Lambda (no SSL verification)
208
- is_lambda = os.environ.get("AWS_EXECUTION_ENV") or os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
209
-
210
- client = clickhouse_connect.get_client(
211
- host=self._config.host,
212
- port=self._config.port,
213
- username=self._config.username,
214
- password=self._config.password,
215
- database=self._config.database,
216
- secure=self._config.secure,
217
- verify=not bool(is_lambda), # Disable SSL verification in Lambda
218
- )
219
- # Test connection with simple query
220
- client.command("SELECT 1")
221
- return client
222
- except Exception as e:
223
- if "authentication" in str(e).lower() or "credential" in str(e).lower():
224
- # Authentication failures should not retry
225
- raise ClickHouseConnectionError(
226
- f"Authentication failed: Invalid credentials for user '{self._config.username}'. "
227
- f"Check CLICKHOUSE_USER and CLICKHOUSE_PASSWORD environment variables."
228
- ) from e
229
- elif attempt < max_retries - 1:
230
- # Network errors: retry once
231
- time.sleep(retry_delay)
232
- continue
233
- else:
234
- # Final attempt failed
235
- raise ClickHouseConnectionError(f"Failed to connect to ClickHouse at {self._config.host}: {e}") from e
236
-
237
- # Should never reach here due to max_retries logic, but for type safety
238
- raise ClickHouseConnectionError("Failed to establish connection after retries")
239
-
240
- def get_config(self) -> ClickHouseConfig:
241
- """Get current configuration.
242
-
243
- Returns:
244
- ClickHouseConfig instance
245
-
246
- Raises:
247
- ClickHouseConfigError: If configuration loading fails
248
- """
249
- if self._config is None:
250
- self._config = ClickHouseConfig.load()
251
- return self._config
252
-
253
- def close(self) -> None:
254
- """Close the current connection.
255
-
256
- Note: Typically not needed for CLI use cases (process termination handles cleanup).
257
- Provided for completeness and testing.
258
- """
259
- if self._client is not None:
260
- try:
261
- self._client.close()
262
- except Exception: # nosec B110 - cleanup, safe to ignore failures
263
- pass # Best effort close
264
- finally:
265
- self._client = None
266
-
267
-
268
- class ClickHouseClient:
269
- """Wrapper for ClickHouse query execution with formatted output."""
270
-
271
- def __init__(self) -> None:
272
- """Initialize ClickHouse client wrapper."""
273
- self.manager = ClickHouseConnectionManager.get_instance()
274
-
275
- def execute_query(self, query: str, format: str = "json") -> Dict[str, Any]:
276
- """Execute query and return formatted results.
277
-
278
- Args:
279
- query: SQL query to execute
280
- format: Output format ('json', 'table', 'csv')
281
-
282
- Returns:
283
- Dictionary with query results and metadata:
284
- {
285
- 'columns': List[str],
286
- 'data': List[List[Any]],
287
- 'rows': int,
288
- 'elapsed': str,
289
- 'query': str
290
- }
291
-
292
- Raises:
293
- ClickHouseQueryError: If query execution fails
294
- """
295
- try:
296
- client = self.manager.get_client()
297
- start_time = time.time()
298
-
299
- # Execute query
300
- result = client.query(query)
301
-
302
- elapsed = time.time() - start_time
303
- elapsed_ms = int(elapsed * 1000)
304
-
305
- # Extract column names and data
306
- columns = result.column_names
307
- data = result.result_rows
308
- rows = len(data)
309
-
310
- # Auto-log metrics to centralized tracker
311
- try:
312
- from athf.core.metrics_tracker import MetricsTracker
313
-
314
- MetricsTracker.get_instance().log_clickhouse_query(
315
- sql=query,
316
- duration_ms=elapsed_ms,
317
- rows=rows,
318
- status="success",
319
- )
320
- except Exception:
321
- pass # Never fail query execution due to metrics logging
322
-
323
- return {
324
- "columns": columns,
325
- "data": data,
326
- "rows": rows,
327
- "elapsed": f"{elapsed:.3f}s",
328
- "query": query,
329
- }
330
-
331
- except Exception as e:
332
- # Log error metrics
333
- try:
334
- from athf.core.metrics_tracker import MetricsTracker
335
-
336
- status = "timeout" if "timeout" in str(e).lower() else "error"
337
- MetricsTracker.get_instance().log_clickhouse_query(
338
- sql=query,
339
- duration_ms=0, # Unknown duration on error
340
- rows=0,
341
- status=status,
342
- )
343
- except Exception:
344
- pass # Never fail due to metrics logging
345
-
346
- # Check for timeout errors
347
- if "timeout" in str(e).lower():
348
- raise ClickHouseQueryError(
349
- f"Query timeout: {e}\n\n"
350
- "Tips to avoid timeouts:\n"
351
- " 1. Add time bounds: WHERE timestamp >= now() - INTERVAL 7 DAY\n"
352
- " 2. Start with small LIMIT: LIMIT 100\n"
353
- " 3. Filter early: Add WHERE clause before aggregations\n"
354
- ' 4. Validate query: athf validate query --sql "..."'
355
- ) from e
356
- else:
357
- raise ClickHouseQueryError(f"Query execution failed: {e}") from e
358
-
359
- def test_connection(self) -> Dict[str, Any]:
360
- """Test ClickHouse connection with simple query.
361
-
362
- Returns:
363
- Dictionary with connection status and details:
364
- {
365
- 'success': bool,
366
- 'host': str,
367
- 'port': int,
368
- 'database': str,
369
- 'message': str
370
- }
371
-
372
- Raises:
373
- ClickHouseConnectionError: If connection test fails
374
- """
375
- try:
376
- client = self.manager.get_client()
377
- client.command("SELECT 1")
378
-
379
- config = self.manager.get_config()
380
-
381
- return {
382
- "success": True,
383
- "host": config.host,
384
- "port": config.port,
385
- "database": config.database,
386
- "message": "Connection successful",
387
- }
388
- except Exception as e:
389
- config = self.manager.get_config()
390
- return {
391
- "success": False,
392
- "host": config.host,
393
- "port": config.port,
394
- "database": config.database,
395
- "message": f"Connection failed: {e}",
396
- }