client.py•6.46 kB
"""
Search Client
High-level client for Elasticsearch entity searches.
"""
import logging
import os
from typing import Optional, List, Dict, Any
from elasticsearch import Elasticsearch
from elasticsearch_search_lib.query_builder import QueryBuilder
from elasticsearch_search_lib.search_handler import SearchHandler
from elasticsearch_search_lib.config.loader import ConfigLoader
from elasticsearch_search_lib.models import SearchResponse, EntityConfig, FieldConfig
from elasticsearch_search_lib.exceptions import ValidationError
logger = logging.getLogger(__name__)
class SearchClient:
"""
High-level client for entity searches.
Provides a simple, intuitive API for searching ITSM entities
in Elasticsearch.
Example:
>>> client = SearchClient(tenant_id="apolo")
>>> results = await client.search("user", "john doe")
>>> for item in results.items:
... print(item.data, item.score)
"""
def __init__(
self,
tenant_id: str,
es_host: Optional[str] = None,
es_port: int = 9200,
config_path: Optional[str] = None,
):
"""
Initialize search client.
Args:
tenant_id: Tenant identifier for index naming
es_host: Elasticsearch host (defaults to ES_HOST env var or 'localhost')
es_port: Elasticsearch port (default: 9200)
config_path: Path to configuration file (uses default if None)
"""
self.tenant_id = tenant_id
self.es_host = es_host or os.getenv("ES_HOST", "localhost")
self.es_port = es_port
# Initialize components
self.config_loader = ConfigLoader(config_path)
self.query_builder = QueryBuilder(self.config_loader)
# Create Elasticsearch client
self.es_client = self._create_es_client()
# Create search handler
self.search_handler = SearchHandler(
es_client=self.es_client,
query_builder=self.query_builder,
)
logger.info(
f"SearchClient initialized for tenant '{tenant_id}' "
f"with {len(self.get_supported_entities())} supported entity types"
)
def _create_es_client(self) -> Elasticsearch:
"""Create Elasticsearch client."""
es_url = f"http://{self.es_host}:{self.es_port}"
try:
client = Elasticsearch(
[es_url],
request_timeout=30,
max_retries=3,
retry_on_timeout=True
)
logger.info(f"Elasticsearch client created for {es_url}")
return client
except Exception as e:
logger.warning(f"Failed to create Elasticsearch client: {e}")
# Return a mock object for testing
from unittest.mock import Mock
return Mock()
async def search(
self,
entity_type: str,
query: str,
limit: Optional[int] = None,
from_offset: int = 0,
) -> SearchResponse:
"""
Search for entities.
Args:
entity_type: Type of entity to search (user, impact, etc.)
query: Search query string
limit: Maximum number of results (uses entity default if None)
from_offset: Offset for pagination
Returns:
SearchResponse with results or error
Example:
>>> results = await client.search("user", "john doe", limit=10)
>>> print(f"Found {results.total_hits} users")
>>> for user in results.items:
... print(user.data['user_name'], user.score)
"""
return await self.search_handler.search(
entity_type=entity_type,
query=query,
tenant_id=self.tenant_id,
limit=limit,
from_offset=from_offset,
)
def get_supported_entities(self) -> List[str]:
"""
Get list of supported entity types.
Returns:
List of entity type names
Example:
>>> entities = client.get_supported_entities()
>>> print(entities)
['impact', 'urgency', 'priority', 'status', ...]
"""
return self.query_builder.get_supported_entities()
def get_entity_config(self, entity_type: str) -> EntityConfig:
"""
Get configuration for an entity type.
Args:
entity_type: Entity type name
Returns:
EntityConfig instance
Example:
>>> config = client.get_entity_config("user")
>>> print(config.default_limit, config.max_limit)
10 100
"""
return self.query_builder.get_entity_config(entity_type)
def get_entity_fields(self, entity_type: str) -> List[FieldConfig]:
"""
Get enabled fields for an entity type.
Args:
entity_type: Entity type name
Returns:
List of FieldConfig instances
Example:
>>> fields = client.get_entity_fields("user")
>>> for field in fields:
... print(field.name, field.boost, field.fuzziness)
"""
config = self.get_entity_config(entity_type)
return config.get_enabled_fields()
def is_entity_supported(self, entity_type: str) -> bool:
"""
Check if an entity type is supported.
Args:
entity_type: Entity type name
Returns:
True if entity is supported
Example:
>>> if client.is_entity_supported("user"):
... results = await client.search("user", "john")
"""
return self.query_builder.is_entity_supported(entity_type)
def get_index_name(self, entity_type: str) -> str:
"""
Get Elasticsearch index name for an entity type.
Args:
entity_type: Entity type name
Returns:
Index name string
Example:
>>> index = client.get_index_name("user")
>>> print(index)
apolo_user
"""
return self.query_builder.get_index_name(entity_type, self.tenant_id)