from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from typing import Optional, List, Dict, Any
import logging
# 配置日志(方便调试)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class ESClientManager:
"""
Elasticsearch 客户端管理类
封装了连接、索引管理、文档增删改查等常用操作
"""
def __init__(self, host: str = "http://localhost:9200", username: Optional[str] = None,
password: Optional[str] = None):
"""
初始化 ES 客户端
Args:
host: ES 服务地址,默认 http://localhost:9200
username: 用户名(如果需要认证)
password: 密码(如果需要认证)
"""
self.host = host
self.client = None
self._connect(username, password)
def _connect(self, username: Optional[str] = None, password: Optional[str] = None):
"""
建立连接(内部方法)
"""
try:
if username and password:
# 带认证的连接
self.client = Elasticsearch(
[self.host],
basic_auth=(username, password),
request_timeout=30
)
else:
# 无认证连接
self.client = Elasticsearch([self.host], request_timeout=30)
# 测试连接
info = self.client.info()
logger.info(f"✅ ES 连接成功!版本: {info['version']['number']}")
except Exception as e:
logger.error(f"❌ ES 连接失败: {e}")
raise
def check_connection(self) -> bool:
"""
检查连接是否正常
Returns:
bool: True 表示连接正常
"""
try:
return self.client.ping()
except:
return False
# ==================== 索引管理 ====================
def create_index(self, index_name: str, mappings: Optional[Dict] = None, settings: Optional[Dict] = None) -> bool:
"""
创建索引
Args:
index_name: 索引名称
mappings: 字段映射,例如 {"properties": {"title": {"type": "text"}}}
settings: 索引设置,例如 {"number_of_shards": 1}
Returns:
bool: 创建成功返回 True
"""
try:
# 如果索引已存在,先删除(谨慎使用!)
if self.client.indices.exists(index=index_name):
logger.warning(f"⚠️ 索引 '{index_name}' 已存在,跳过创建")
return False
body = {}
if mappings:
body["mappings"] = mappings
if settings:
body["settings"] = settings
self.client.indices.create(index=index_name, body=body)
logger.info(f"✅ 索引 '{index_name}' 创建成功")
return True
except Exception as e:
logger.error(f"❌ 创建索引失败: {e}")
return False
def delete_index(self, index_name: str, force: bool = False) -> bool:
"""
删除索引(慎用!)
Args:
index_name: 索引名称
force: 是否强制删除(设为 True 跳过确认)
Returns:
bool: 删除成功返回 True
"""
if not force:
confirm = input(f"⚠️ 确认要删除索引 '{index_name}' 吗?(yes/no): ")
if confirm.lower() != 'yes':
logger.info("已取消删除")
return False
try:
self.client.indices.delete(index=index_name)
logger.info(f"✅ 索引 '{index_name}' 删除成功")
return True
except Exception as e:
logger.error(f"❌ 删除索引失败: {e}")
return False
def list_indices(self) -> List[str]:
"""
列出所有索引
Returns:
List[str]: 索引名称列表
"""
try:
indices = self.client.indices.get_alias(index="*")
# 过滤掉系统索引(以 . 开头的)
return [idx for idx in indices.keys() if not idx.startswith('.')]
except Exception as e:
logger.error(f"❌ 获取索引列表失败: {e}")
return []
def index_exists(self, index_name: str) -> bool:
"""
检查索引是否存在
Returns:
bool: 存在返回 True
"""
try:
return self.client.indices.exists(index=index_name)
except:
return False
# ==================== 文档操作 ====================
def insert_doc(self, index_name: str, doc: Dict, doc_id: Optional[str] = None) -> Optional[str]:
"""
插入单条文档
Args:
index_name: 索引名称
doc: 文档数据(字典)
doc_id: 文档ID(可选,不传则自动生成)
Returns:
Optional[str]: 文档ID,失败返回 None
"""
try:
result = self.client.index(index=index_name, id=doc_id, document=doc)
doc_id = result['_id']
logger.info(f"✅ 文档插入成功,ID: {doc_id}")
return doc_id
except Exception as e:
logger.error(f"❌ 插入文档失败: {e}")
return None
# def bulk_insert(self, index_name: str, docs: List[Dict]) -> int:
# """
# 批量插入文档(更高效)
#
# Args:
# index_name: 索引名称
# docs: 文档列表,每个元素是一个字典
#
# Returns:
# int: 成功插入的数量
# """
# try:
# actions = [
# {"_index": index_name, "_source": doc}
# for doc in docs
# ]
# success, failed = bulk(self.client, actions, stats_only=True)
# logger.info(f"✅ 批量插入成功: {success} 条,失败: {failed} 条")
# return success
# except Exception as e:
# logger.error(f"❌ 批量插入失败: {e}")
# return 0
# 在 bulk_insert 方法中添加 refresh 参数
def bulk_insert(self, index_name: str, docs: List[Dict], refresh: bool = True) -> int:
"""
批量插入文档(更高效)
Args:
index_name: 索引名称
docs: 文档列表,每个元素是一个字典
refresh: 是否立即刷新索引,默认 True
"""
try:
actions = [
{"_index": index_name, "_source": doc}
for doc in docs
]
success, failed = bulk(
self.client,
actions,
stats_only=True,
refresh=refresh # 添加这个参数
)
logger.info(f"✅ 批量插入成功: {success} 条,失败: {failed} 条")
return success
except Exception as e:
logger.error(f"❌ 批量插入失败: {e}")
return 0
def search(self, index_name: str, query: Dict, size: int = 10, from_: int = 0) -> List[Dict]:
"""
搜索文档
Args:
index_name: 索引名称
query: 查询 DSL,例如 {"match": {"title": "python"}}
size: 返回条数
from_: 偏移量(分页用)
Returns:
List[Dict]: 文档列表,每个文档包含 _id, _score, _source
"""
try:
result = self.client.search(
index=index_name,
query=query,
size=size,
from_=from_
)
hits = result['hits']['hits']
logger.info(f"🔍 搜索完成,共找到 {result['hits']['total']['value']} 条,返回 {len(hits)} 条")
# 格式化返回结果
docs = []
for hit in hits:
docs.append({
'_id': hit['_id'],
'_score': hit['_score'],
'_source': hit['_source']
})
return docs
except Exception as e:
logger.error(f"❌ 搜索失败: {e}")
return []
def get_doc(self, index_name: str, doc_id: str) -> Optional[Dict]:
"""
根据 ID 获取文档
Args:
index_name: 索引名称
doc_id: 文档ID
Returns:
Optional[Dict]: 文档数据,不存在返回 None
"""
try:
result = self.client.get(index=index_name, id=doc_id)
return result['_source']
except Exception as e:
logger.error(f"❌ 获取文档失败: {e}")
return None
def update_doc(self, index_name: str, doc_id: str, doc: Dict) -> bool:
"""
更新文档(部分更新)
Args:
index_name: 索引名称
doc_id: 文档ID
doc: 要更新的字段(字典)
Returns:
bool: 更新成功返回 True
"""
try:
self.client.update(index=index_name, id=doc_id, doc=doc)
logger.info(f"✅ 文档 '{doc_id}' 更新成功")
return True
except Exception as e:
logger.error(f"❌ 更新文档失败: {e}")
return False
def delete_doc(self, index_name: str, doc_id: str) -> bool:
"""
删除文档
Args:
index_name: 索引名称
doc_id: 文档ID
Returns:
bool: 删除成功返回 True
"""
try:
self.client.delete(index=index_name, id=doc_id)
logger.info(f"✅ 文档 '{doc_id}' 删除成功")
return True
except Exception as e:
logger.error(f"❌ 删除文档失败: {e}")
return False
def count_docs(self, index_name: str) -> int:
"""
统计索引中的文档数量
Returns:
int: 文档总数
"""
try:
result = self.client.count(index=index_name)
return result['count']
except Exception as e:
logger.error(f"❌ 统计文档数失败: {e}")
return 0
def close(self):
"""
关闭客户端连接
"""
if self.client:
self.client.close()
logger.info("🔒 ES 连接已关闭")
# ==================== 使用示例 ====================
if __name__ == "__main__":
# 1. 创建客户端
es_manager = ESClientManager("http://localhost:9200")
# 2. 创建索引
index_name = "test_articles"
mappings = {
"properties": {
"title": {"type": "text"},
"content": {"type": "text"},
"views": {"type": "integer"}
}
}
es_manager.create_index(index_name, mappings=mappings)
# 3. 插入文档
doc1 = {"title": "Python 入门", "content": "Python 是一门好语言", "views": 100}
doc_id = es_manager.insert_doc(index_name, doc1)
print(f"✅ 插入的文档 ID: {doc_id}")
# 4. 批量插入
docs = [
{"title": "ES 教程", "content": "Elasticsearch 搜索引擎", "views": 200},
{"title": "Docker 基础", "content": "容器化技术", "views": 150}
]
es_manager.bulk_insert(index_name, docs)
# 🌟 新增:强制刷新索引,让数据立即可见
es_manager.client.indices.refresh(index=index_name)
print("🔄 索引已刷新")
# 5. 搜索
results = es_manager.search(
index_name,
query={"multi_match": {"query": "Python", "fields": ["title", "content"]}}
)
print(f"\n🔍 搜索 'Python' 的结果:")
for doc in results:
print(f" 标题: {doc['_source']['title']}, 得分: {doc['_score']}")
# 6. 统计数量
total = es_manager.count_docs(index_name)
print(f"\n📊 总文档数: {total}")
# 7. 获取所有文档(验证)
all_docs = es_manager.search(
index_name,
query={"match_all": {}},
size=10
)
print(f"\n📄 所有文档:")
for doc in all_docs:
print(f" - {doc['_source']['title']} (浏览量: {doc['_source']['views']})")
# 8. 关闭连接
es_manager.close()
# docker run -d --name elasticvue -p 8080:8080 cars10/elasticvue